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() }