From 202df95282208067cf3556a2fecfd2bdd8e72600 Mon Sep 17 00:00:00 2001 From: Rafi Shamim Date: Wed, 22 Jul 2026 13:10:38 -0400 Subject: [PATCH] drpcserver: bound the ServeOne TLS handshake with an optional timeout ServeOne performs the TLS handshake explicitly so it can populate the peer connection info before serving requests. That handshake had no deadline, so a peer that completes the TCP connection but then stalls mid-handshake (never sending its ClientHello flight) would pin the serving goroutine and its file descriptor indefinitely. Add an Options.TLSHandshakeTimeout field. When positive, ServeOne sets a deadline on the connection for the duration of the handshake and clears it once the handshake returns, so it does not affect subsequent reads and writes on the established connection. A zero value preserves the existing behavior (no deadline), keeping the change backward compatible. Informs: https://github.com/cockroachdb/cockroach/issues/144754 Co-Authored-By: roachdev-claude --- drpcserver/server.go | 25 +++++++++++++ drpcserver/server_test.go | 75 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+) diff --git a/drpcserver/server.go b/drpcserver/server.go index 83c11525..04690005 100644 --- a/drpcserver/server.go +++ b/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) diff --git a/drpcserver/server_test.go b/drpcserver/server_test.go index 968d7cea..205ce2a5 100644 --- a/drpcserver/server_test.go +++ b/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() }