diff --git a/context/NContext.go b/context/NContext.go index 876e268c..6bd93b76 100644 --- a/context/NContext.go +++ b/context/NContext.go @@ -38,7 +38,7 @@ type threeValueCtx[T1 any, T2 any, T3 any] struct { // and executes that function. if context is cancelled before function returns // it will return context error otherwise it will return nil func ExecFunc(ctx context.Context, fn func()) error { - ch := make(chan struct{}) + ch := make(chan struct{}, 1) // buffered so the worker can send-and-exit if ctx is cancelled first go func() { fn() ch <- struct{}{} @@ -56,7 +56,7 @@ func ExecFunc(ctx context.Context, fn func()) error { // if context is cancelled before function returns it will return context error // otherwise it will return function's return values func ExecFuncWithTwoReturns[T1 any](ctx context.Context, fn func() (T1, error)) (T1, error) { - ch := make(chan twoValueCtx[T1, error]) + ch := make(chan twoValueCtx[T1, error], 1) // buffered so the worker can send-and-exit if ctx is cancelled first go func() { x, y := fn() ch <- twoValueCtx[T1, error]{var1: x, var2: y} @@ -75,7 +75,7 @@ func ExecFuncWithTwoReturns[T1 any](ctx context.Context, fn func() (T1, error)) // if context is cancelled before function returns it will return context error // otherwise it will return function's return values func ExecFuncWithThreeReturns[T1 any, T2 any](ctx context.Context, fn func() (T1, T2, error)) (T1, T2, error) { - ch := make(chan threeValueCtx[T1, T2, error]) + ch := make(chan threeValueCtx[T1, T2, error], 1) // buffered so the worker can send-and-exit if ctx is cancelled first go func() { x, y, z := fn() ch <- threeValueCtx[T1, T2, error]{var1: x, var2: y, var3: z} diff --git a/context/leak_test.go b/context/leak_test.go new file mode 100644 index 00000000..700337ad --- /dev/null +++ b/context/leak_test.go @@ -0,0 +1,88 @@ +package contextutil_test + +import ( + "context" + "runtime" + "testing" + "time" + + contextutil "github.com/projectdiscovery/utils/context" +) + +// When the context is cancelled before fn returns, ExecFunc* returns the +// context error but the worker goroutine keeps running fn. Once fn finishes it +// sends its result on the internal channel. If that channel is unbuffered and +// the caller has already returned via ctx.Done(), nobody ever receives, so the +// worker blocks forever on the send — one leaked goroutine per cancelled call. +// +// These tests drive many cancelled calls whose fn outlives the context, then +// assert the goroutine count returns to baseline. They fail on an unbuffered +// result channel and pass once it is buffered (cap 1) so the worker can always +// send-and-exit. + +const leakCalls = 50 + +// assertNoGoroutineLeak polls (workers finish and exit asynchronously) until the +// live goroutine count returns to ~base, failing if it does not. +func assertNoGoroutineLeak(t *testing.T, base int) { + t.Helper() + deadline := time.Now().Add(5 * time.Second) + for { + runtime.GC() + leaked := runtime.NumGoroutine() - base + if leaked <= 2 { // tolerance for transient runtime/test goroutines + return + } + if time.Now().After(deadline) { + t.Fatalf("goroutine leak: ~%d goroutines still alive after %d context-cancelled calls (base=%d, now=%d)", + leaked, leakCalls, base, runtime.NumGoroutine()) + } + time.Sleep(25 * time.Millisecond) + } +} + +// baseline settles outstanding goroutines then records the count. +func baseline() int { + runtime.GC() + time.Sleep(50 * time.Millisecond) + runtime.GC() + return runtime.NumGoroutine() +} + +func TestExecFunc_NoGoroutineLeakOnCancel(t *testing.T) { + base := baseline() + for range leakCalls { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond) + _ = contextutil.ExecFunc(ctx, func() { + time.Sleep(40 * time.Millisecond) // outlives the context + }) + cancel() + } + assertNoGoroutineLeak(t, base) +} + +func TestExecFuncWithTwoReturns_NoGoroutineLeakOnCancel(t *testing.T) { + base := baseline() + for range leakCalls { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond) + _, _ = contextutil.ExecFuncWithTwoReturns(ctx, func() (int, error) { + time.Sleep(40 * time.Millisecond) // outlives the context + return 42, nil + }) + cancel() + } + assertNoGoroutineLeak(t, base) +} + +func TestExecFuncWithThreeReturns_NoGoroutineLeakOnCancel(t *testing.T) { + base := baseline() + for range leakCalls { + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond) + _, _, _ = contextutil.ExecFuncWithThreeReturns(ctx, func() (int, string, error) { + time.Sleep(40 * time.Millisecond) // outlives the context + return 42, "hello", nil + }) + cancel() + } + assertNoGoroutineLeak(t, base) +} diff --git a/reader/conn_read.go b/reader/conn_read.go index 95424bb4..af0570f0 100644 --- a/reader/conn_read.go +++ b/reader/conn_read.go @@ -26,6 +26,10 @@ var ( // for file/buffer based reading, ConnReadN should be preferred // instead of 'conn.Read() without loop' . It ignores EOF, UnexpectedEOF and timeout errors // Note: you are responsible for adding a timeout to context +// Note: if the reader supports SetReadDeadline (e.g. net.Conn), a cancelled +// context expires its read deadline to unblock a pending read, leaving the +// deadline in the past; a connection whose read was cancelled should be +// discarded rather than reused. func ConnReadN(ctx context.Context, reader io.Reader, N int64) ([]byte, error) { if N == -1 { N = MaxReadSize @@ -52,8 +56,21 @@ func ConnReadN(ctx context.Context, reader io.Reader, N int64) ([]byte, error) { fn := func() (int64, error) { return io.CopyN(pw, io.LimitReader(reader, N), N) } - // ExecFuncWithTwoReturns will execute the function but errors if context is done + // A context deadline can't interrupt a blocking Read; if the reader + // supports deadlines (net.Conn and friends) expire it when ctx is + // cancelled so the read returns instead of leaking the goroutine and the + // connection. The deferred stop cancels this on the normal path, leaving + // the deadline untouched when the read finishes in time. + if rd, ok := reader.(interface{ SetReadDeadline(time.Time) error }); ok { + defer context.AfterFunc(ctx, func() { _ = rd.SetReadDeadline(time.Now()) })() + } _, readErr = contextutil.ExecFuncWithTwoReturns(ctx, fn) + // On cancellation report the context error rather than the net timeout + // produced by expiring the deadline, so the timeout handling below is + // deterministic instead of depending on which goroutine wins. + if readErr != nil && ctx.Err() != nil { + readErr = ctx.Err() + } }() // read from pipe and return diff --git a/reader/conn_read_cancel_test.go b/reader/conn_read_cancel_test.go new file mode 100644 index 00000000..768417a6 --- /dev/null +++ b/reader/conn_read_cancel_test.go @@ -0,0 +1,115 @@ +package reader + +import ( + "context" + "io" + "net" + "runtime" + "testing" + "time" +) + +// TestConnReadN_NoGoroutineLeakOnCancel guards against ConnReadN leaking its +// read goroutine when the context is cancelled before the peer sends data. +// +// ConnReadN performs the read in a goroutine and returns the context error when +// ctx is done. A context deadline can't interrupt a blocking socket read, so +// unless ConnReadN actively unblocks the read on cancellation the goroutine +// stays parked in Read for the connection's lifetime, leaking the goroutine, +// its buffers, and the connection on every cancelled read. +func TestConnReadN_NoGoroutineLeakOnCancel(t *testing.T) { + addr := newSilentServer(t) + + runtime.GC() + time.Sleep(50 * time.Millisecond) + runtime.GC() + base := runtime.NumGoroutine() + + const calls = 50 + for range calls { + conn, err := net.Dial("tcp", addr) + if err != nil { + t.Fatal(err) + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond) + _, _ = ConnReadN(ctx, conn, 16) // peer sends nothing; ctx expires first + cancel() + t.Cleanup(func() { _ = conn.Close() }) + } + + deadline := time.Now().Add(5 * time.Second) + for { + runtime.GC() + leaked := runtime.NumGoroutine() - base + if leaked <= 2 { // tolerance for transient runtime goroutines + return + } + if time.Now().After(deadline) { + t.Fatalf("goroutine leak: ~%d goroutines still alive after %d context-cancelled ConnReadN calls (base=%d, now=%d)", + leaked, calls, base, runtime.NumGoroutine()) + } + time.Sleep(25 * time.Millisecond) + } +} + +// TestConnReadN_ReturnsPartialDataOnCancel verifies that data already received +// before the context is cancelled is returned rather than dropped. Expiring the +// read deadline to unblock the read produces a net timeout error; ConnReadN must +// report the cancellation as the context error so the partial data is returned. +func TestConnReadN_ReturnsPartialDataOnCancel(t *testing.T) { + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = ln.Close() }) + go func() { + c, err := ln.Accept() + if err != nil { + return + } + defer c.Close() + _, _ = c.Write([]byte("hi")) // send partial data, then stall + _, _ = io.Copy(io.Discard, c) // block until the client closes + }() + + client, err := net.Dial("tcp", ln.Addr().String()) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = client.Close() }) + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + + // asks for 16 bytes but only 2 arrive before the context expires + data, err := ConnReadN(ctx, client, 16) + if err != nil { + t.Fatalf("expected partial data with no error, got err: %v", err) + } + if string(data) != "hi" { + t.Fatalf("expected %q, got %q", "hi", string(data)) + } +} + +// newSilentServer returns the address of a TCP server that accepts connections +// and holds them open without ever writing, so reads against it block until the +// reader's deadline or close. +func newSilentServer(t *testing.T) string { + t.Helper() + ln, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatal(err) + } + held := make(chan net.Conn, 1024) + go func() { + for { + c, err := ln.Accept() + if err != nil { + return + } + held <- c // hold the server end open; never write + } + }() + t.Cleanup(func() { _ = ln.Close() }) + return ln.Addr().String() +}