Repository: envoyproxy/ai-gateway
Pull request #461 (2025-03-05 18:48:02)

backport: fix requeue bug and extproc tag sync
**Commit Message**

Backporting fix for extproc image tag not being sync'd properly and case when requeue is negative.

For more info, please check out:
- controller: fix syncing extproc image tag (https://github.com/envoyproxy/ai-gateway/pull/447)
- controller: update rotate function to return expiration time (https://github.com/envoyproxy/ai-gateway/pull/455)

mathetake: why they are three commits vs two commits mentioned? what about https://github.com/envoyproxy/ai-gateway/pull/461/commits/c35e9c738dbee076c2eaecb28a94166408817685 ?

Patch:
diff --git a/internal/controller/ai_gateway_route.go b/internal/controller/ai_gateway_route.go
index af5aba4..154e870 100644
--- a/internal/controller/ai_gateway_route.go
+++ b/internal/controller/ai_gateway_route.go
@@ -179,7 +179,7 @@ func extProcName(route *aigv1a1.AIGatewayRoute) string {
 	return fmt.Sprintf("ai-eg-route-extproc-%s", route.Name)
 }
 
-func applyExtProcDeploymentConfigUpdate(d *appsv1.DeploymentSpec, filterConfig *aigv1a1.AIGatewayFilterConfig) {
+func (c *AIGatewayRouteController) applyExtProcDeploymentConfigUpdate(d *appsv1.DeploymentSpec, filterConfig *aigv1a1.AIGatewayFilterConfig) {
 	if filterConfig == nil || filterConfig.ExternalProcessor == nil {
 		d.Replicas = nil
 		d.Template.Spec.Containers[0].Resources = corev1.ResourceRequirements{}
@@ -191,6 +191,7 @@ func applyExtProcDeploymentConfigUpdate(d *appsv1.DeploymentSpec, filterConfig *
 	} else {
 		d.Template.Spec.Containers[0].Resources = corev1.ResourceRequirements{}
 	}
+	d.Template.Spec.Containers[0].Image = c.extProcImage
 	d.Replicas = extProc.Replicas
 }
 
@@ -549,7 +550,7 @@ func (c *AIGatewayRouteController) syncExtProcDeployment(ctx context.Context, ai
 			if err == nil {
 				deployment.Spec.Template.Spec = *updatedSpec
 			}
-			applyExtProcDeploymentConfigUpdate(&deployment.Spec, aiGatewayRoute.Spec.FilterConfig)
+			c.applyExtProcDeploymentConfigUpdate(&deployment.Spec, aiGatewayRoute.Spec.FilterConfig)
 			_, err = c.kube.AppsV1().Deployments(aiGatewayRoute.Namespace).Create(ctx, deployment, metav1.CreateOptions{})
 			if err != nil {
 				return fmt.Errorf("failed to create deployment: %w", err)
@@ -564,7 +565,7 @@ func (c *AIGatewayRouteController) syncExtProcDeployment(ctx context.Context, ai
 		if err == nil {
 			deployment.Spec.Template.Spec = *updatedSpec
 		}
-		applyExtProcDeploymentConfigUpdate(&deployment.Spec, aiGatewayRoute.Spec.FilterConfig)
+		c.applyExtProcDeploymentConfigUpdate(&deployment.Spec, aiGatewayRoute.Spec.FilterConfig)
 		if _, err = c.kube.AppsV1().Deployments(aiGatewayRoute.Namespace).Update(ctx, deployment, metav1.UpdateOptions{}); err != nil {
 			return fmt.Errorf("failed to update deployment: %w", err)
 		}
diff --git a/internal/controller/backend_security_policy.go b/internal/controller/backend_security_policy.go
index 2f07ee3..dbaf17f 100644
--- a/internal/controller/backend_security_policy.go
+++ b/internal/controller/backend_security_policy.go
@@ -105,7 +105,6 @@ func (c *BackendSecurityPolicyController) Reconcile(ctx context.Context, req ctr
 // 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]
@@ -122,15 +121,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.
diff --git a/manifests/charts/ai-gateway-helm/values.yaml b/manifests/charts/ai-gateway-helm/values.yaml
index 4d2bda4..2e66d81 100644
--- a/manifests/charts/ai-gateway-helm/values.yaml
+++ b/manifests/charts/ai-gateway-helm/values.yaml
@@ -50,7 +50,7 @@ controller:
   #    configmap:
   #      defaultMode: placeholder
   #      name: configmap-name
-  volumes: {}
+  volumes: []
   service:
     type: ClusterIP
     ports:


Test patch:
diff --git a/internal/controller/ai_gateway_route_test.go b/internal/controller/ai_gateway_route_test.go
index 03f0964..405e341 100644
--- a/internal/controller/ai_gateway_route_test.go
+++ b/internal/controller/ai_gateway_route_test.go
@@ -169,10 +169,12 @@ func Test_applyExtProcDeploymentConfigUpdate(t *testing.T) {
 			},
 		},
 	}
+	extProcImage := "extproc:v0.1.0"
+	c := &AIGatewayRouteController{client: fake.NewClientBuilder().WithScheme(scheme).Build(), extProcImage: extProcImage}
 	t.Run("not panic", func(_ *testing.T) {
-		applyExtProcDeploymentConfigUpdate(dep, nil)
-		applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{})
-		applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{
+		c.applyExtProcDeploymentConfigUpdate(dep, nil)
+		c.applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{})
+		c.applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{
 			ExternalProcessor: &aigv1a1.AIGatewayFilterConfigExternalProcessor{},
 		})
 	})
@@ -183,7 +185,7 @@ func Test_applyExtProcDeploymentConfigUpdate(t *testing.T) {
 				corev1.ResourceMemory: resource.MustParse("100Mi"),
 			},
 		}
-		applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{
+		c.applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{
 			ExternalProcessor: &aigv1a1.AIGatewayFilterConfigExternalProcessor{
 				Resources: &req,
 				Replicas:  ptr.To[int32](123),
@@ -192,11 +194,12 @@ func Test_applyExtProcDeploymentConfigUpdate(t *testing.T) {
 		)
 		require.Equal(t, req, dep.Template.Spec.Containers[0].Resources)
 		require.Equal(t, int32(123), *dep.Replicas)
+		require.Equal(t, extProcImage, dep.Template.Spec.Containers[0].Image)
 	})
 	t.Run("remove partial config", func(t *testing.T) {
 		t.Run("replicas", func(t *testing.T) {
 			dep.Replicas = ptr.To[int32](123)
-			applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{
+			c.applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{
 				ExternalProcessor: &aigv1a1.AIGatewayFilterConfigExternalProcessor{},
 			})
 			require.Nil(t, dep.Replicas)
@@ -209,7 +212,7 @@ func Test_applyExtProcDeploymentConfigUpdate(t *testing.T) {
 				},
 			}
 			dep.Replicas = ptr.To[int32](123)
-			applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{
+			c.applyExtProcDeploymentConfigUpdate(dep, &aigv1a1.AIGatewayFilterConfig{
 				ExternalProcessor: &aigv1a1.AIGatewayFilterConfigExternalProcessor{Replicas: ptr.To[int32](123)},
 			})
 			require.Empty(t, dep.Template.Spec.Containers[0].Resources.Limits)
@@ -218,7 +221,7 @@ func Test_applyExtProcDeploymentConfigUpdate(t *testing.T) {
 		})
 	})
 	t.Run("remove the whole config", func(t *testing.T) {
-		for _, c := range []*aigv1a1.AIGatewayFilterConfig{nil, {}} {
+		for _, filterConfig := range []*aigv1a1.AIGatewayFilterConfig{nil, {}} {
 			dep.Replicas = ptr.To[int32](123)
 			dep.Template.Spec.Containers[0].Resources = corev1.ResourceRequirements{
 				Limits: corev1.ResourceList{
@@ -230,7 +233,7 @@ func Test_applyExtProcDeploymentConfigUpdate(t *testing.T) {
 					corev1.ResourceMemory: resource.MustParse("50Mi"),
 				},
 			}
-			applyExtProcDeploymentConfigUpdate(dep, c)
+			c.applyExtProcDeploymentConfigUpdate(dep, filterConfig)
 			require.Nil(t, dep.Replicas)
 			require.Empty(t, dep.Template.Spec.Containers[0].Resources.Limits)
 			require.Empty(t, dep.Template.Spec.Containers[0].Resources.Requests)
diff --git a/internal/controller/backend_security_policy_test.go b/internal/controller/backend_security_policy_test.go
index b6db3de..fc5990a 100644
--- a/internal/controller/backend_security_policy_test.go
+++ b/internal/controller/backend_security_policy_test.go
@@ -84,7 +84,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.
 //
@@ -95,7 +97,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
 }
@@ -195,7 +197,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
@@ -240,6 +242,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)
 	})
 }