Deep Review: 20260426-165750-pr-4889

	}
	// The previous client (if any) is unreachable: close its underlying gRPC
	// ClientConn before replacing it, otherwise the resolver/transport
	// goroutines and the dialed net.Conn (e.g. forwarded ga.sock) leak across
	// every guest agent restart or VM reboot.
	if a.client != nil {
		if err := a.client.Close(); err != nil {
			logrus.WithError(err).Debug("failed to close stale guest agent client")
		}
		a.client = nil
	}
	var err error
	a.client, err = guestagentclient.NewGuestAgentClient(a.createConnection)
	return a.client, err
}

processGuestAgentEvents builds the dialer once per loop iteration as dialContext := portfwd.DialContextToGRPCTunnel(client) and hands it to a.grpcPortForwarder.OnEvent (pkg/hostagent/hostagent.go:898-899). DialContextToGRPCTunnel captures client in a closure (pkg/portfwd/client.go:46-67); each future TCP/UDP accept calls client.Tunnel(...) against that captured pointer.

			}
		}
		if useSSHFwd {
			a.portForwarder.OnEvent(ctx, ev)
		} else {
			dialContext := portfwd.DialContextToGRPCTunnel(client)
			a.grpcPortForwarder.OnEvent(ctx, dialContext, ev)
		}
	}

	if err := client.Events(ctx, onEvent); err != nil {
		if status.Code(err) == codes.Canceled {

	}()
	proxy.Run()
	logrus.Debugf("udp proxy for guestAddr: %s closed", guestAddr)
}

func DialContextToGRPCTunnel(client *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) {
		// Passed context.Context is used for timeout on initiate connection, not for the lifetime of the connection.
		// We use context.Background() here to avoid unexpected cancellation.
		stream, err := client.Tunnel(context.Background())
		if err != nil {
			return nil, fmt.Errorf("could not open tunnel for addr: %s error:%w", addr, err)
		}
		// Handshake message to start tunnel
		id := fmt.Sprintf("%s-%s-%d", network, addr, connectionCounter.Add(1))
		if err := stream.Send(&api.TunnelMessage{Id: id, Protocol: network, GuestAddr: addr}); err != nil {
			return nil, fmt.Errorf("could not start tunnel for id: %s addr: %s error:%w", id, addr, err)
		}
		rw := &GrpcClientRW{stream: stream, id: id, addr: addr, protocol: network}
		return rw, nil
	}
}

type GrpcClientRW struct {
	id   string
	addr string

ClosableListeners.forwardTCP keys the listener map by tcp/host/guest and early-returns when the same key already exists, so the dialer is not refreshed on a follow-up event:

_, ok := p.listeners[key]
if ok {
    p.listenersRW.Unlock()
    return
}

(pkg/portfwd/listener.go:99-103)

}

func (p *ClosableListeners) forwardTCP(ctx context.Context, dialContext func(ctx context.Context, network string, addr string) (net.Conn, error), hostAddress, guestAddress string) {
	key := key("tcp", hostAddress, guestAddress)

	p.listenersRW.Lock()
	_, ok := p.listeners[key]
	if ok {
		p.listenersRW.Unlock()
		return
	}
	tcpLis, err := Listen(ctx, p.listenConfig, hostAddress)
	if err != nil {
		logListenError(err, "tcp", hostAddress)
		p.listenersRW.Unlock()

After a guest-agent restart with an unchanged guest port set, the existing host listener keeps its old dialContext; new TCP accepts call Tunnel() on the now-Close()d ClientConn and fail immediately with ErrClientConnClosing. Pre-PR the same scenario worked because gRPC's automatic reconnect (via the WithContextDialer that re-runs a.createConnection) re-dialed the forwarded ga.sock on transport failure — the cost was the leak this PR fixes.

Fix: have the dialer resolve the current client on every dial, so existing listeners pick up new clients automatically:

dialContext := func(ctx context.Context, network, addr string) (net.Conn, error) {
    client, err := a.getOrCreateClient(ctx)
    if err != nil {
        return nil, err
    }
    return portfwd.DialContextToGRPCTunnel(client)(ctx, network, addr)
}
a.grpcPortForwarder.OnEvent(ctx, dialContext, ev)

The same problem applies more weakly to startInotify (pkg/hostagent/inotify.go:34-38), which captures the client once and never re-fetches. That stream is bound to a single ctx and dies naturally on guest-agent restart, so it is less urgent, but the re-fetch pattern would close that gap too.

		return errors.Join(errs...)
	})

	go func() {
		if a.instConfig.MountInotify != nil && *a.instConfig.MountInotify {
			if a.client == nil || !isGuestAgentSocketAccessible(ctx, a.client) {
				if a.driver.ForwardGuestAgent() {
					sshAddress, sshPort := a.sshAddressPort()
					_ = forwardSSH(ctx, a.sshConfig, sshAddress, sshPort, localUnix, remoteUnix, verbForward, false)
				}
			}

(pkg/hostagent/hostagent.go:758)

		a.client = nil
		return err
	})

	for {
		if a.client == nil || !isGuestAgentSocketAccessible(ctx, a.client) {
			if a.driver.ForwardGuestAgent() {
				sshAddress, sshPort := a.sshAddressPort()
				_ = forwardSSH(ctx, a.sshConfig, sshAddress, sshPort, localUnix, remoteUnix, verbForward, false)
			}
		}

Three goroutines call getOrCreateClient concurrently: watchGuestAgentEvents (hostagent.go:764), startInotify (pkg/hostagent/inotify.go:34), and the timesync loop (pkg/hostagent/timesync.go:49). Inside getOrCreateClient the new code at lines 843-847 explicitly nils a.client between Close() and the new NewGuestAgentClient assignment — a window of microseconds inside the locked region.

	mountWatchCh := make(chan notify.EventInfo, 128)
	err := a.setupWatchers(mountWatchCh)
	if err != nil {
		return err
	}
	client, err := a.getOrCreateClient(ctx)
	if err != nil {
		logrus.WithError(err).Error("failed to create client for inotify")
	}
	inotifyClient, err := client.Inotify(ctx)
	if err != nil {

		}
	}
}

