Merge pull request #5558 from aka-rider/portfwd-current-client

portfwd: dial through the current guest agent client, not the (poissibly stale one) captured at listen time

pkg/hostagent/hostagent.go

@@ -981,7 +981,7 @@ func (a *HostAgent) processGuestAgentEvents(ctx context.Context, client *guestag
 		if useSSHFwd {
 			a.portForwarder.OnEvent(ctx, ev)
 		} else {
-			dialContext := portfwd.DialContextToGRPCTunnel(client)
+			dialContext := portfwd.DialContextToGRPCTunnel(a.getClient)
 			a.grpcPortForwarder.OnEvent(ctx, dialContext, ev)
 		}
 		return nil

pkg/portfwd/client.go

@@ -44,13 +44,21 @@ func HandleUDPConnection(ctx context.Context, dialContext func(ctx context.Conte
 	logrus.Debugf("udp proxy for guestAddr: %s closed", guestAddr)
 }
 
-func DialContextToGRPCTunnel(client *guestagentclient.GuestAgentClient) func(ctx context.Context, network, addr string) (net.Conn, error) {
+// DialContextToGRPCTunnel returns a dialer that opens a new gRPC tunnel stream per connection.
+// The guest agent client is resolved through getClient on every dial, not captured here,
+// because the host agent replaces the client when the guest agent reconnects and a
+// listener created for an earlier client outlives it.
+func DialContextToGRPCTunnel(getClient func() *guestagentclient.GuestAgentClient) func(ctx context.Context, network, addr string) (net.Conn, error) {
 	// gvisor-tap-vsock's UDPProxy demultiplexes client connections internally based on their source address.
 	// It calls this dialer function only when it receives a datagram from a new, unrecognized client.
 	// For each new client, we must return a new net.Conn, which in our case is a new gRPC stream.
 	// The atomic counter ensures that each stream has a unique ID to distinguish them on the server side.
 	var connectionCounter atomic.Uint32
 	return func(_ context.Context, network, addr string) (net.Conn, error) {
+		client := getClient()
+		if client == nil {
+			return nil, fmt.Errorf("could not open tunnel for addr: %s: guest agent client is not available", addr)
+		}
 		// Passed context.Context is used for timeout on initiate connection, not for the lifetime of the connection.
 		// Use a dedicated context so CloseRead can unblock a goroutine waiting in stream.Recv().
 		streamCtx, cancel := context.WithCancel(context.Background())

pkg/portfwd/client_test.go

@@ -8,6 +8,7 @@ import (
 	"errors"
 	"io"
 	"net"
+	"sync"
 	"testing"
 	"time"
 
@@ -27,7 +28,7 @@ func TestDialContextToGRPCTunnelClosesStreamFromTCPProxy(t *testing.T) {
 	clientConn, proxyConn := net.Pipe()
 	proxyDone := make(chan struct{})
 	go func() {
-		HandleTCPConnection(t.Context(), DialContextToGRPCTunnel(client), proxyConn, "127.0.0.1:80")
+		HandleTCPConnection(t.Context(), DialContextToGRPCTunnel(func() *guestagentclient.GuestAgentClient { return client }), proxyConn, "127.0.0.1:80")
 		close(proxyDone)
 	}()
 
@@ -121,3 +122,30 @@ func (s *testGuestService) Tunnel(stream api.GuestService_TunnelServer) error {
 		return err
 	}
 }
+
+func TestDialContextToGRPCTunnelDialsThroughCurrentClient(t *testing.T) {
+	replaced := newTestGuestAgentClient(t, &testGuestService{tunnelDone: make(chan error, 1)})
+	current := newTestGuestAgentClient(t, &testGuestService{tunnelDone: make(chan error, 1)})
+	defer current.Close()
+
+	var mu sync.Mutex
+	client := replaced
+	dial := DialContextToGRPCTunnel(func() *guestagentclient.GuestAgentClient {
+		mu.Lock()
+		defer mu.Unlock()
+		return client
+	})
+
+	conn, err := dial(t.Context(), "tcp", "127.0.0.1:80")
+	assert.NilError(t, err)
+	assert.NilError(t, conn.Close())
+
+	assert.NilError(t, replaced.Close())
+	mu.Lock()
+	client = current
+	mu.Unlock()
+
+	conn, err = dial(t.Context(), "tcp", "127.0.0.1:80")
+	assert.NilError(t, err)
+	assert.NilError(t, conn.Close())
+}