From 4ce00ea521d22dff010223fa0b7c276234dc63c2 Mon Sep 17 00:00:00 2001 From: Paul-Julien Vauthier Date: Wed, 26 Aug 2026 17:24:05 -0400 Subject: [PATCH 1/3] websocket: make the input dial honor context cancellation getConn received a context but discarded it: dialer.Dial always used context.Background, so a peer that accepted the TCP connection and never completed the HTTP upgrade handshake held the input for gorilla's 45s default handshake timeout. AsyncReader has no other way to stop the reader than the context it passes to Connect, so this could outlast the Kubernetes default terminationGracePeriodSeconds (30s) and turn a graceful shutdown into a kill. Add a dialContext helper that keeps a handle on the raw connection via NetDialContext and uses context.AfterFunc to close it once the outer context is done, waking a handshake read that gorilla's own timeout wouldn't catch in time. The resulting error is mapped back to ctx.Err so callers can tell a shutdown from a flaky peer. The helper hides the context's deadline from gorilla so gorilla's own connection deadline doesn't race the watcher. dialContext and its test helpers are free functions with no reader state, so they live in shared websocket.go/websocket_test.go files rather than the input-only ones, ready for the output side to reuse. --- CHANGELOG.md | 7 ++ internal/impl/io/input_websocket.go | 8 +- internal/impl/io/input_websocket_test.go | 106 +++++++++++++++++++++++ internal/impl/io/websocket.go | 69 +++++++++++++++ internal/impl/io/websocket_test.go | 103 ++++++++++++++++++++++ 5 files changed, 288 insertions(+), 5 deletions(-) create mode 100644 internal/impl/io/websocket.go create mode 100644 internal/impl/io/websocket_test.go diff --git a/CHANGELOG.md b/CHANGELOG.md index cf97255dc..107a39ab6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,13 @@ Changelog All notable changes to this project will be documented in this file. +## Unreleased + +### Fixed + +- Input `websocket`: Connection attempts (dial and upgrade handshake) now honor context cancellation. Previously, unresponsive servers could + block graceful shutdown for up to the default 45s handshake timeout. (@Leward) + ## 4.78.0 - 2026-08-20 ### Added diff --git a/internal/impl/io/input_websocket.go b/internal/impl/io/input_websocket.go index 2982fb94f..0e9af925a 100644 --- a/internal/impl/io/input_websocket.go +++ b/internal/impl/io/input_websocket.go @@ -185,13 +185,11 @@ func (w *websocketReader) getConn(ctx context.Context) (*websocket.Conn, error) if w.proxyURLParsed != nil { dialer.Proxy = http.ProxyURL(w.proxyURLParsed) } - if w.tlsEnabled { dialer.TLSClientConfig = w.tlsConf - if client, res, err = dialer.Dial(w.urlStr, headers); err != nil { - return nil, err - } - } else if client, res, err = dialer.Dial(w.urlStr, headers); err != nil { + } + + if client, res, err = dialContext(ctx, dialer, w.urlStr, headers); err != nil { return nil, err } diff --git a/internal/impl/io/input_websocket_test.go b/internal/impl/io/input_websocket_test.go index 187ffc261..810020e9e 100644 --- a/internal/impl/io/input_websocket_test.go +++ b/internal/impl/io/input_websocket_test.go @@ -6,6 +6,7 @@ import ( "bytes" "context" "fmt" + "net" "net/http" "net/http/httptest" "net/url" @@ -304,3 +305,108 @@ url: %v wg.Wait() close(closeChan) } + +// newWebsocketReader returns a websocketReader configured to connect to the websocket server listening on addr. +func newWebsocketReader(t *testing.T, addr net.Addr) *websocketReader { + t.Helper() + + pConf, err := websocketInputSpec().ParseYAML("url: ws://"+addr.String()+"\n", nil) + require.NoError(t, err) + + m, err := newWebsocketReaderFromParsed(pConf, mock.NewManager()) + require.NoError(t, err) + + return m +} + +// newHangingWebsocketReader returns a reader pointed at a listener that accepts +// TCP connections but never answers the handshake, along with the accept channel +// of that listener. +func newHangingWebsocketReader(t *testing.T) (*websocketReader, <-chan struct{}) { + t.Helper() + + addr, accepted := newHangingListener(t) + return newWebsocketReader(t, addr), accepted +} + +// newUnreachableWebsocketReader returns a reader pointed at a closed port. A dial +// there fails immediately, so a context error proves no dial was attempted. +func newUnreachableWebsocketReader(t *testing.T) *websocketReader { + t.Helper() + + return newWebsocketReader(t, newUnreachableAddr(t)) +} + +// TestWebsocketConnectContextDone tests that Connect reports the context error +// when the context is done before or during the websocket handshake. +func TestWebsocketConnectContextDone(t *testing.T) { + tests := []struct { + name string + // setUp returns a reader and a context that is already done, or that becomes + // done while the dial waits for the handshake response. It also returns the + // accept channel to check after Connect returns, or nil for a case where no + // accept is expected. + setUp func(*testing.T) (*websocketReader, context.Context, <-chan struct{}) + wantErr error + }{ + { + name: "canceled before dialing", + setUp: func(t *testing.T) (*websocketReader, context.Context, <-chan struct{}) { + ctx, cancel := context.WithCancel(t.Context()) + cancel() + // The port is closed, so a dial fails at once. Only a context check + // before the dial can give context.Canceled here. + return newUnreachableWebsocketReader(t), ctx, nil + }, + wantErr: context.Canceled, + }, + { + name: "canceled while dialing", + setUp: func(t *testing.T) (*websocketReader, context.Context, <-chan struct{}) { + m, accepted := newHangingWebsocketReader(t) + + ctx, cancel := context.WithCancel(t.Context()) + t.Cleanup(cancel) + // Cancel from the accept, not from a timer, so the dial always waits + // for the handshake response when the context becomes done. + go func() { + <-accepted + cancel() + }() + + return m, ctx, nil + }, + wantErr: context.Canceled, + }, + { + name: "deadline exceeded while dialing", + setUp: func(t *testing.T) (*websocketReader, context.Context, <-chan struct{}) { + m, accepted := newHangingWebsocketReader(t) + + // A context carries its deadline from the start, so this case cannot + // take its trigger from the accept. The accept check after Connect + // returns proves the deadline expired during the handshake. + ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) + t.Cleanup(cancel) + + return m, ctx, accepted + }, + wantErr: context.DeadlineExceeded, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + m, ctx, accepted := test.setUp(t) + + err := awaitWithin(t, 500*time.Millisecond, "Connect", func() error { + return m.Connect(ctx) + }) + require.ErrorIs(t, err, test.wantErr) + + if accepted != nil { + requireAccepted(t, 5*time.Second, accepted) + } + }) + } +} diff --git a/internal/impl/io/websocket.go b/internal/impl/io/websocket.go new file mode 100644 index 000000000..59c8428c3 --- /dev/null +++ b/internal/impl/io/websocket.go @@ -0,0 +1,69 @@ +// Copyright 2025 Redpanda Data, Inc. + +package io + +import ( + "context" + "net" + "net/http" + "time" + + "github.com/gorilla/websocket" +) + +// hiddenDeadlineContext strips ctx.Deadline() while preserving cancellation. +// +// Gorilla sets a socket deadline matching ctx.Deadline(). Because the runtime's +// socket deadline timer and Go's context timer run independently, a socket read +// can time out (returning a generic "i/o timeout") slightly before ctx.Err() is set. +// +// Concealing the deadline prevents Gorilla from setting this competing socket timer. +// This ensures our context.AfterFunc watcher exclusively controls socket interruption +// and deterministically returns context errors. +type hiddenDeadlineContext struct { + context.Context +} + +func (hiddenDeadlineContext) Deadline() (time.Time, bool) { + return time.Time{}, false +} + +// dialContext is dialer.DialContext plus cancellation of the HTTP upgrade exchange. +// +// Gorilla applies the context to the TCP and TLS handshakes only. +// It then performs the websocket upgrade on a bare connection, which does not respond to context cancellation. +// A peer that accepts TCP and then stays silent blocks the dial until the handshake timeout. +func dialContext(ctx context.Context, dialer websocket.Dialer, urlStr string, headers http.Header) (*websocket.Conn, *http.Response, error) { + netDialer := &net.Dialer{} + var stop func() bool + + // 1. Intercept network connection creation to attach a context watcher (context.AfterFunc) that closes the conn once ctx is done. + dialer.NetDialContext = func(dialCtx context.Context, network, addr string) (net.Conn, error) { + conn, err := netDialer.DialContext(dialCtx, network, addr) + if err != nil { + return nil, err + } + // Watch ctx, not dialCtx: gorilla wraps dialCtx with HandshakeTimeout and + // cancels it as DialContext returns, which would fire this on success. + stop = context.AfterFunc(ctx, func() { _ = conn.Close() }) + return conn, nil + } + + // 2. Perform dial with hidden deadline context. + // The context we give gorilla keeps the cancellation of ctx (stopping TCP/TLS + // handshakes), but conceals ctx.Deadline() so gorilla does not set a competing + // socket timer that could race with ctx cancellation. + client, res, err := dialer.DialContext(hiddenDeadlineContext{ctx}, urlStr, headers) + + // 3. Clean up and prioritize context error returns. + if stop != nil { + _ = stop() + } + if ctxErr := ctx.Err(); ctxErr != nil { + if client != nil { + _ = client.Close() + } + return nil, res, ctxErr + } + return client, res, err +} diff --git a/internal/impl/io/websocket_test.go b/internal/impl/io/websocket_test.go new file mode 100644 index 000000000..4fafba94c --- /dev/null +++ b/internal/impl/io/websocket_test.go @@ -0,0 +1,103 @@ +// Copyright 2025 Redpanda Data, Inc. + +package io + +import ( + "net" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// newHangingListener returns the address of a listener that accepts TCP +// connections but never answers the handshake, so a dial there stays blocked +// until the dialer times out. The returned channel reports each accept, which +// tells the caller that a handshake is in flight. +func newHangingListener(t *testing.T) (net.Addr, <-chan struct{}) { + t.Helper() + + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + accepted := make(chan struct{}, 1) + var ( + mut sync.Mutex + conns []net.Conn + ) + + go func() { + for { + conn, err := lis.Accept() + if err != nil { + return + } + mut.Lock() + conns = append(conns, conn) + mut.Unlock() + select { + case accepted <- struct{}{}: + default: + } + } + }() + + // Close the accepted connections as well as the listener, so that an + // abandoned dial fails instead of running to the handshake timeout. + t.Cleanup(func() { + _ = lis.Close() + mut.Lock() + defer mut.Unlock() + for _, conn := range conns { + _ = conn.Close() + } + }) + + return lis.Addr(), accepted +} + +// newUnreachableAddr returns the address of a closed port. A dial there fails +// immediately, so a context error proves no dial was attempted. +func newUnreachableAddr(t *testing.T) net.Addr { + t.Helper() + + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + addr := lis.Addr() + require.NoError(t, lis.Close()) + + return addr +} + +// awaitWithin runs fn in the background and returns its error, or fails the test +// if fn has not returned within d. +func awaitWithin(t *testing.T, d time.Duration, what string, fn func() error) error { + t.Helper() + + done := make(chan error, 1) + go func() { + done <- fn() + }() + + select { + case err := <-done: + return err + case <-time.After(d): + t.Fatalf("%v stayed blocked for %v", what, d) + return nil + } +} + +// requireAccepted fails the test if the listener reported no accept within d. It +// proves that the dial reached the handshake, so a test that expects the context +// to become mid-handshake cannot pass through an earlier failure instead. +func requireAccepted(t *testing.T, d time.Duration, accepted <-chan struct{}) { + t.Helper() + + select { + case <-accepted: + case <-time.After(d): + t.Fatal("the listener accepted no connection, so the dial never reached the handshake") + } +} From e55e513a85348e6aa0856c59b509c9cd6e9d83e2 Mon Sep 17 00:00:00 2001 From: Paul-Julien Vauthier Date: Tue, 1 Sep 2026 14:01:01 -0400 Subject: [PATCH 2/3] Set copyright header to current year Co-authored-by: Joseph Woodward --- internal/impl/io/websocket.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/impl/io/websocket.go b/internal/impl/io/websocket.go index 59c8428c3..b93059318 100644 --- a/internal/impl/io/websocket.go +++ b/internal/impl/io/websocket.go @@ -1,4 +1,4 @@ -// Copyright 2025 Redpanda Data, Inc. +// Copyright 2026 Redpanda Data, Inc. package io From f47750c5102b0d6a335e73cfc70af54b15b5a750 Mon Sep 17 00:00:00 2001 From: Paul-Julien Vauthier Date: Tue, 1 Sep 2026 14:01:10 -0400 Subject: [PATCH 3/3] Set copyright header to current year Co-authored-by: Joseph Woodward --- internal/impl/io/websocket_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/impl/io/websocket_test.go b/internal/impl/io/websocket_test.go index 4fafba94c..ef33b7c97 100644 --- a/internal/impl/io/websocket_test.go +++ b/internal/impl/io/websocket_test.go @@ -1,4 +1,4 @@ -// Copyright 2025 Redpanda Data, Inc. +// Copyright 2026 Redpanda Data, Inc. package io