fix(credentials): leak in wrappedOnClose (#3785)

* fix leak in wrappedOnClose

* address linter warning

auth/auth.go

@@ -9,12 +9,30 @@ type StreamingCredentialsProvider interface {
 	// Subscribe subscribes to the credentials provider for updates.
 	// It returns the current credentials, a cancel function to unsubscribe from the provider,
 	// and an error if any.
+	//
+	// Implementations MUST be idempotent with respect to listener identity:
+	// subscribing the same listener value more than once must not produce
+	// duplicate notifications and must not create multiple independent
+	// subscriptions that each need to be cancelled separately. Every
+	// UnsubscribeFunc returned for a given listener must cancel that
+	// listener's subscription; calling any one of them must be sufficient to
+	// stop updates to that listener, and calling subsequent ones must be a
+	// safe no-op. Callers (including go-redis internals) may retain only
+	// the most recently returned UnsubscribeFunc and rely on it to fully
+	// unsubscribe the listener.
+	//
 	// TODO(ndyakov): Should we add context to the Subscribe method?
 	Subscribe(listener CredentialsListener) (Credentials, UnsubscribeFunc, error)
 }
 
 // UnsubscribeFunc is a function that is used to cancel the subscription to the credentials provider.
 // It is used to unsubscribe from the provider when the credentials are no longer needed.
+//
+// Per the StreamingCredentialsProvider.Subscribe contract, if the same
+// listener is subscribed multiple times, every UnsubscribeFunc returned for
+// that listener must fully unsubscribe it on first invocation, and
+// subsequent invocations (from any of the equivalent UnsubscribeFuncs) must
+// be a safe no-op.
 type UnsubscribeFunc func() error
 
 // CredentialsListener is an interface that defines the methods for a credentials listener.

internal/pool/conn.go

@@ -575,6 +575,41 @@ func (cn *Conn) getEffectiveWriteTimeout(normalTimeout time.Duration) time.Durat
 	}
 }
 
+// SetOnClose installs fn as the callback invoked exactly once when this
+// connection is closed (via Conn.Close).
+//
+// IMPORTANT: SetOnClose OVERWRITES any previously installed callback — it
+// does not compose, chain, or deduplicate. A Conn has room for a single
+// onClose hook by design, because its lifecycle is bounded (a Conn is
+// created, optionally re-initialized on its own net.Conn, and then closed
+// once) and the pool's OnRemove hooks handle any registry-level cleanup
+// that must survive the net.Conn being swapped.
+//
+// This has a subtle implication for per-connection subscriptions such as
+// the unsubscribe function returned by StreamingCredentialsProvider
+// (e.g. EntraID token rotation): if SetOnClose is called twice on the
+// same Conn with DIFFERENT unsubscribe closures — for example because
+// initConn ran a second time and obtained a fresh Subscribe() —
+// the previous unsubscribe is dropped and will NEVER run, leaking a
+// subscription on the provider. Callers must therefore ensure either:
+//
+//   - the provider's Subscribe is idempotent for the same listener (the
+//     streaming credentials Manager deduplicates listeners by connection
+//     id, so re-Subscribe returns an equivalent unsubscribe), OR
+//   - the previous callback has already been invoked before SetOnClose is
+//     called again.
+//
+// Design note: unlike the client-level onCloseHooks registry (see
+// redis.baseClient), there is intentionally NO named-hook dedup or
+// multi-callback support on Conn. This is a deliberate trade-off to keep
+// the Conn object slim — a pool can hold thousands of Conn values and
+// each one is a hot allocation, so paying for a sync.Mutex plus a
+// map[string]func() error per connection to support a feature that would
+// only be used by at most one subsystem today (streaming credentials) is
+// not worth the per-connection memory and allocation cost. For a single
+// Conn there is at most one meaningful close callback at any point in
+// time, and a richer registry here would not even solve the "stale
+// closure" hazard described above.
 func (cn *Conn) SetOnClose(fn func() error) {
 	cn.onClose = fn
 }

internal_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/redis/go-redis/v9/auth"
 	"github.com/redis/go-redis/v9/internal/pool"
 	"github.com/redis/go-redis/v9/internal/proto"
 
