Repository: envoyproxy/ai-gateway
Pull request #455 (2025-03-04 03:06:29)

controller: update rotate function to return expiration time
**Commit Message**

Update `Rotate` to return the expiration time in addition to the error. This will allow us to get the expiration time without having to make a `GET` API request to Kubernetes.



**Related Issues/PRs (if applicable)**

https://github.com/envoyproxy/ai-gateway/blob/8f38c672ee11e6e724fff7b1e430979417fd3ce6/internal/controller/backend_security_policy.go#L106-L108

```
2025-03-02T20:12:23Z    INFO    controller.backend-security-policy      successfully rotated credentials for credentials in namespace default of auth type AWSCredentials, renewing in -0.009733 minutes
```

The requeue time after rotating a set of credentials (AWS in this case) returns a negative number of minutes occasionally.  It requires a manual action/global sync to restart the rotation logic.

Upon further investigation, the issue found lies with `rotator.GetPreRotationTime` returning the incorrect time. GetPreRotation gets the expiration time from the secret's annotation, and the assumption is it's getting a cached previous secret's annotation instead of the new expiration time.

https://github.com/envoyproxy/ai-gateway/blob/8f38c672ee11e6e724fff7b1e430979417fd3ce6/internal/controller/backend_security_policy.go#L143-L147




