diff --git a/core/providers/utils/idle_timeout_reader_test.go b/core/providers/utils/idle_timeout_reader_test.go index c2465729cfb..08f1e424e93 100644 --- a/core/providers/utils/idle_timeout_reader_test.go +++ b/core/providers/utils/idle_timeout_reader_test.go @@ -16,6 +16,41 @@ import ( // helpers // --------------------------------------------------------------------------- +// syncedPanicBody is a controlled io.ReadCloser used to deterministically +// reproduce the race in SetupStreamCancellation. +// +// Read blocks until Close() is called, then panics — simulating what fasthttp's +// streaming body does when the underlying TCP connection is force-closed. +// +// Close() triggers the panic (by closing panicTrigger) and then blocks on +// allowReturn until the test signals it. This keeps SetupStreamCancellation +// stuck inside Close() so BifrostContextKeyConnectionClosed is guaranteed to +// be unset when idleTimeoutReader.Read's recover block runs — the exact race +// window that exists in the current code. +type syncedPanicBody struct { + panicTrigger chan struct{} + allowReturn chan struct{} + closeOnce sync.Once +} + +func newSyncedPanicBody() *syncedPanicBody { + return &syncedPanicBody{ + panicTrigger: make(chan struct{}), + allowReturn: make(chan struct{}), + } +} + +func (s *syncedPanicBody) Read(_ []byte) (int, error) { + <-s.panicTrigger + panic("use of closed network connection") +} + +func (s *syncedPanicBody) Close() error { + s.closeOnce.Do(func() { close(s.panicTrigger) }) + <-s.allowReturn + return nil +} + var errPanicReader = errors.New("panic reader called") type panicReader struct{} @@ -387,3 +422,67 @@ func TestIdleTimeoutReader_CleanupWaitsForRunningTimerCallback(t *testing.T) { t.Fatal("cleanup did not return after timer callback finished") } } + +// TestSetupStreamCancellation_NoPanicOnCancelledContext reproduces the race where +// SetupStreamCancellation calls Close() before setting BifrostContextKeyConnectionClosed. +// +// Timeline with current (unfixed) code: +// 1. cancel() fires → SetupStreamCancellation goroutine calls closer.Close() +// 2. Close() closes panicTrigger → Read goroutine unblocks and panics +// 3. Close() blocks on allowReturn → flag is NOT yet set +// 4. idleTimeoutReader.Read recover: r.fired=false, connectionClosed=false → re-panics (BUG) +// +// Timeline after fix (set flag before Close()): +// 1. cancel() fires → SetupStreamCancellation sets flag → calls closer.Close() +// 2. Close() closes panicTrigger → Read goroutine unblocks and panics +// 3. idleTimeoutReader.Read recover: connectionClosed=true → returns ErrStreamClosed (OK) +func TestSetupStreamCancellation_NoPanicOnCancelledContext(t *testing.T) { + t.Parallel() + + body := newSyncedPanicBody() + goCtx, cancel := context.WithCancel(context.Background()) + bifrostCtx := schemas.NewBifrostContext(goCtx, time.Time{}) + + reader, cleanupReader := NewIdleTimeoutReader(body, body, time.Minute, bifrostCtx) + // Defers run LIFO: allowReturn first, then stopCancel, then cleanupReader. + // This order ensures Close() can return before we wait for the goroutine. + defer cleanupReader() + stopCancel := SetupStreamCancellation(bifrostCtx, body, getLogger()) + defer stopCancel() + defer close(body.allowReturn) + + type readResult struct { + err error + panicked any + } + resultCh := make(chan readResult, 1) + go func() { + var res readResult + defer func() { + if r := recover(); r != nil { + res.panicked = r + } + resultCh <- res + }() + buf := make([]byte, 1) + _, res.err = reader.Read(buf) + }() + + // Cancelling triggers SetupStreamCancellation → Close() → Read panics. + // Close() is blocked on allowReturn so the flag has not been set yet when + // the recover block in idleTimeoutReader.Read runs — worst-case ordering. + cancel() + + select { + case res := <-resultCh: + if res.panicked != nil { + t.Errorf("Read re-panicked (flag not set before Close — BUG): %v", res.panicked) + return + } + if !errors.Is(res.err, ErrStreamClosed) { + t.Errorf("expected ErrStreamClosed, got: %v", res.err) + } + case <-time.After(2 * time.Second): + t.Error("Read goroutine did not unblock after context cancellation") + } +} diff --git a/core/providers/utils/utils.go b/core/providers/utils/utils.go index 96cfd1c51bb..f2c351cf5fc 100644 --- a/core/providers/utils/utils.go +++ b/core/providers/utils/utils.go @@ -2054,15 +2054,15 @@ func SetupStreamCancellation(ctx *schemas.BifrostContext, bodyStream io.Reader, } // Context cancelled or deadline exceeded - close the body stream to unblock reads if closer, ok := bodyStream.(io.Closer); ok { + ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) if err := closer.Close(); err != nil { getLogger().Debug(fmt.Sprintf("Error closing body stream on context done: %v", err)) } - ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) } else if wce, ok := bodyStream.(streamCloserWithError); ok { + ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) if err := wce.CloseWithError(ctx.Err()); err != nil { getLogger().Debug(fmt.Sprintf("Error closing body stream on context done: %v", err)) } - ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) } case <-done: // Race between done and ctx.Done: the streaming goroutine has reached its defer @@ -2077,15 +2077,15 @@ func SetupStreamCancellation(ctx *schemas.BifrostContext, bodyStream io.Reader, return } if closer, ok := bodyStream.(io.Closer); ok { + ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) if err := closer.Close(); err != nil { getLogger().Debug(fmt.Sprintf("Error closing body stream on done with cancelled context: %v", err)) } - ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) } else if wce, ok := bodyStream.(streamCloserWithError); ok { + ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) if err := wce.CloseWithError(ctx.Err()); err != nil { getLogger().Debug(fmt.Sprintf("Error closing body stream on done with cancelled context: %v", err)) } - ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) } } } @@ -2236,7 +2236,7 @@ func (r *idleTimeoutReader) Read(p []byte) (n int, err error) { // Checking if stream is already closed if r.connectionClosed() { - return 0, nil + return 0, r.closedReadError() } n, err = r.reader.Read(p) if n > 0 {