server: fix data race on handshakeDeadline after the handshake

conn.go

@@ -3,16 +3,18 @@ package ssh
 import (
 	"context"
 	"net"
+	"sync/atomic"
 	"time"
 )
 
 type serverConn struct {
 	net.Conn
 
-	idleTimeout       time.Duration
-	handshakeDeadline time.Time
-	maxDeadline       time.Time
-	closeCanceler     context.CancelFunc
+	idleTimeout   time.Duration
+	maxDeadline   time.Time
+	closeCanceler context.CancelFunc
+
+	handshakeDeadline atomic.Pointer[time.Time] // during the handshake; nil once it completes
 }
 
 func (c *serverConn) Write(p []byte) (n int, err error) {
@@ -48,8 +50,9 @@ func (c *serverConn) Close() (err error) {
 func (c *serverConn) updateDeadline() {
 	deadline := c.maxDeadline
 
-	if !c.handshakeDeadline.IsZero() && (deadline.IsZero() || c.handshakeDeadline.Before(deadline)) {
-		deadline = c.handshakeDeadline
+	// handshakeDeadline is nil once the handshake completes.
+	if hsDeadline := c.handshakeDeadline.Load(); hsDeadline != nil && (deadline.IsZero() || hsDeadline.Before(deadline)) {
+		deadline = *hsDeadline
 	}
 
 	if c.idleTimeout > 0 {

server.go

@@ -298,7 +298,8 @@ func (srv *Server) HandleConn(newConn net.Conn) {
 		conn.maxDeadline = time.Now().Add(srv.MaxTimeout)
 	}
 	if srv.HandshakeTimeout > 0 {
-		conn.handshakeDeadline = time.Now().Add(srv.HandshakeTimeout)
+		handshakeDeadline := time.Now().Add(srv.HandshakeTimeout)
+		conn.handshakeDeadline.Store(&handshakeDeadline)
 	}
 	conn.updateDeadline()
 	defer conn.Close()
@@ -309,7 +310,7 @@ func (srv *Server) HandleConn(newConn net.Conn) {
 		}
 		return
 	}
-	conn.handshakeDeadline = time.Time{}
+	conn.handshakeDeadline.Store(nil)
 	conn.updateDeadline()
 	srv.trackConn(sshConn, true)
 	defer srv.trackConn(sshConn, false)

server_test.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"io"
 	"net"
+	"sync/atomic"
 	"testing"
 	"time"
 )
@@ -159,3 +160,72 @@ func TestServerHandshakeTimeout(t *testing.T) {
 		return
 	}
 }
+
+// TestServerConnHandshakeDeadlineRace checks that clearing the
+// handshake deadline is safe while other goroutines call
+// updateDeadline, which is what Read and Write do on every packet when
+// IdleTimeout is set. It requires -race to detect the race, and fails
+// if handshakeDeadline is a non-atomic *time.Time field.
+func TestServerConnHandshakeDeadlineRace(t *testing.T) {
+	c1, c2 := net.Pipe()
+	defer c1.Close()
+	defer c2.Close()
+	deadline := time.Now().Add(time.Minute)
+	conn := &serverConn{
+		Conn:        c1,
+		idleTimeout: time.Minute,
+	}
+	conn.handshakeDeadline.Store(&deadline)
+
+	stop := make(chan struct{})
+	var reads atomic.Int64
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for {
+			select {
+			case <-stop:
+				return
+			default:
+				// Read and Write call updateDeadline on every
+				// packet once idleTimeout is set.
+				conn.updateDeadline()
+				reads.Add(1)
+			}
+		}
+	}()
+
+	// Store nil, as HandleConn does after the handshake, until
+	// concurrent updateDeadline calls are under way.
+	for reads.Load() < 100 {
+		conn.handshakeDeadline.Store(nil)
+	}
+	close(stop)
+	<-done
+}
+
+// TestServerHandshakeTimeoutClearedAfterHandshake checks that the
+// handshake deadline does not apply after the handshake completes:
+// with HandshakeTimeout set, a session whose handler runs longer than
+// the grace period still completes. The host key is generated before
+// the server is started so the handshake itself stays well inside the
+// grace period.
+func TestServerHandshakeTimeoutClearedAfterHandshake(t *testing.T) {
+	signer, err := generateSigner()
+	if err != nil {
+		t.Fatal(err)
+	}
+	s := &Server{
+		HandshakeTimeout: time.Second,
+		IdleTimeout:      time.Minute,
+		Handler: func(s Session) {
+			time.Sleep(1500 * time.Millisecond)
+		},
+	}
+	s.AddHostKey(signer)
+	sess, _, cleanup := newTestSession(t, s, nil)
+	defer cleanup()
+	if err := sess.Run(""); err != nil {
+		t.Fatalf("session run: %v", err)
+	}
+}