refactor(manager): small refactors around the manager and token logic

internal/utils.go

@@ -1,16 +1,18 @@
 package internal
 
 // IsClosed checks if a channel is closed.
+// Returns true only if the channel is actually closed, not just if it has data available.
 //
-// NOTE: It returns true if the channel is closed as well
-// as if the channel is not empty. Used internally
-// to check if the channel is closed.
+// WARNING: This function will consume one value from the channel if it has pending data.
+// Use with caution on channels where consuming data might cause issues.
 func IsClosed(ch <-chan struct{}) bool {
 	select {
-	case <-ch:
-		return true
+	case _, ok := <-ch:
+		// If ok is false, the channel is closed
+		// If ok is true, the channel had data (which we just consumed)
+		return !ok
 	default:
+		// Channel is open but has no data available
+		return false
 	}
-
-	return false
 }

manager/entraid_manager.go

@@ -148,6 +148,16 @@ func (e *entraidTokenManager) Start(listener TokenListener) (StopFunc, error) {
 	e.listener = listener
 
 	go func(listener TokenListener, closed <-chan struct{}) {
+		// Add panic recovery to prevent crashes
+		defer func() {
+			if r := recover(); r != nil {
+				// Attempt to notify listener of panic, but don't panic again if that fails
+				func() {
+					defer func() { _ = recover() }()
+					listener.OnError(fmt.Errorf("token manager goroutine panic: %v", r))
+				}()
+			}
+		}()
 		maxDelay := e.retryOptions.MaxDelay
 		initialDelay := e.retryOptions.InitialDelay
 
@@ -223,6 +233,7 @@ func (e *entraidTokenManager) stop() (err error) {
 			err = fmt.Errorf("failed to stop token manager: %s", r)
 		}
 	}()
+
 	if e.ctxCancel != nil {
 		e.ctxCancel()
 	}
@@ -232,7 +243,11 @@ func (e *entraidTokenManager) stop() (err error) {
 	}
 
 	e.listener = nil
-	close(e.closedChan)
+
+	// Safely close the channel - only close if not already closed
+	if !internal.IsClosed(e.closedChan) {
+		close(e.closedChan)
+	}
 
 	return nil
 }

token/token.go

@@ -10,12 +10,25 @@ import (
 var _ auth.Credentials = (*Token)(nil)
 
 // New creates a new token with the specified username, password, raw token, expiration time, received at time, and time to live.
-// NOTE: This won't do any validation on the token, expiresOn, receivedAt, or ttl. It will simply create a new token instance.
-// The caller is responsible for ensuring the token is valid.
+// NOTE: The caller is responsible for ensuring the token is valid.
+// If the token is invalid, the behavior is undefined.
+// - if expiresOn is zero, New returns nil
+// - if receivedAt is zero, it will be set to the current time and TTL will be recalculated
 // Expiration time and TTL are used to determine when the token should be refreshed.
 // TTL is in milliseconds.
 // receivedAt + ttl should be within a millisecond of expiresOn
+// If receivedAt is zero, it will be set to the current time and TTL will be recalculated.
 func New(username, password, rawToken string, expiresOn, receivedAt time.Time, ttl int64) *Token {
+	// If expiresOn is zero, return nil to avoid panic in ReceivedAt()
+	if expiresOn.IsZero() {
+		return nil
+	}
+	// If receivedAt is zero, set it to now and recalculate TTL to avoid race conditions in ReceivedAt()
+	if receivedAt.IsZero() {
+		receivedAt = time.Now()
+		ttl = expiresOn.Sub(receivedAt).Milliseconds()
+	}
+
 	return &Token{
 		username:   username,
 		password:   password,
@@ -28,6 +41,10 @@ func New(username, password, rawToken string, expiresOn, receivedAt time.Time, t
 
 // Token represents parsed authentication token used to access the Redis server.
 // It implements the auth.Credentials interface.
+//
+// WARNING: Use New() to create a new token.
+// Creating a token with Token{} is invalid and will undefined behavior in the TokenManager.
+// The zero value of Token is not valid.
 type Token struct {
 	// username is the username of the user.
 	username string
@@ -60,11 +77,6 @@ func (t *Token) RawToken() string {
 
 // ReceivedAt returns the time when the token was received.
 func (t *Token) ReceivedAt() time.Time {
-	if t.receivedAt.IsZero() {
-		// set it to now, recalculate ttl
-		t.receivedAt = time.Now()
-		t.ttl = t.expiresOn.Sub(t.receivedAt).Milliseconds()
-	}
 	return t.receivedAt
 }
 

token/token_test.go

@@ -94,14 +94,15 @@ func TestCopyToken(t *testing.T) {
 	assert.NotEqual(t, token.expiresOn, copiedToken.expiresOn)
 
 	// copy nil
-	copiedToken = copyToken(nil)
-	assert.Nil(t, copiedToken)
+	nilToken := copyToken(nil)
+	assert.Nil(t, nilToken)
 	// copy empty token
-	copiedToken = copyToken(&Token{})
-	assert.NotNil(t, copiedToken)
+	emptyToken := copyToken(&Token{})
+	assert.Nil(t, emptyToken)
 	anotherCopy := copiedToken.Copy()
 	anotherCopy.rawToken = "changed"
 	assert.NotEqual(t, copiedToken, anotherCopy)
+	assert.NotEqual(t, copiedToken.rawToken, anotherCopy.rawToken)
 }
 
 func TestTokenReceivedAt(t *testing.T) {
@@ -124,7 +125,7 @@ func TestTokenReceivedAt(t *testing.T) {
 	// Check if the copied token is a new instance
 	assert.NotNil(t, tcopiedToken)
 
-	emptyRecievedAt := &Token{}
+	emptyRecievedAt := New("username", "password", "rawToken", time.Now(), time.Time{}, time.Hour.Milliseconds())
 	assert.True(t, emptyRecievedAt.ReceivedAt().After(time.Now().Add(-1*time.Hour)))
 	assert.True(t, emptyRecievedAt.ReceivedAt().Before(time.Now().Add(1*time.Hour)))
 }