@@ -460,6 +461,10 @@ func (ct *testCounter) expect(values map[string]int) {
 	}
 }
 
+// testOnCloseHookID is the id used by the ring-shard cleanup tests when
+// registering a close hook against the internal onCloseHooks registry.
+const testOnCloseHookID = "test-close-counter"
+
 func TestRingShardsCleanup(t *testing.T) {
 	const (
 		ringShard1Name = "ringShardOne"
@@ -479,7 +484,7 @@ func TestRingShardsCleanup(t *testing.T) {
 			},
 			NewClient: func(opt *Options) *Client {
 				c := NewClient(opt)
-				c.baseClient.onClose = c.baseClient.wrappedOnClose(func() error {
+				c.baseClient.onClose.register(testOnCloseHookID, func() error {
 					closeCounter.increment(opt.Addr)
 					return nil
 				})
@@ -528,7 +533,7 @@ func TestRingShardsCleanup(t *testing.T) {
 				}
 				createCounter.increment(opt.Addr)
 				c := NewClient(opt)
-				c.baseClient.onClose = c.baseClient.wrappedOnClose(func() error {
+				c.baseClient.onClose.register(testOnCloseHookID, func() error {
 					closeCounter.increment(opt.Addr)
 					return nil
 				})
@@ -685,3 +690,280 @@ var _ = Describe("isLoopback", func() {
 		Entry("partial docker internal", "docker.internal", false),
 	)
 })
+
+
+// TestOnCloseHooks_RunInRegistrationOrder verifies that hooks registered under
+// distinct ids are all invoked on run() in the order they were registered.
+func TestOnCloseHooks_RunInRegistrationOrder(t *testing.T) {
+	h := &onCloseHooks{}
+	var calls []string
+
+	h.register("a", func() error { calls = append(calls, "a"); return nil })
+	h.register("b", func() error { calls = append(calls, "b"); return nil })
+	h.register("c", func() error { calls = append(calls, "c"); return nil })
+
+	if err := h.run(); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	want := []string{"a", "b", "c"}
+	if !reflect.DeepEqual(calls, want) {
+		t.Fatalf("run order = %v, want %v", calls, want)
+	}
+}
+
+// TestOnCloseHooks_RegisterSameIDReplaces is the regression test for issue
+// #3772. Registering the same id repeatedly must replace the existing
+// callback rather than chain onto it, so the registry stays bounded even
+// under storm-like re-registration (the exact scenario that previously leaked
+// when initConn re-wrapped c.onClose on every connection init).
+func TestOnCloseHooks_RegisterSameIDReplaces(t *testing.T) {
+	h := &onCloseHooks{}
+	const id = "same-id"
+	const iterations = 10_000
+
+	var lastSeen int32
+	for i := 0; i < iterations; i++ {
+		i := int32(i)
+		h.register(id, func() error { atomic.StoreInt32(&lastSeen, i); return nil })
+	}
+
+	if got := len(h.order); got != 1 {
+		t.Fatalf("order length after %d re-registrations = %d, want 1", iterations, got)
+	}
+	if got := len(h.hooks); got != 1 {
+		t.Fatalf("hooks map size after %d re-registrations = %d, want 1", iterations, got)
+	}
+
+	if err := h.run(); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if got := atomic.LoadInt32(&lastSeen); got != iterations-1 {
+		t.Fatalf("last-registered callback not invoked: lastSeen = %d, want %d", got, iterations-1)
+	}
+}
+
+// TestOnCloseHooks_DistinctIDsCoexist guarantees the dedup behavior does not
+// discard hooks from other callers: registering new ids must never drop
+// previously registered ids.
+func TestOnCloseHooks_DistinctIDsCoexist(t *testing.T) {
+	h := &onCloseHooks{}
+	var aCount, bCount int32
+
+	h.register("a", func() error { atomic.AddInt32(&aCount, 1); return nil })
+	h.register("b", func() error { atomic.AddInt32(&bCount, 1); return nil })
+	// Re-registering "a" must not drop "b".
+	h.register("a", func() error { atomic.AddInt32(&aCount, 1); return nil })
+
+	if err := h.run(); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if a, b := atomic.LoadInt32(&aCount), atomic.LoadInt32(&bCount); a != 1 || b != 1 {
+		t.Fatalf("call counts a=%d b=%d, want a=1 b=1", a, b)
+	}
+}
+
+// TestOnCloseHooks_Unregister verifies that unregister removes a hook and
+// that running after unregister does not invoke it.
+func TestOnCloseHooks_Unregister(t *testing.T) {
+	h := &onCloseHooks{}
+	var aCalled, bCalled bool
+
+	h.register("a", func() error { aCalled = true; return nil })
+	h.register("b", func() error { bCalled = true; return nil })
+	h.unregister("a")
+	h.unregister("missing") // no-op must not panic
+
+	if err := h.run(); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	if aCalled {
+		t.Fatal("unregistered hook was invoked")
+	}
+	if !bCalled {
+		t.Fatal("remaining hook was not invoked")
+	}
+	if got := len(h.order); got != 1 {
+		t.Fatalf("order length = %d, want 1", got)
+	}
+}
+
+
+// TestOnCloseHooks_AllRunOnError confirms every hook is invoked even if an
+// earlier one returns an error, and that the first error is returned.
+func TestOnCloseHooks_AllRunOnError(t *testing.T) {
+	h := &onCloseHooks{}
+	var called [3]bool
+	errFirst := fmt.Errorf("first")
+	errSecond := fmt.Errorf("second")
+
+	h.register("a", func() error { called[0] = true; return errFirst })
+	h.register("b", func() error { called[1] = true; return errSecond })
+	h.register("c", func() error { called[2] = true; return nil })
+
+	err := h.run()
+	if err != errFirst {
+		t.Fatalf("run() err = %v, want %v", err, errFirst)
+	}
+	for i, c := range called {
+		if !c {
+			t.Fatalf("hook %d was not invoked", i)
+		}
+	}
+}
+
+// TestOnCloseHooks_NilReceiver ensures run() on a nil registry is a safe
+// no-op. baseClient embedded in Conn/Tx does initialize the registry, but
+// defensive nil-safety lets future constructors add the field without
+// breaking Close().
+func TestOnCloseHooks_NilReceiver(t *testing.T) {
+	var h *onCloseHooks
+	if err := h.run(); err != nil {
+		t.Fatalf("run() on nil = %v, want nil", err)
+	}
+}
+
+// TestOnCloseHooks_ConcurrentRegisterSameID hammers the registry with many
+// goroutines re-registering under the same id. The registry must remain
+// bounded and the surviving callback must still be invoked exactly once.
+func TestOnCloseHooks_ConcurrentRegisterSameID(t *testing.T) {
+	h := &onCloseHooks{}
+	const id = "hot"
+	const goroutines = 64
+	const perG = 1_000
+
+	var wg sync.WaitGroup
+	wg.Add(goroutines)
+	for g := 0; g < goroutines; g++ {
+		go func() {
+			defer wg.Done()
+			for i := 0; i < perG; i++ {
+				h.register(id, func() error { return nil })
+			}
+		}()
+	}
+	wg.Wait()
+
+	if got := len(h.order); got != 1 {
+		t.Fatalf("order length after concurrent storm = %d, want 1", got)
+	}
+}
+
+
+// entraidLikeProvider mimics the exact semantics of
+// github.com/redis/go-redis-entraid's StreamingCredentialsProvider relevant
+// to issue #3772: it deduplicates subscriptions by listener pointer
+// identity, and every call to Subscribe returns a FRESH UnsubscribeFunc
+// closure that removes the listener by pointer match from the shared
+// listeners slice. Two unsubs obtained for the same listener are therefore
+// equivalent: the first one called removes the entry, any subsequent call
+// is a safe no-op.
+type entraidLikeProvider struct {
+	mu          sync.Mutex
+	listeners   []auth.CredentialsListener
+	subscribeN  int32
+	unsubCalls  int32
+}
+
+func (p *entraidLikeProvider) Subscribe(listener auth.CredentialsListener) (auth.Credentials, auth.UnsubscribeFunc, error) {
+	atomic.AddInt32(&p.subscribeN, 1)
+
+	p.mu.Lock()
+	already := false
+	for _, l := range p.listeners {
+		if l == listener {
+			already = true
+			break
+		}
+	}
+	if !already {
+		p.listeners = append(p.listeners, listener)
+	}
+	p.mu.Unlock()
+
+	unsub := func() error {
+		atomic.AddInt32(&p.unsubCalls, 1)
+		p.mu.Lock()
+		defer p.mu.Unlock()
+		for i, l := range p.listeners {
+			if l == listener {
+				p.listeners = append(p.listeners[:i], p.listeners[i+1:]...)
+				return nil
+			}
+		}
+		return nil
+	}
+	return auth.NewBasicCredentials("u", "p"), unsub, nil
+}
+
+func (p *entraidLikeProvider) listenerCount() int {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return len(p.listeners)
+}
+
+// stubCredentialsListener is a minimal auth.CredentialsListener used only
+// for pointer-identity in the entraid-mimicking test.
+type stubCredentialsListener struct{}
+
+func (*stubCredentialsListener) OnNext(auth.Credentials) {}
+func (*stubCredentialsListener) OnError(error)           {}
+
+// TestInitConn_EntraidLike_NoLeakAcrossReinits is the targeted regression
+// test for issue #3772 against the real-world StreamingCredentialsProvider
+// behavior implemented by go-redis-entraid. It simulates N re-initializations
+// on the same logical connection (i.e. the same CredentialsListener pointer,
+// which is what streaming.Manager.Listener returns from its per-connId cache),
+// replacing cn.SetOnClose with each new unsubscribe closure.
+//
+// Invariants the fix must uphold:
+//  1. Subscribe dedups: the provider's listener list stays at size 1.
+//  2. Calling only the MOST RECENT unsub fully removes the listener.
+//  3. All prior (orphaned) unsubs are safe no-ops after that.
+//  4. No registration remains on the provider after close.
+func TestInitConn_EntraidLike_NoLeakAcrossReinits(t *testing.T) {
+	const reinits = 1000
+
+	provider := &entraidLikeProvider{}
+	listener := &stubCredentialsListener{}
+
+	var latestUnsub auth.UnsubscribeFunc
+	orphaned := make([]auth.UnsubscribeFunc, 0, reinits-1)
+
+	for i := 0; i < reinits; i++ {
+		_, unsub, err := provider.Subscribe(listener)
+		if err != nil {
+			t.Fatalf("Subscribe #%d: %v", i, err)
+		}
+		if latestUnsub != nil {
+			// Mirror the pool.Conn.SetOnClose behavior: the previous unsub
+			// is dropped on the floor, only the latest one is retained.
+			orphaned = append(orphaned, latestUnsub)
+		}
+		latestUnsub = unsub
+	}
+
+	if got := provider.listenerCount(); got != 1 {
+		t.Fatalf("after %d Subscribes with same listener, listener count = %d, want 1", reinits, got)
+	}
+	if got := atomic.LoadInt32(&provider.subscribeN); got != int32(reinits) {
+		t.Fatalf("Subscribe call count = %d, want %d", got, reinits)
+	}
+
+	// Only the latest unsub is invoked, matching the post-fix cn.onClose.
+	if err := latestUnsub(); err != nil {
+		t.Fatalf("latest unsub returned error: %v", err)
+	}
+	if got := provider.listenerCount(); got != 0 {
+		t.Fatalf("listener count after latest unsub = %d, want 0", got)
+	}
+
+	// Every orphaned unsub must be a safe no-op (contract in auth.UnsubscribeFunc).
+	for i, u := range orphaned {
+		if err := u(); err != nil {
+			t.Fatalf("orphaned unsub #%d returned error: %v", i, err)
+		}
+	}
+	if got := provider.listenerCount(); got != 0 {
+		t.Fatalf("listener count after orphaned unsubs = %d, want 0", got)
+	}
+}

redis.go

@@ -215,6 +215,96 @@ func (hs *hooksMixin) processTxPipelineHook(ctx context.Context, cmds []Cmder) e
 
 //------------------------------------------------------------------------------
 
+// Stable identifiers for baseClient.onClose hooks. Each component that
+// registers a close callback owns a dedicated id here so the set of known
+// hooks is discoverable in one place and id collisions are caught at
+// compile time. New ids should be added as additional constants.
+const (
+	// onCloseHookIDSentinelFailover identifies the close callback installed
+	// by NewFailoverClient to tear down sentinel failover background work.
+	onCloseHookIDSentinelFailover = "sentinel-failover"
+)
+
+// onCloseHooks is a small registry of named close callbacks attached to a
+// baseClient. Each callback is identified by a stable string id; registering
+// the same id twice replaces the previous callback rather than chaining onto
+// it. This guarantees the registry stays bounded regardless of how often a
+// hook is (re)registered and avoids the unbounded closure chain that
+// motivated issue #3772.
+//
+// Hooks are invoked in registration order. All hooks run regardless of
+// individual errors; the first non-nil error is returned.
+//
+// A zero-value onCloseHooks is ready to use. It is safe for concurrent use.
+// Clones of a baseClient share the same *onCloseHooks so registrations and
+// close semantics are preserved across WithTimeout / WithContext / etc.
+type onCloseHooks struct {
+	mu    sync.Mutex
+	order []string
+	hooks map[string]func() error
+}
+
+// register adds or replaces the callback associated with id. Re-registering
+// an existing id overwrites the previous callback in place; new ids are
+// appended to the invocation order.
+func (h *onCloseHooks) register(id string, fn func() error) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	if h.hooks == nil {
+		h.hooks = make(map[string]func() error)
+	}
+	if _, exists := h.hooks[id]; !exists {
+		h.order = append(h.order, id)
+	}
+	h.hooks[id] = fn
+}
+
+// unregister removes the callback associated with id, if any. It is kept
+// for API symmetry with register so future callers (e.g. dynamic hook
+// owners that need to detach before client Close) do not have to
+// reinvent it.
+//
+//nolint:unused // kept for API symmetry with register; see comment above.
+func (h *onCloseHooks) unregister(id string) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	if _, exists := h.hooks[id]; !exists {
+		return
+	}
+	delete(h.hooks, id)
+	for i, x := range h.order {
+		if x == id {
+			h.order = append(h.order[:i], h.order[i+1:]...)
+			break
+		}
+	}
+}
+
+// run invokes all registered callbacks in registration order and returns
+// the first non-nil error encountered. All callbacks are executed even if
+// an earlier one returns an error.
+func (h *onCloseHooks) run() error {
+	if h == nil {
+		return nil
+	}
+	h.mu.Lock()
+	fns := make([]func() error, 0, len(h.order))
+	for _, id := range h.order {
+		if fn := h.hooks[id]; fn != nil {
+			fns = append(fns, fn)
+		}
+	}
+	h.mu.Unlock()
+
+	var firstErr error
+	for _, fn := range fns {
+		if err := fn(); err != nil && firstErr == nil {
+			firstErr = err
+		}
+	}
+	return firstErr
+}
+
 type baseClient struct {
 	opt        *Options
 	optLock    sync.RWMutex
@@ -222,7 +312,13 @@ type baseClient struct {
 	pubSubPool *pool.PubSubPool
 	hooksMixin
 
-	onClose func() error // hook called when client is closed
+	// onClose holds named callbacks invoked when the client is closed.
+	// Registering a new callback never removes previously registered ones;
+	// only re-registering the same id replaces the existing callback. This
+	// lets composing components (e.g. sentinel failover) add close logic
+	// safely without fear of overwriting each other and without building
+	// unbounded closure chains on repeated registration.
+	onClose *onCloseHooks
 
 	// Push notification processing
 	pushProcessor push.NotificationProcessor
@@ -376,27 +472,6 @@ func (c *baseClient) onAuthenticationErr() func(poolCn *pool.Conn, err error) {
 	}
 }
 
-func (c *baseClient) wrappedOnClose(newOnClose func() error) func() error {
-	onClose := c.onClose
-	return func() error {
-		var firstErr error
-		err := newOnClose()
-		// Even if we have an error we would like to execute the onClose hook
-		// if it exists. We will return the first error that occurred.
-		// This is to keep error handling consistent with the rest of the code.
-		if err != nil {
-			firstErr = err
-		}
-		if onClose != nil {
-			err = onClose()
-			if err != nil && firstErr == nil {
-				firstErr = err
-			}
-		}
-		return firstErr
-	}
-}
-
 func (c *baseClient) initConn(ctx context.Context, cn *pool.Conn) error {
 	// This function is called in two scenarios:
 	// 1. First-time init: Connection is in CREATED state (from pool.Get())
@@ -496,7 +571,22 @@ func (c *baseClient) initConn(ctx context.Context, cn *pool.Conn) error {
 			return fmt.Errorf("failed to subscribe to streaming credentials: %w", initErr)
 		}
 
-		c.onClose = c.wrappedOnClose(unsubscribeFromCredentialsProvider)
+		// Per-connection unsubscribe is attached to the connection itself so it
+		// runs when this specific connection is closed. Do not register it on
+		// c.onClose: initConn runs for every (re)initialized connection, and
+		// attaching per-connection state to the shared baseClient registry would
+		// either leak entries (one per connection id, never trimmed) or — with
+		// the pre-fix wrappedOnClose approach — build an unbounded closure chain
+		// retaining every prior connection's unsubscribe (see issue #3772).
+		//
+		// Note: pool.Conn.SetOnClose OVERWRITES any prior callback (see the
+		// doc on that method). That is safe here because the streaming
+		// credentials Manager deduplicates listeners by connection id, so a
+		// second initConn on the same cn re-Subscribes the SAME listener and
+		// the returned unsubscribe is equivalent to the one already installed.
+		// Any future code path that could hand out a distinct unsubscribe on
+		// re-initialization must first invoke the existing one to avoid
+		// orphaning the old subscription on the credentials provider.
 		cn.SetOnClose(unsubscribeFromCredentialsProvider)
 
 		username, password = credentials.BasicAuth()
@@ -963,10 +1053,8 @@ func (c *baseClient) Close() error {
 		firstErr = err
 	}
 
-	if c.onClose != nil {
-		if err := c.onClose(); err != nil && firstErr == nil {
-			firstErr = err
-		}
+	if err := c.onClose.run(); err != nil && firstErr == nil {
+		firstErr = err
 	}
 
 	// Unregister pools from OTel before closing them
@@ -1227,7 +1315,8 @@ func NewClient(opt *Options) *Client {
 
 	c := Client{
 		baseClient: &baseClient{
-			opt: opt,
+			opt:     opt,
+			onClose: &onCloseHooks{},
 		},
 	}
 	c.init()
@@ -1510,6 +1599,7 @@ func newConn(opt *Options, connPool pool.Pooler, parentHooks *hooksMixin) *Conn
 		baseClient: baseClient{
 			opt:      opt,
 			connPool: connPool,
+			onClose:  &onCloseHooks{},
 		},
 	}
 

sentinel.go

@@ -533,7 +533,8 @@ func NewFailoverClient(failoverOpt *FailoverOptions) *Client {
 
 	rdb := &Client{
 		baseClient: &baseClient{
-			opt: opt,
+			opt:     opt,
+			onClose: &onCloseHooks{},
 		},
 	}
 	rdb.init()
@@ -557,7 +558,7 @@ func NewFailoverClient(failoverOpt *FailoverOptions) *Client {
 		panic(fmt.Errorf("redis: failed to create pubsub pool: %w", err))
 	}
 
-	rdb.onClose = rdb.wrappedOnClose(failover.Close)
+	rdb.onClose.register(onCloseHookIDSentinelFailover, failover.Close)
 
 	failover.mu.Lock()
 	failover.onFailover = func(ctx context.Context, addr string) {
@@ -621,7 +622,8 @@ func NewSentinelClient(opt *Options) *SentinelClient {
 	opt.init()
 	c := &SentinelClient{
 		baseClient: &baseClient{
-			opt: opt,
+			opt:     opt,
+			onClose: &onCloseHooks{},
 		},
 	}
 

tx.go

@@ -28,6 +28,7 @@ func (c *Client) newTx() *Tx {
 			connPool:      pool.NewStickyConnPool(c.connPool),
 			hooksMixin:    c.hooksMixin.clone(),
 			pushProcessor: c.pushProcessor, // Copy push processor from parent client
+			onClose:       &onCloseHooks{},
 		},
 	}
 	tx.init()