[client] Close SSH connections still in the handshake on Stop and reject ones accepted by a stopped server (#8392)

* Close SSH connections still in the handshake on Stop and reject ones accepted by a stopped server

* Add a test that a connection dialed after Stop is refused

* [client] Clarify SSH server lifecycle and connection validation comments

* Clarify trackConn's doc comment on connection ownership after rejection

* Stop asserting a specific outcome of the Stop/Serve registration race in the post-Stop dial test

* Explicitly close the SSH listener on Stop and avoid logging the resulting race as an error

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

client/ssh/server/server.go

@@ -168,6 +168,10 @@ type Server struct {
 	// Populated at authentication time, stores JWT username and port forwards for status display.
 	connections map[connKey]*connState
 
+	// openConns tracks every accepted connection until it is closed.
+	// The SSH library only tracks a connection once it has authenticated, so Stop closes the others from here.
+	openConns map[*trackedConn]struct{}
+
 	allowLocalPortForwarding  bool
 	allowRemotePortForwarding bool
 	allowRootLogin            bool
@@ -223,7 +227,7 @@ type SessionInfo struct {
 }
 
 // New creates an SSH server instance with the provided host key and optional JWT configuration
-// If jwtConfig is nil, JWT authentication is disabled
+// If config.JWT is nil, JWT authentication is disabled.
 func New(config *Config) *Server {
 	s := &Server{
 		mu:                     sync.RWMutex{},
@@ -232,6 +236,7 @@ func New(config *Config) *Server {
 		pendingAuthJWT:         make(map[authKey]string),
 		remoteForwardListeners: make(map[forwardKey]net.Listener),
 		connections:            make(map[connKey]*connState),
+		openConns:              make(map[*trackedConn]struct{}),
 		jwtEnabled:             config.JWT != nil,
 		jwtConfig:              config.JWT,
 		authorizer:             sshauth.NewAuthorizer(),
@@ -274,12 +279,30 @@ func (s *Server) Start(ctx context.Context, addr netip.AddrPort) error {
 
 	go func() {
 		if err := sshServer.Serve(ln); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
-			log.Errorf("SSH server error: %v", err)
+			s.logServeError("SSH server error", sshServer, err)
 		}
 	}()
 	return nil
 }
 
+// logServeError logs an error returned by a gliderlabs/ssh Server's Serve call. Closing our
+// own listener explicitly as a Stop safeguard can race gliderlabs/ssh's own bookkeeping (it
+// resets its shutdown signal when a listener registers as the first one), so Serve can return
+// a raw "closed network connection" error instead of ssh.ErrServerClosed even on an intentional
+// stop. Log at Debug when srv is no longer the server's current one (Stop or a new Start already
+// ran), and at Error only when Serve failed while still expected to be running.
+func (s *Server) logServeError(msg string, srv *ssh.Server, err error) {
+	s.mu.RLock()
+	stopped := s.sshServer != srv
+	s.mu.RUnlock()
+
+	if stopped {
+		log.Debugf("%s: %v", msg, err)
+		return
+	}
+	log.Errorf("%s: %v", msg, err)
+}
+
 // AddListener starts serving SSH on an additional address (e.g. IPv6).
 // Must be called after Start.
 func (s *Server) AddListener(ctx context.Context, addr netip.AddrPort) error {
@@ -303,7 +326,7 @@ func (s *Server) AddListener(ctx context.Context, addr netip.AddrPort) error {
 
 	go func() {
 		if err := srv.Serve(ln); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
-			log.Errorf("SSH server error on %s: %v", addrDesc, err)
+			s.logServeError(fmt.Sprintf("SSH server error on %s", addrDesc), srv, err)
 		}
 	}()
 	return nil
@@ -336,7 +359,9 @@ func (s *Server) closeListener(ln net.Listener) {
 	}
 }
 
-// Stop closes the SSH server
+// Stop closes the SSH server, including connections still authenticating and
+// remote forwarding listeners, and clears session and connection state.
+// It does nothing if already stopped and always returns nil, even if closing fails.
 func (s *Server) Stop() error {
 	s.mu.Lock()
 	sshServer := s.sshServer
@@ -345,9 +370,11 @@ func (s *Server) Stop() error {
 		return nil
 	}
 	s.sshServer = nil
+	listener := s.listener
 	s.listener = nil
 	extraListeners := s.extraListeners
 	s.extraListeners = nil
+	openConns := maps.Keys(s.openConns)
 	s.mu.Unlock()
 
 	// Close outside the lock: session handlers need s.mu for unregisterSession.
@@ -355,6 +382,20 @@ func (s *Server) Stop() error {
 		log.Debugf("close SSH server: %v", err)
 	}
 
+	// sshServer.Close only closes listeners its Serve goroutine has already registered via
+	// trackListener. A fast Start/Stop flap can race that registration, leaving the listener
+	// open even though sshServer.Close returned; close it explicitly as a safeguard, the same
+	// way extraListeners are.
+	s.closeListener(listener)
+
+	// sshServer.Close only closes authenticated connections. One still in the handshake
+	// would otherwise stay open and be able to log in to the stopped server.
+	for _, conn := range openConns {
+		if err := conn.Close(); err != nil && !errors.Is(err, net.ErrClosed) {
+			log.Debugf("close SSH connection from %s: %v", conn.remoteAddr, err)
+		}
+	}
+
 	for _, ln := range extraListeners {
 		if err := ln.Close(); err != nil {
 			log.Debugf("close extra SSH listener: %v", err)
@@ -790,19 +831,47 @@ type trackedConn struct {
 	onceClose  sync.Once
 }
 
+// Close closes the underlying connection and removes its server tracking once,
+// even if closing fails. It returns the underlying connection's close error.
 func (c *trackedConn) Close() error {
 	err := c.Conn.Close()
 	c.onceClose.Do(func() {
-		c.server.handleConnectionClose(c.remoteAddr)
+		c.server.handleConnectionClose(c)
 	})
 	return err
 }
 
-func (s *Server) handleConnectionClose(remoteAddr string) {
+// trackConn wraps an accepted connection and registers it so that Stop can close it.
+// It returns nil if ctx identifies a stopped or replaced SSH server. The SSH server
+// closes the original connection when this callback returns nil.
+func (s *Server) trackConn(ctx ssh.Context, conn net.Conn) net.Conn {
+	tracked := &trackedConn{
+		Conn:       conn,
+		server:     s,
+		remoteAddr: conn.RemoteAddr().String(),
+	}
+
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	key := connKey(remoteAddr)
+	if s.sshServer == nil || ctx.Value(ssh.ContextKeyServer) != s.sshServer {
+		log.Debugf("SSH connection from %s rejected: server stopped", tracked.remoteAddr)
+		return nil
+	}
+
+	s.openConns[tracked] = struct{}{}
+	return tracked
+}
+
+// handleConnectionClose removes conn from shutdown tracking and removes the
+// connection status associated with its remote address.
+func (s *Server) handleConnectionClose(conn *trackedConn) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	delete(s.openConns, conn)
+
+	key := connKey(conn.remoteAddr)
 	state, exists := s.connections[key]
 	if exists && len(state.portForwards) > 0 {
 		s.connLogger(state).Info("port forwarding connection closed")
@@ -844,13 +913,17 @@ func (s *Server) findSessionKeyByContext(ctx ssh.Context) sessionKey {
 	return "unknown"
 }
 
-func (s *Server) connectionValidator(_ ssh.Context, conn net.Conn) net.Conn {
+// connectionValidator returns a tracked connection, or nil to reject it.
+// With a valid configured IPv4 address and network, only TCP peers in the NetBird
+// IPv4 or IPv6 subnet are allowed, excluding the server's own IPs. Otherwise,
+// network filtering is skipped. Stopped or replaced SSH servers are always rejected.
+func (s *Server) connectionValidator(ctx ssh.Context, conn net.Conn) net.Conn {
 	s.mu.RLock()
 	wgAddr := s.wgAddress
 	s.mu.RUnlock()
 
 	if !wgAddr.Network.IsValid() || !wgAddr.IP.IsValid() {
-		return conn
+		return s.trackConn(ctx, conn)
 	}
 
 	remoteAddr := conn.RemoteAddr()
@@ -881,11 +954,7 @@ func (s *Server) connectionValidator(_ ssh.Context, conn net.Conn) net.Conn {
 	}
 
 	log.Infof("SSH connection from NetBird peer %s allowed", tcpAddr)
-	return &trackedConn{
+	return s.trackConn(ctx, conn)
-		Conn:       conn,
-		server:     s,
-		remoteAddr: conn.RemoteAddr().String(),
-	}
 }
 
 func (s *Server) createSSHServer(addr net.Addr) (*ssh.Server, error) {

client/ssh/server/server_test.go

@@ -1,8 +1,10 @@
 package server
 
 import (
+	"bufio"
 	"context"
 	"fmt"
+	"io"
 	"net"
 	"net/netip"
 	"os/user"
@@ -31,6 +33,60 @@ func TestServer_StartStop(t *testing.T) {
 	assert.NoError(t, err)
 }
 
+func TestServer_StopClosesUnauthenticatedConnections(t *testing.T) {
+	key, err := nbssh.GeneratePrivateKey(nbssh.ED25519)
+	require.NoError(t, err)
+
+	server := New(&Config{HostKeyPEM: key})
+	serverAddr := StartTestServer(t, server)
+
+	conn, err := net.Dial("tcp", serverAddr)
+	require.NoError(t, err)
+	defer func() { _ = conn.Close() }()
+
+	// The server sends its version first, so reading it means the connection is being served.
+	_, err = bufio.NewReader(conn).ReadString('\n')
+	require.NoError(t, err)
+
+	require.NoError(t, server.Stop())
+
+	require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second)))
+	_, err = conn.Read(make([]byte, 1))
+	require.ErrorIs(t, err, io.EOF, "connection still in the handshake must be closed by Stop")
+
+	server.mu.RLock()
+	defer server.mu.RUnlock()
+	assert.Empty(t, server.openConns)
+}
+
+// TestServer_RejectsConnectionAfterStop verifies that a connection dialed after Stop has
+// returned is never served by the stopped server's stale state. Whether the listener was
+// already closed by Stop, or was still briefly open because Serve had not yet registered
+// it, is a scheduling detail outside the test's control: either outcome is acceptable, as
+// long as the connection is never served.
+func TestServer_RejectsConnectionAfterStop(t *testing.T) {
+	key, err := nbssh.GeneratePrivateKey(nbssh.ED25519)
+	require.NoError(t, err)
+
+	server := New(&Config{HostKeyPEM: key})
+	serverAddr := StartTestServer(t, server)
+	require.NoError(t, server.Stop())
+
+	conn, err := net.Dial("tcp", serverAddr)
+	if err != nil {
+		// The listener was already closed by Stop; the connection was correctly refused.
+		return
+	}
+	defer func() { _ = conn.Close() }()
+
+	// The listener was still open because Serve had not yet registered it with the
+	// embedded SSH server when Stop ran. trackConn must still refuse the connection
+	// once Serve's accept loop picks it up, since s.sshServer was already cleared.
+	require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second)))
+	_, err = conn.Read(make([]byte, 1))
+	require.ErrorIs(t, err, io.EOF, "a connection accepted after Stop must be closed, not served")
+}
+
 func TestSSHServerIntegration(t *testing.T) {
 	// Generate host key for server
 	hostKey, err := nbssh.GeneratePrivateKey(nbssh.ED25519)