drpcserver: bound the ServeOne TLS handshake with an optional timeout

drpcserver/server.go

@@ -45,6 +45,13 @@ type Options struct {
 	// restrictions. If it returns a non-nil error the connection is rejected.
 	TLSCipherRestrict func(conn net.Conn) error
 
+	// TLSHandshakeTimeout, if positive, bounds the TLS handshake performed in
+	// ServeOne by setting a deadline on the connection for the duration of the
+	// handshake. The deadline is cleared once the handshake completes. If zero,
+	// no deadline is set and a peer that stalls mid-handshake can hold its
+	// connection (and serving goroutine) open indefinitely.
+	TLSHandshakeTimeout time.Duration
+
 	// Metrics holds optional metrics the server will populate.
 	Metrics ServerMetrics
 }
@@ -142,7 +149,25 @@ func (s *Server) ServeOne(ctx context.Context, tr drpc.Transport) (err error) {
 		// interrupting any ongoing communication. Even if we didn't call it
 		// explicitly, the first read/write operation would call it internally
 		// anyway.
+		//
+		// Bound the handshake with a deadline so a peer that stalls
+		// mid-handshake cannot hold the connection open indefinitely, then
+		// clear it so it does not affect reads and writes once the connection is
+		// serving requests.
+		//
+		// The reset cannot be deferred: ServeOne serves the connection for its
+		// entire lifetime after the handshake, so a deferred reset would leave
+		// the handshake deadline in force for every request. Clearing to the
+		// zero time (rather than a prior deadline) is correct: there is no API
+		// to read a pre-existing deadline, and the accept path sets none.
+		handshakeTimeout := s.opts.TLSHandshakeTimeout
+		if handshakeTimeout > 0 {
+			_ = tlsConn.SetDeadline(time.Now().Add(handshakeTimeout))
+		}
 		err := tlsConn.HandshakeContext(ctx)
+		if handshakeTimeout > 0 {
+			_ = tlsConn.SetDeadline(time.Time{})
+		}
 		if err != nil {
 			s.recordTLSHandshakeError()
 			return drpc.ConnectionError.New("server handshake [%q] failed: %w", tlsConn.RemoteAddr(), err)

drpcserver/server_test.go

@@ -4,8 +4,12 @@
 package drpcserver
 
 import (
+	"context"
+	"crypto/tls"
 	"net"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/zeebo/assert"
 
@@ -35,6 +39,77 @@ func TestServerTemporarySleep(t *testing.T) {
 	assert.NoError(t, New(nil).Serve(ctx, l))
 }
 
+// TestServeOneTLSHandshakeTimeout verifies that TLSHandshakeTimeout bounds the
+// handshake ServeOne drives, and that a zero timeout sets no deadline.
+func TestServeOneTLSHandshakeTimeout(t *testing.T) {
+	t.Run("times out a stalled handshake", func(t *testing.T) {
+		ctx := drpctest.NewTracker(t)
+		defer ctx.Close()
+
+		// The client end never writes, so the server handshake blocks reading
+		// the ClientHello until the deadline fires.
+		srv, clt := net.Pipe()
+		defer func() { _ = srv.Close() }()
+		defer func() { _ = clt.Close() }()
+
+		s := NewWithOptions(nil, Options{TLSHandshakeTimeout: 50 * time.Millisecond})
+
+		done := make(chan error, 1)
+		ctx.Run(func(ctx context.Context) {
+			done <- s.ServeOne(ctx, tls.Server(srv, &tls.Config{}))
+		})
+
+		select {
+		case err := <-done:
+			assert.Error(t, err) // handshake deadline fired
+		case <-time.After(30 * time.Second):
+			t.Fatal("ServeOne did not return; the handshake was not bounded")
+		}
+	})
+
+	t.Run("no deadline when unset", func(t *testing.T) {
+		ctx := drpctest.NewTracker(t)
+		defer ctx.Close()
+
+		srv, clt := net.Pipe()
+		defer func() { _ = srv.Close() }()
+		defer func() { _ = clt.Close() }()
+
+		rec := &deadlineRecorder{Conn: srv}
+		s := NewWithOptions(nil, Options{}) // zero timeout
+
+		done := make(chan error, 1)
+		ctx.Run(func(ctx context.Context) {
+			done <- s.ServeOne(ctx, tls.Server(rec, &tls.Config{}))
+		})
+
+		// Let ServeOne reach the (blocking) handshake, then confirm no deadline
+		// was set. Cancel to unblock so the goroutine exits.
+		time.Sleep(50 * time.Millisecond)
+		assert.Equal(t, rec.setDeadlineCalls(), int32(0))
+		ctx.Cancel() // HandshakeContext observes cancellation and returns
+
+		select {
+		case <-done:
+		case <-time.After(30 * time.Second):
+			t.Fatal("ServeOne did not return after context cancellation")
+		}
+	})
+}
+
+// deadlineRecorder wraps a net.Conn and counts SetDeadline calls.
+type deadlineRecorder struct {
+	net.Conn
+	calls atomic.Int32
+}
+
+func (r *deadlineRecorder) SetDeadline(t time.Time) error {
+	r.calls.Add(1)
+	return r.Conn.SetDeadline(t)
+}
+
+func (r *deadlineRecorder) setDeadlineCalls() int32 { return r.calls.Load() }
+
 type listener func() (net.Conn, error)
 
 func (l listener) Accept() (net.Conn, error) { return l() }