fix: prevent goroutine leak on context-cancelled reads (contextutil + reader.ConnReadN)

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}

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

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

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