diff --git a/pkg/hostagent/hostagent.go b/pkg/hostagent/hostagent.go index be15e2c6..77e46d77 100644 --- a/pkg/hostagent/hostagent.go +++ b/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 diff --git a/pkg/portfwd/client.go b/pkg/portfwd/client.go index dcb40ad2..c66279a0 100644 --- a/pkg/portfwd/client.go +++ b/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()) diff --git a/pkg/portfwd/client_test.go b/pkg/portfwd/client_test.go index f465856c..3e348297 100644 --- a/pkg/portfwd/client_test.go +++ b/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()) +}