wengyao04: @mathetake , in our DEV setup we find the the negative minutes used in the requeueAfter filed. In the controller-runtime [implementation](https://github.com/kubernetes-sigs/controller-runtime/blob/main/pkg/internal/controller/controller.go#L311-L360), if the requeqAfter is false, the event will be dropped to the reconciliation loop. We have to wait till the next global reconcile period (default is 10 hours)

We use mgr.Client() to get the secret, where the default read client read the secret from the informer cache and it is common that the informer cache is async with the k8s etcd. 

Beside this change, I think maybe we can use controller generic events with a time ticker to reconcile the backendsecuritypolicy every 10 minutes in case we miss any requeue event.


Patch:
diff --git a/internal/controller/backend_security_policy.go b/internal/controller/backend_security_policy.go
index 8f3e1d0..d7d4790 100644
--- a/internal/controller/backend_security_policy.go
+++ b/internal/controller/backend_security_policy.go
@@ -119,7 +119,6 @@ func (c *BackendSecurityPolicyController) reconcile(ctx context.Context, backend
 // rotateCredential rotates the credentials using the access token from OIDC provider and return the requeue time for next rotation.
 func (c *BackendSecurityPolicyController) rotateCredential(ctx context.Context, policy *aigv1a1.BackendSecurityPolicy, oidcCreds egv1a1.OIDC, rotator rotators.Rotator) (time.Duration, error) {
 	bspKey := backendSecurityPolicyKey(policy.Namespace, policy.Name)
-
 	var err error
 	c.oidcTokenCacheMutex.RLock()
 	validToken, ok := c.oidcTokenCache[bspKey]
@@ -136,15 +135,15 @@ func (c *BackendSecurityPolicyController) rotateCredential(ctx context.Context,
 	}
 
 	token := validToken.AccessToken
-	err = rotator.Rotate(ctx, token)
+	expiration, err := rotator.Rotate(ctx, token)
 	if err != nil {
 		return time.Minute, err
 	}
-	rotationTime, err := rotator.GetPreRotationTime(ctx)
-	if err != nil {
-		return time.Minute, err
+	rotationTime := expiration.Add(-preRotationWindow)
+	if requeue := time.Until(rotationTime); requeue > 0 {
+		return requeue, nil
 	}
-	return time.Until(rotationTime), nil
+	return time.Minute, fmt.Errorf("newly rotate credentials is already expired (%v) for policy %s in %s", rotationTime, policy.Name, policy.Namespace)
 }
 
 // getBackendSecurityPolicyAuthOIDC returns the backendSecurityPolicy's OIDC pointer or nil.
diff --git a/internal/controller/rotators/aws_oidc_rotator.go b/internal/controller/rotators/aws_oidc_rotator.go
index f8347dd..5198127 100644
--- a/internal/controller/rotators/aws_oidc_rotator.go
+++ b/internal/controller/rotators/aws_oidc_rotator.go
@@ -134,7 +134,7 @@ func populateSecretWithAwsIdentity(secret *corev1.Secret, awsIdentity *sts.Assum
 // Rotate implements aws credential secret upsert operation to k8s secret store.
 //
 // This implements [Rotator.Rotate].
-func (r *AWSOIDCRotator) Rotate(ctx context.Context, token string) error {
+func (r *AWSOIDCRotator) Rotate(ctx context.Context, token string) (time.Time, error) {
 	bspNamespace := r.backendSecurityPolicyNamespace
 	bspName := r.backendSecurityPolicyName
 	secretName := GetBSPSecretName(bspName)
@@ -143,8 +143,11 @@ func (r *AWSOIDCRotator) Rotate(ctx context.Context, token string) error {
 	awsIdentity, err := r.assumeRoleWithToken(ctx, token)
 	if err != nil {
 		r.logger.Error(err, "failed to assume role", "role", r.roleArn, "access token", token)
-		return err
+		return time.Time{}, err
+	} else if awsIdentity.Credentials == nil {
+		return time.Time{}, fmt.Errorf("unexpected nil awsIdentity credentials for %s in %s", r.backendSecurityPolicyName, r.backendSecurityPolicyNamespace)
 	}
+	r.logger.Info(fmt.Sprintf("awsIdentity Credentials will expire in '%s'", awsIdentity.Credentials.Expiration.String()), "namespace", bspNamespace, "name", bspName)
 
 	secret, err := LookupSecret(ctx, r.client, bspNamespace, secretName)
 	if err != nil {
@@ -159,14 +162,14 @@ func (r *AWSOIDCRotator) Rotate(ctx context.Context, token string) error {
 				Data: make(map[string][]byte),
 			}
 			populateSecretWithAwsIdentity(secret, awsIdentity, r.region)
-			return r.client.Create(ctx, secret)
+			return *awsIdentity.Credentials.Expiration, r.client.Create(ctx, secret)
 		}
 		r.logger.Error(err, "failed to lookup aws credentials secret", "namespace", bspNamespace, "name", bspName)
-		return err
+		return time.Time{}, err
 	}
 	r.logger.Info("updating existing aws credential secret", "namespace", bspNamespace, "name", bspName)
 	populateSecretWithAwsIdentity(secret, awsIdentity, r.region)
-	return r.client.Update(ctx, secret)
+	return *awsIdentity.Credentials.Expiration, r.client.Update(ctx, secret)
 }
 
 // assumeRoleWithToken exchanges an OIDC token for AWS credentials.
diff --git a/internal/controller/rotators/common.go b/internal/controller/rotators/common.go
index c43049a..9f37c26 100644
--- a/internal/controller/rotators/common.go
+++ b/internal/controller/rotators/common.go
@@ -26,8 +26,8 @@ type Rotator interface {
 	IsExpired(preRotationExpirationTime time.Time) bool
 	// GetPreRotationTime gets the time when the credentials need to be renewed.
 	GetPreRotationTime(ctx context.Context) (time.Time, error)
-	// Rotate will update the credential secret file with new credentials.
-	Rotate(ctx context.Context, token string) error
+	// Rotate will update the credential secret file with new credentials and return expiration time.
+	Rotate(ctx context.Context, token string) (time.Time, error)
 }
 
 // LookupSecret retrieves an existing secret.


Test patch:
diff --git a/internal/controller/backend_security_policy_test.go b/internal/controller/backend_security_policy_test.go
index d5ba19e..386fb4b 100644
--- a/internal/controller/backend_security_policy_test.go
+++ b/internal/controller/backend_security_policy_test.go
@@ -91,7 +91,9 @@ func TestBackendSecurityController_Reconcile(t *testing.T) {
 }
 
 // mockSTSClient implements the STSOperations interface for testing
-type mockSTSClient struct{}
+type mockSTSClient struct {
+	expTime time.Time
+}
 
 // AssumeRoleWithWebIdentity will return placeholder of type aws credentials.
 //
@@ -102,7 +104,7 @@ func (m *mockSTSClient) AssumeRoleWithWebIdentity(_ context.Context, _ *sts.Assu
 			AccessKeyId:     aws.String("NEWKEY"),
 			SecretAccessKey: aws.String("NEWSECRET"),
 			SessionToken:    aws.String("NEWTOKEN"),
-			Expiration:      aws.Time(time.Now().Add(1 * time.Hour)),
+			Expiration:      &m.expTime,
 		},
 	}, nil
 }
@@ -202,7 +204,7 @@ func TestBackendSecurityController_RotateCredentials(t *testing.T) {
 
 	// new aws oidc rotator
 	ctx := oidcv3.InsecureIssuerURLContext(t.Context(), discoveryServer.URL)
-	rotator, err := rotators.NewAWSOIDCRotator(ctx, cl, &mockSTSClient{}, fake2.NewClientset(), ctrl.Log, bspNamespace, bsp.Name, preRotationWindow, "placeholder", "us-east-1")
+	rotator, err := rotators.NewAWSOIDCRotator(ctx, cl, &mockSTSClient{time.Now().Add(time.Hour)}, fake2.NewClientset(), ctrl.Log, bspNamespace, bsp.Name, preRotationWindow, "placeholder", "us-east-1")
 	require.NoError(t, err)
 
 	// ensure aws credentials secret do not exist
@@ -247,6 +249,31 @@ func TestBackendSecurityController_RotateCredentials(t *testing.T) {
 	require.NotEqual(t, t1, t2)
 }
 
+func TestBackendSecurityPolicyController_RotateExpiredCredential(t *testing.T) {
+	cl := fake.NewClientBuilder().WithScheme(Scheme).Build()
+	c := NewBackendSecurityPolicyController(cl, fake2.NewClientset(), ctrl.Log, internaltesting.NewSyncFnImpl[aigv1a1.AIServiceBackend]().Sync)
+	bspName := "mybackendSecurityPolicy"
+	bspNamespace := "default"
+
+	c.oidcTokenCache[backendSecurityPolicyKey(bspNamespace, bspName)] = &oauth2.Token{AccessToken: "some-access-token", Expiry: time.Now().Add(time.Hour)}
+	rotator, err := rotators.NewAWSOIDCRotator(t.Context(), cl, &mockSTSClient{time.Now().Add(-time.Hour)}, fake2.NewClientset(), ctrl.Log, bspNamespace, bspName, preRotationWindow, "placeholder", "us-east-1")
+	require.NoError(t, err)
+
+	oidcCreds := egv1a1.OIDC{}
+	policy := &aigv1a1.BackendSecurityPolicy{
+		ObjectMeta: metav1.ObjectMeta{
+			Name:      bspName,
+			Namespace: bspNamespace,
+		},
+	}
+
+	// Expiration time will be before current time, so the requeue will be changed to one minute.
+	requeue, err := c.rotateCredential(t.Context(), policy, oidcCreds, rotator)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "newly rotate credentials is already expired")
+	require.Equal(t, time.Minute, requeue)
+}
+
 func TestBackendSecurityController_GetBackendSecurityPolicyAuthOIDC(t *testing.T) {
 	// API Key type does not support OIDC.
 	require.Nil(t, getBackendSecurityPolicyAuthOIDC(aigv1a1.BackendSecurityPolicySpec{Type: aigv1a1.BackendSecurityPolicyTypeAPIKey}))
diff --git a/internal/controller/rotators/aws_oidc_rotator_test.go b/internal/controller/rotators/aws_oidc_rotator_test.go
index 65311c6..1d2fa99 100644
--- a/internal/controller/rotators/aws_oidc_rotator_test.go
+++ b/internal/controller/rotators/aws_oidc_rotator_test.go
@@ -105,6 +105,7 @@ func (m *mockStsOperations) AssumeRoleWithWebIdentity(ctx context.Context, param
 
 func TestAWS_OIDCRotator(t *testing.T) {
 	t.Run("basic rotation", func(t *testing.T) {
+		startTime := time.Now()
 		var mockSTS STSClient = &mockStsOperations{
 			assumeRoleWithWebIdentityFunc: func(_ context.Context, _ *sts.AssumeRoleWithWebIdentityInput, _ ...func(*sts.Options)) (*sts.AssumeRoleWithWebIdentityOutput, error) {
 				return &sts.AssumeRoleWithWebIdentityOutput{
@@ -112,7 +113,7 @@ func TestAWS_OIDCRotator(t *testing.T) {
 						AccessKeyId:     aws.String(newAwsAccessKey),
 						SecretAccessKey: aws.String(newAwsSecretKey),
 						SessionToken:    aws.String(newAwsSessionToken),
-						Expiration:      aws.Time(time.Now().Add(1 * time.Hour)),
+						Expiration:      aws.Time(startTime.Add(1 * time.Hour)),
 					},
 				}, nil
 			},
@@ -135,7 +136,10 @@ func TestAWS_OIDCRotator(t *testing.T) {
 			roleArn:                        awsRoleArn,
 		}
 
-		require.NoError(t, awsOidcRotator.Rotate(t.Context(), newOidcToken))
+		expiration, err := awsOidcRotator.Rotate(t.Context(), newOidcToken)
+		require.NoError(t, err)
+		require.NotNil(t, expiration)
+		require.WithinRange(t, expiration, startTime, startTime.Add(1*time.Hour))
 		verifyAwsCredentialsSecret(t, client, policyNameSpace, policyName, newAwsAccessKey, newAwsSecretKey, newAwsSessionToken, awsProfileName, awsRegion)
 	})
 
@@ -160,12 +164,15 @@ func TestAWS_OIDCRotator(t *testing.T) {
 			region:                         awsRegion,
 			roleArn:                        awsRoleArn,
 		}
-		err := awsOidcRotator.Rotate(t.Context(), newOidcToken)
+
+		expiration, err := awsOidcRotator.Rotate(t.Context(), newOidcToken)
 		require.Error(t, err)
+		require.True(t, expiration.IsZero())
 		assert.Contains(t, err.Error(), "failed to assume role")
 	})
 
 	t.Run("rotation - create when aws credential secret does not exist", func(t *testing.T) {
+		startTime := time.Now()
 		scheme := runtime.NewScheme()
 		scheme.AddKnownTypes(corev1.SchemeGroupVersion,
 			&corev1.Secret{},
@@ -179,7 +186,7 @@ func TestAWS_OIDCRotator(t *testing.T) {
 						AccessKeyId:     aws.String(newAwsAccessKey),
 						SecretAccessKey: aws.String(newAwsSecretKey),
 						SessionToken:    aws.String(newAwsSessionToken),
-						Expiration:      aws.Time(time.Now().Add(1 * time.Hour)),
+						Expiration:      aws.Time(startTime.Add(1 * time.Hour)),
 					},
 				}, nil
 			},
@@ -192,12 +199,15 @@ func TestAWS_OIDCRotator(t *testing.T) {
 			region:                         awsRegion,
 			roleArn:                        awsRoleArn,
 		}
-		err := rotator.Rotate(t.Context(), newOidcToken)
+		expiration, err := rotator.Rotate(t.Context(), newOidcToken)
 		require.NoError(t, err)
+		require.NotNil(t, expiration)
+		require.WithinRange(t, expiration, startTime, startTime.Add(1*time.Hour))
 		verifyAwsCredentialsSecret(t, client, policyNameSpace, policyName, newAwsAccessKey, newAwsSecretKey, newAwsSessionToken, awsProfileName, awsRegion)
 	})
 
 	t.Run("rotation - update when aws credential secret exists", func(t *testing.T) {
+		startTime := time.Now()
 		scheme := runtime.NewScheme()
 		scheme.AddKnownTypes(corev1.SchemeGroupVersion,
 			&corev1.Secret{},
@@ -215,7 +225,7 @@ func TestAWS_OIDCRotator(t *testing.T) {
 						AccessKeyId:     aws.String(newAwsAccessKey),
 						SecretAccessKey: aws.String(newAwsSecretKey),
 						SessionToken:    aws.String(newAwsSessionToken),
-						Expiration:      aws.Time(time.Now().Add(1 * time.Hour)),
+						Expiration:      aws.Time(startTime.Add(1 * time.Hour)),
 					},
 				}, nil
 			},
@@ -228,8 +238,11 @@ func TestAWS_OIDCRotator(t *testing.T) {
 			region:                         awsRegion,
 			roleArn:                        awsRoleArn,
 		}
-		err := rotator.Rotate(t.Context(), newOidcToken)
+
+		expiration, err := rotator.Rotate(t.Context(), newOidcToken)
 		require.NoError(t, err)
+		require.NotNil(t, expiration)
+		require.WithinRange(t, expiration, startTime, startTime.Add(1*time.Hour))
 		verifyAwsCredentialsSecret(t, client, policyNameSpace, policyName, newAwsAccessKey, newAwsSecretKey, newAwsSessionToken, awsProfileName, awsRegion)
 	})
 }