improve calculation of time to renewal

Author: ndyakov

entraid_test.go

@@ -82,6 +82,9 @@ type mockIdentityProvider struct {
 
 func (m *mockIdentityProvider) RequestToken() (IdentityProviderResponse, error) {
 	args := m.Called()
+	if args.Get(0) == nil {
+		return nil, args.Error(1)
+	}
 	return args.Get(0).(IdentityProviderResponse), args.Error(1)
 }
 

token.go

@@ -60,6 +60,9 @@ func NewToken(username, password, rawToken string, expiresOn, receivedAt time.Ti
 
 // copyToken creates a copy of the token.
 func copyToken(token *Token) *Token {
+	if token == nil {
+		return nil
+	}
 	return NewToken(token.username, token.password, token.rawToken, token.expiresOn, token.receivedAt, token.ttl)
 }
 

token_manager.go

@@ -23,7 +23,7 @@ type TokenManagerOptions struct {
 	// default: 0.7
 	ExpirationRefreshRatio float64
 	// LowerRefreshBoundMs is the lower bound for the refresh time in milliseconds.
-	// Represents the minimum time in milliseconds before token expiration to trigger a refresh, in milliseconds.
+	// Represents the minimum time in milliseconds before token expiration to trigger a refresh.
 	// This value sets a fixed lower bound for when a token refresh should occur, regardless
 	// of the token's total lifetime.
 	//
@@ -239,7 +239,8 @@ type entraidTokenManager struct {
 }
 
 func (e *entraidTokenManager) GetToken() (*Token, error) {
-	if e.token != nil && e.token.expiresOn.Before(time.Now().Add(e.lowerBoundDuration)) {
+	// check if the token is nil and if it is not expired
+	if e.token != nil && e.token.expiresOn.After(time.Now()) && e.durationToRenewal() > e.lowerBoundDuration {
 		// copy the token so the caller can't modify it
 		return copyToken(e.token), nil
 	}
@@ -256,6 +257,10 @@ func (e *entraidTokenManager) GetToken() (*Token, error) {
 
 	// copy the token so the caller can't modify it
 	e.token = copyToken(token)
+
+	if e.token == nil {
+		return nil, fmt.Errorf("failed to get token: token is nil")
+	}
 	return token, nil
 }
 
@@ -276,12 +281,25 @@ func (e *entraidTokenManager) durationToRenewal() time.Duration {
 	if e.token == nil {
 		return 0
 	}
+	timeTillExpiration := time.Until(e.token.expiresOn)
+
+	// if lower bound has passed, do it NOW
+	if timeTillExpiration <= e.lowerBoundDuration {
+		return 0
+	}
+
 	// Calculate the time to renew the token based on the expiration refresh ratio
-	duration := time.Duration(float64(time.Until(e.token.expiresOn)) * e.expirationRefreshRatio)
-	if duration < e.lowerBoundDuration {
-		return e.lowerBoundDuration
+	duration := time.Duration(float64(timeTillExpiration) * e.expirationRefreshRatio)
+	if duration <= 0 {
+		return 0
+	}
+
+	// if the duration will take us past the lower bound, return the duration to lower bound
+	if timeTillExpiration-e.lowerBoundDuration < duration {
+		return timeTillExpiration - e.lowerBoundDuration
 	}
 
+	// return the calculated duration
 	return duration
 }
 
@@ -309,13 +327,19 @@ func (e *entraidTokenManager) Start(listener TokenListener) (cancelFunc, error)
 	go func(listener TokenListener) {
 		// Simulate token refresh
 		for {
+			timeToRenewal := e.durationToRenewal()
 			select {
 			case <-e.closed:
 				// Token manager is closed, stop the loop
 				return
-			case <-time.After(e.durationToRenewal()):
+			case <-time.After(timeToRenewal):
 				// Token is about to expire, refresh it
 				delay := time.Duration(e.retryOptions.InitialDelayMs) * time.Millisecond
+				// Token asked to be refreshed asap, but let's make sure we wait a bit
+				if timeToRenewal == 0 {
+					time.Sleep(delay)
+				}
+
 				for i := 0; i < e.retryOptions.MaxAttempts; i++ {
 					select {
 					case <-e.closed:

token_manager_test.go

@@ -308,7 +308,7 @@ func TestTokenManager_Start(t *testing.T) {
 				go func() {
 					defer wg.Done()
 					time.Sleep(time.Duration(int64(rand.Intn(100)) * int64(time.Millisecond)))
-					_, err = tokenManager.Start(listener)
+					_, err := tokenManager.Start(listener)
 					if err == nil {
 						hasStarted += 1
 						return
@@ -554,6 +554,7 @@ func TestEntraidTokenManager_GetToken(t *testing.T) {
 		assert.NotNil(t, token)
 
 	})
+
 	t.Run("GetToken with parse error", func(t *testing.T) {
 		idp := &mockIdentityProvider{}
 		listener := &mockTokenListener{}
@@ -580,7 +581,93 @@ func TestEntraidTokenManager_GetToken(t *testing.T) {
 		assert.Error(t, err)
 		assert.Nil(t, cancel)
 		assert.NotNil(t, tm.listener)
+	})
+	t.Run("GetToken with expired token", func(t *testing.T) {
+		idp := &mockIdentityProvider{}
+		tokenManager, err := NewTokenManager(idp,
+			TokenManagerOptions{},
+		)
+		assert.NoError(t, err)
+
+		authResult := &public.AuthResult{
+			ExpiresOn: time.Now().Add(-time.Hour).UTC(),
+		}
+		idpResponse, err := NewIDPResponse(ResponseTypeAuthResult,
+			authResult)
+		assert.NoError(t, err)
+		assert.NotNil(t, tokenManager)
+		tm, ok := tokenManager.(*entraidTokenManager)
+		assert.True(t, ok)
+		assert.Nil(t, tm.listener)
+
+		idp.On("RequestToken").Return(idpResponse, nil)
+
+		token, err := tokenManager.GetToken()
+		assert.Error(t, err)
+		assert.Nil(t, token)
+	})
+
+	t.Run("GetToken with nil token", func(t *testing.T) {
+		idp := &mockIdentityProvider{}
+		tokenManager, err := NewTokenManager(idp,
+			TokenManagerOptions{},
+		)
+		assert.NoError(t, err)
+		assert.NotNil(t, tokenManager)
+		_, ok := tokenManager.(*entraidTokenManager)
+		assert.True(t, ok)
+
+		rawResponse, err := NewIDPResponse(ResponseTypeRawToken, "test")
+		assert.NoError(t, err)
+
+		idp.On("RequestToken").Return(rawResponse, nil)
+
+		token, err := tokenManager.GetToken()
+		assert.Error(t, err)
+		assert.Nil(t, token)
+	})
+
+	t.Run("GetToken with nil from parser", func(t *testing.T) {
+		idp := &mockIdentityProvider{}
+		mParser := &mockIdentityProviderResponseParser{}
+		tokenManager, err := NewTokenManager(idp,
+			TokenManagerOptions{
+				IdentityProviderResponseParser: mParser,
+			},
+		)
+		assert.NoError(t, err)
+		assert.NotNil(t, tokenManager)
+		_, ok := tokenManager.(*entraidTokenManager)
+		assert.True(t, ok)
+
+		idpResponse, err := NewIDPResponse(ResponseTypeRawToken, "test")
+		assert.NoError(t, err)
+		idp.On("RequestToken").Return(idpResponse, nil)
+		mParser.On("ParseResponse", idpResponse).Return(nil, nil)
+
+		token, err := tokenManager.GetToken()
+		assert.Error(t, err)
+		assert.Nil(t, token)
+	})
 
+	t.Run("GetToken with idp error", func(t *testing.T) {
+		idp := &mockIdentityProvider{}
+		mParser := &mockIdentityProviderResponseParser{}
+		tokenManager, err := NewTokenManager(idp,
+			TokenManagerOptions{
+				IdentityProviderResponseParser: mParser,
+			},
+		)
+		assert.NoError(t, err)
+		assert.NotNil(t, tokenManager)
+		_, ok := tokenManager.(*entraidTokenManager)
+		assert.True(t, ok)
+
+		idp.On("RequestToken").Return(nil, fmt.Errorf("idp error"))
+
+		token, err := tokenManager.GetToken()
+		assert.Error(t, err)
+		assert.Nil(t, token)
 	})
 }
 
@@ -591,7 +678,6 @@ func TestEntraidTokenManager_durationToRenewal(t *testing.T) {
 		idp := &mockIdentityProvider{}
 		tokenManager, err := NewTokenManager(idp, TokenManagerOptions{
 			LowerRefreshBoundMs: 1000 * 60 * 60, // 1 hour
-
 		})
 		assert.NoError(t, err)
 		assert.NotNil(t, tokenManager)
@@ -610,15 +696,89 @@ func TestEntraidTokenManager_durationToRenewal(t *testing.T) {
 			idpResponse, err := NewIDPResponse(ResponseTypeAuthResult,
 				expiresSoon)
 			assert.NoError(t, err)
-			idp.On("RequestToken").Return(idpResponse, nil)
+			idp.On("RequestToken").Return(idpResponse, nil).Once()
+			tm.token = nil
+			_, err = tm.GetToken()
+			assert.NoError(t, err)
+			assert.NotNil(t, tm.token)
+
+			// return zero, should happen now since it expires before the lower bound
+			result = tm.durationToRenewal()
+			assert.Equal(t, time.Duration(0), result)
+		})
 
+		// get token that expires after the lower bound and expirationRefreshRatio to 1
+		assert.NotPanics(t, func() {
+			tm.expirationRefreshRatio = 1
+			expiresAfterlb := &public.AuthResult{
+				ExpiresOn: time.Now().Add(time.Duration(tm.lowerBoundDuration) + time.Hour).UTC(),
+			}
+			idpResponse, err := NewIDPResponse(ResponseTypeAuthResult,
+				expiresAfterlb)
+			assert.NoError(t, err)
+			idp.On("RequestToken").Return(idpResponse, nil).Once()
+			tm.token = nil
 			_, err = tm.GetToken()
 			assert.NoError(t, err)
 			assert.NotNil(t, tm.token)
+
+			// return time to lower bound, if the returned time will be after the lower bound
+			result = tm.durationToRenewal()
+			assert.InDelta(t, time.Until(tm.token.expiresOn.Add(-1*tm.lowerBoundDuration)), result, float64(time.Second))
 		})
 
-		// return the lower bound
-		result = tm.durationToRenewal()
-		assert.Equal(t, tm.lowerBoundDuration, result)
+	})
+}
+
+func TestEntraidTokenManager_Streaming(t *testing.T) {
+	// write a test that will cover the goroutine in the Start
+	t.Parallel()
+	t.Run("Streaming", func(t *testing.T) {
+		idp := &mockIdentityProvider{}
+		listener := &mockTokenListener{}
+		mParser := &mockIdentityProviderResponseParser{}
+		tokenManager, err := NewTokenManager(idp,
+			TokenManagerOptions{
+				IdentityProviderResponseParser: mParser,
+			},
+		)
+		assert.NoError(t, err)
+		assert.NotNil(t, tokenManager)
+		tm, ok := tokenManager.(*entraidTokenManager)
+		assert.True(t, ok)
+		assert.Nil(t, tm.listener)
+
+		expiresIn := 10 * time.Millisecond
+		expiresOn := time.Now().Add(expiresIn).UTC()
+		authResult := &public.AuthResult{
+			ExpiresOn: expiresOn,
+		}
+		idpResponse, err := NewIDPResponse(ResponseTypeAuthResult,
+			authResult)
+		assert.NoError(t, err)
+
+		idp.On("RequestToken").Return(idpResponse, nil)
+		token := NewToken(
+			"test",
+			"test",
+			"test",
+			expiresOn,
+			time.Now(),
+			int64(time.Until(expiresOn)),
+		)
+
+		mParser.On("ParseResponse", idpResponse).Return(token, nil)
+		listener.On("OnTokenNext", token).Return()
+
+		cancel, err := tokenManager.Start(listener)
+		assert.NotNil(t, cancel)
+		assert.NoError(t, err)
+		assert.NotNil(t, tm.listener)
+
+		toRenewal := tm.durationToRenewal()
+		assert.NotEqual(t, 0, toRenewal)
+		assert.NotEqual(t, expiresIn, toRenewal)
+		assert.True(t, expiresIn > toRenewal)
+		// should fail on mocks
 	})
 }