func (a *HostAgent) syncTimeOnce(ctx context.Context) {
	client, err := a.getOrCreateClient(ctx)
	if err != nil {
		logrus.WithError(err).Debug("Time sync: failed to get client")
		return
	}

	}
	// The previous client (if any) is unreachable: close its underlying gRPC
	// ClientConn before replacing it, otherwise the resolver/transport
	// goroutines and the dialed net.Conn (e.g. forwarded ga.sock) leak across
	// every guest agent restart or VM reboot.
	if a.client != nil {
		if err := a.client.Close(); err != nil {
			logrus.WithError(err).Debug("failed to close stale guest agent client")
		}
		a.client = nil
	}
	var err error
	a.client, err = guestagentclient.NewGuestAgentClient(a.createConnection)
	return a.client, err
}

The watcher loop at lines 731 and 758 reads a.client twice in the short-circuit || expression without taking clientMu. If goroutine A enters that window while goroutine B's loop is mid- expression, the first read can observe the prior pointer (a.client == nil returns false) and the second read can observe the in-window nil. isGuestAgentSocketAccessible(ctx, nil) calls client.Info(ctx) → c.cli.GetInfo(...), which dereferences nil and panics, taking the hostagent down.

Pre-PR this race already existed for line 731 (the inotify goroutine), but a.client was never assigned nil between successful constructions, so the second read could only observe a stale-but- non-nil pointer — a benign data race. The new a.client = nil in the Close path turns it into a crashable race.

Fix: serialize every read of a.client through the same mutex. A small accessor avoids touching the loop body shape:

func (a *HostAgent) getClient() *guestagentclient.GuestAgentClient {
    a.clientMu.RLock()
    defer a.clientMu.RUnlock()
    return a.client
}

Then in watchGuestAgentEvents:

-        if a.client == nil || !isGuestAgentSocketAccessible(ctx, a.client) {
+        client := a.getClient()
+        if client == nil || !isGuestAgentSocketAccessible(ctx, client) {

Apply the same change at line 731 (inside the inotify-startup goroutine). Note: keeping a.client = nil in getOrCreateClient is fine once readers go through the accessor — the assignment at line 850 will repopulate it before the lock is released, but the explicit nil makes the failure mode clear if NewGuestAgentClient itself returns an error.

Severity calibration: Codex flagged this as Important; Gemini escalated to Critical citing the panic. We keep it at Important because the race window is narrow (Close + re-assign happen back-to- back inside one critical section) and the second read at line 758 is the same goroutine that calls getOrCreateClient, so only the inotify/timesync interleaving exposes it. The crash is real, but not deterministic.

		}
	}()

	// ensure close before ctx is cancelled
	a.cleanUp(a.grpcPortForwarder.Close)
	a.cleanUp(func() error {
		a.clientMu.Lock()
		defer a.clientMu.Unlock()
		if a.client == nil {
			return nil
		}
		err := a.client.Close()
		a.client = nil
		return err
	})

	for {
		if a.client == nil || !isGuestAgentSocketAccessible(ctx, a.client) {
			if a.driver.ForwardGuestAgent() {
				sshAddress, sshPort := a.sshAddressPort()

		logrus.Infof("Driver stopped due to error: %q", driverErr)
	case sig := <-a.signalCh:
		logrus.Infof("Received %s, shutting down the host agent", osutil.SignalName(sig))
	}
	// close the host agent routines before cancelling the context
	if closeErr := a.close(); closeErr != nil {
		logrus.WithError(closeErr).Warn("an error during shutting down the host agent")
	}
	cancelHA()
	return a.driver.Stop(ctx)
}

func (a *HostAgent) Info(_ context.Context) (*hostagentapi.Info, error) {
	info := &hostagentapi.Info{

		cli:  client,
		conn: clientConn,
	}, nil
}

// Close releases the underlying gRPC connection. It is safe to call on a nil
// receiver or a client whose connection has already been closed.
func (c *GuestAgentClient) Close() error {
	if c == nil || c.conn == nil {
		return nil
	}
	return c.conn.Close()
}

func (c *GuestAgentClient) Info(ctx context.Context) (*api.Info, error) {
	return c.cli.GetInfo(ctx, &emptypb.Empty{})
}

 func (c *GuestAgentClient) Close() error {
     if c == nil || c.conn == nil {
         return nil
     }
-    return c.conn.Close()
+    conn := c.conn
+    c.conn = nil
+    return conn.Close()
 }