bring back bed rock auth

Signed-off-by: Yuval Kohavi <[email protected]>

Author: yuval-k

api/v1alpha1/agentgateway/agentgateway_backend_types.go

@@ -1,6 +1,7 @@
 package agentgateway
 
 import (
+	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	gwv1 "sigs.k8s.io/gateway-api/apis/v1"
 )
@@ -288,6 +289,20 @@ type BedrockConfig struct {
 	// If not specified, the AWS Guardrail policy will not be used.
 	// +optional
 	Guardrail *AWSGuardrailConfig `json:"guardrail,omitempty"`
+
+	// Auth specifies an explicit AWS authentication method for the backend.
+	// When omitted, we will try to use the default AWS SDK authentication methods.
+	//
+	// +optional
+	Auth *AwsAuth `json:"auth,omitempty"`
+}
+
+// AwsAuth specifies the authentication method to use for the backend.
+type AwsAuth struct {
+	// SecretRef references a Kubernetes Secret containing the AWS credentials.
+	// The Secret must have keys "accessKey", "secretKey", and optionally "sessionToken".
+	// +required
+	SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"`
 }
 
 type AWSGuardrailConfig struct {

api/v1alpha1/agentgateway/zz_generated.deepcopy.go

@@ -137,6 +137,30 @@ spec:
                               bedrock:
                                 description: Bedrock provider
                                 properties:
+                                  auth:
+                                    description: |-
+                                      Auth specifies an explicit AWS authentication method for the backend.
+                                      When omitted, we will try to use the default AWS SDK authentication methods.
+                                    properties:
+                                      secretRef:
+                                        description: |-
+                                          SecretRef references a Kubernetes Secret containing the AWS credentials.
+                                          The Secret must have keys "accessKey", "secretKey", and optionally "sessionToken".
+                                        properties:
+                                          name:
+                                            default: ""
+                                            description: |-
+                                              Name of the referent.
+                                              This field is effectively required, but due to backwards compatibility is
+                                              allowed to be empty. Instances of this type with an empty value here are
+                                              almost certainly wrong.
+                                              More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names
+                                            type: string
+                                        type: object
+                                        x-kubernetes-map-type: atomic
+                                    required:
+                                    - secretRef
+                                    type: object
                                   guardrail:
                                     description: |-
                                       Guardrail configures the Guardrail policy to use for the backend. See <https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails.html>
@@ -1545,6 +1569,30 @@ spec:
                       bedrock:
                         description: Bedrock provider
                         properties:
+                          auth:
+                            description: |-
+                              Auth specifies an explicit AWS authentication method for the backend.
+                              When omitted, we will try to use the default AWS SDK authentication methods.
+                            properties:
+                              secretRef:
+                                description: |-
+                                  SecretRef references a Kubernetes Secret containing the AWS credentials.
+                                  The Secret must have keys "accessKey", "secretKey", and optionally "sessionToken".
+                                properties:
+                                  name:
+                                    default: ""
+                                    description: |-
+                                      Name of the referent.
+                                      This field is effectively required, but due to backwards compatibility is
+                                      allowed to be empty. Instances of this type with an empty value here are
+                                      almost certainly wrong.
+                                      More info: https://kubernetes.io/docs/concepts/overview/working-with-objects/names/#names
+                                    type: string
+                                type: object
+                                x-kubernetes-map-type: atomic
+                            required:
+                            - secretRef
+                            type: object
                           guardrail:
                             description: |-
                               Guardrail configures the Guardrail policy to use for the backend. See <https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails.html>

install/helm/kgateway-crds/templates/agentgateway.dev_agentgatewaybackends.yaml

@@ -12,6 +12,9 @@ inlinePolicies:
       /v1/messages/count_tokens: ANTHROPIC_TOKEN_COUNT
       /v1/models: MODELS
       /v1/responses: RESPONSES
+- auth:
+    aws:
+      implicit: {}
 key: test-ns/bedrock-with-new-routes
 name:
   name: bedrock-with-new-routes

pkg/kgateway/agentgatewaysyncer/backend/testdata/Bedrock_backend_with_new_route_types_(responses_and_anthropic_token_count).yaml

@@ -0,0 +1,18 @@
+ai:
+  providerGroups:
+  - providers:
+    - bedrock:
+        region: us-east-1
+      name: backend
+inlinePolicies:
+- auth:
+    aws:
+      explicitConfig:
+        accessKeyId: secret-accessKey
+        region: us-east-1
+        secretAccessKey: secret-secretKey
+        sessionToken: secret-sessionToken
+key: test-ns/bedrock-with-secret-ref
+name:
+  name: bedrock-with-secret-ref
+  namespace: test-ns

pkg/kgateway/agentgatewaysyncer/backend/testdata/Bedrock_backend_with_new_secret_ref.yaml

@@ -7,6 +7,10 @@ ai:
         model: anthropic.claude-3-haiku-20240307-v1:0
         region: eu-west-1
       name: backend
+inlinePolicies:
+- auth:
+    aws:
+      implicit: {}
 key: test-ns/bedrock-backend-custom
 name:
   name: bedrock-backend-custom

pkg/kgateway/agentgatewaysyncer/backend/testdata/Valid_Bedrock_backend_with_custom_region_and_guardrail.yaml

@@ -15,6 +15,7 @@ import (
 	"github.com/kgateway-dev/kgateway/v2/api/v1alpha1/agentgateway"
 	"github.com/kgateway-dev/kgateway/v2/pkg/agentgateway/plugins"
 	"github.com/kgateway-dev/kgateway/v2/pkg/agentgateway/utils"
+	"github.com/kgateway-dev/kgateway/v2/pkg/kgateway/wellknown"
 	"github.com/kgateway-dev/kgateway/v2/pkg/logging"
 	"github.com/kgateway-dev/kgateway/v2/pkg/utils/kubeutils"
 )
@@ -224,11 +225,17 @@ func translateAIBackends(ctx plugins.PolicyCtx, be *agentgateway.AgentgatewayBac
 
 	aiBackend := &api.AIBackend{}
 	if llm := ai.LLM; llm != nil {
-		provider, err := translateLLMProvider(llm, utils.SingularLLMProviderSubBackendName)
+		provider, auth, err := translateLLMProvider(ctx, llm, utils.SingularLLMProviderSubBackendName, be.Namespace)
 		if err != nil {
 			return nil, fmt.Errorf("failed to translate LLM provider: %w", err)
 		}
-
+		if auth != nil {
+			inlinePolicies = append(inlinePolicies, &api.BackendPolicySpec{
+				Kind: &api.BackendPolicySpec_Auth{
+					Auth: auth,
+				},
+			})
+		}
 		aiBackend.ProviderGroups = []*api.AIBackend_ProviderGroup{{
 			Providers: []*api.AIBackend_Provider{provider},
 		}}
@@ -237,7 +244,7 @@ func translateAIBackends(ctx plugins.PolicyCtx, be *agentgateway.AgentgatewayBac
 			providerGroup := &api.AIBackend_ProviderGroup{}
 
 			for _, provider := range group.Providers {
-				tp, err := translateLLMProvider(&provider.LLMProvider, string(provider.Name))
+				tp, auth, err := translateLLMProvider(ctx, &provider.LLMProvider, string(provider.Name), be.Namespace)
 				if err != nil {
 					return nil, fmt.Errorf("failed to translate LLM provider: %w", err)
 				}
@@ -247,6 +254,13 @@ func translateAIBackends(ctx plugins.PolicyCtx, be *agentgateway.AgentgatewayBac
 					logger.Warn("failed to translate AI backend policies", "err", err)
 				}
 				tp.InlinePolicies = pol
+				if auth != nil {
+					tp.InlinePolicies = append(tp.InlinePolicies, &api.BackendPolicySpec{
+						Kind: &api.BackendPolicySpec_Auth{
+							Auth: auth,
+						},
+					})
+				}
 
 				providerGroup.Providers = append(providerGroup.Providers, tp)
 			}
@@ -303,7 +317,7 @@ func translateAIBackendPolicies(
 	})
 }
 
-func translateLLMProvider(llm *agentgateway.LLMProvider, providerName string) (*api.AIBackend_Provider, error) {
+func translateLLMProvider(ctx plugins.PolicyCtx, llm *agentgateway.LLMProvider, providerName, namespace string) (*api.AIBackend_Provider, *api.BackendAuthPolicy, error) {
 	provider := &api.AIBackend_Provider{
 		Name: providerName,
 	}
@@ -318,6 +332,7 @@ func translateLLMProvider(llm *agentgateway.LLMProvider, providerName string) (*
 	if llm.Path != "" {
 		provider.PathOverride = &llm.Path
 	}
+	var auth *api.BackendAuthPolicy
 
 	// Extract auth token and model based on provider
 	if llm.OpenAI != nil {
@@ -363,6 +378,12 @@ func translateLLMProvider(llm *agentgateway.LLMProvider, providerName string) (*
 			guardrailVersion = &llm.Bedrock.Guardrail.GuardrailVersion
 		}
 
+		var err error
+		auth, err = buildBedrockAuthPolicy(ctx.Krt, region, llm.Bedrock.Auth, ctx.Collections.Secrets, namespace)
+		if err != nil {
+			return nil, nil, err
+		}
+
 		provider.Provider = &api.AIBackend_Provider_Bedrock{
 			Bedrock: &api.AIBackend_Bedrock{
 				Model:               llm.Bedrock.Model,
@@ -372,10 +393,10 @@ func translateLLMProvider(llm *agentgateway.LLMProvider, providerName string) (*
 			},
 		}
 	} else {
-		return nil, fmt.Errorf("no supported LLM provider configured")
+		return nil, nil, fmt.Errorf("no supported LLM provider configured")
 	}
 
-	return provider, nil
+	return provider, auth, nil
 }
 
 func toMCPProtocol(appProtocol string) api.MCPTarget_Protocol {
@@ -391,3 +412,70 @@ func toMCPProtocol(appProtocol string) api.MCPTarget_Protocol {
 		return api.MCPTarget_UNDEFINED
 	}
 }
+
+func buildBedrockAuthPolicy(krtctx krt.HandlerContext, region string, auth *agentgateway.AwsAuth, secrets krt.Collection[*corev1.Secret], namespace string) (*api.BackendAuthPolicy, error) {
+	var errs []error
+	if auth == nil {
+		logger.Warn("using implicit AWS auth for AI backend")
+		return &api.BackendAuthPolicy{
+			Kind: &api.BackendAuthPolicy_Aws{
+				Aws: &api.Aws{
+					Kind: &api.Aws_Implicit{
+						Implicit: &api.AwsImplicit{},
+					},
+				},
+			},
+		}, nil
+	}
+
+	if auth.SecretRef == nil {
+		logger.Warn("not using any auth for AWS - it's most likely not what you want")
+		return nil, nil
+	}
+
+	// Get secret using the SecretIndex
+	secret, err := kubeutils.GetSecret(secrets, krtctx, auth.SecretRef.Name, namespace)
+	if err != nil {
+		// Return nil auth policy if secret not found - this will be handled upstream
+		// TODO(npolshak): Add backend status errors https://github.com/kgateway-dev/kgateway/issues/11966
+		return nil, err
+	}
+
+	var accessKeyId, secretAccessKey string
+	var sessionToken *string
+
+	// Extract access key
+	if value, exists := kubeutils.GetSecretValue(secret, wellknown.AccessKey); !exists {
+		errs = append(errs, errors.New("accessKey is missing or not a valid string"))
+	} else {
+		accessKeyId = value
+	}
+
+	// Extract secret key
+	if value, exists := kubeutils.GetSecretValue(secret, wellknown.SecretKey); !exists {
+		errs = append(errs, errors.New("secretKey is missing or not a valid string"))
+	} else {
+		secretAccessKey = value
+	}
+
+	// Extract session token (optional)
+	if value, exists := kubeutils.GetSecretValue(secret, wellknown.SessionToken); exists {
+		sessionToken = ptr.Of(value)
+	}
+
+	return &api.BackendAuthPolicy{
+		Kind: &api.BackendAuthPolicy_Aws{
+			Aws: &api.Aws{
+				Kind: &api.Aws_ExplicitConfig{
+					ExplicitConfig: &api.AwsExplicitConfig{
+						AccessKeyId:     accessKeyId,
+						SecretAccessKey: secretAccessKey,
+						SessionToken:    sessionToken,
+						Region:          region,
+					},
+				},
+			},
+		},
+	}, errors.Join(errs...)
+
+}

pkg/kgateway/agentgatewaysyncer/backend/translate.go

@@ -491,6 +491,36 @@ func TestBuildAIBackend(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "Bedrock backend with new secret ref",
+			backend: &agentgateway.AgentgatewayBackend{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "bedrock-with-secret-ref",
+					Namespace: "test-ns",
+				},
+				Spec: agentgateway.AgentgatewayBackendSpec{
+					AI: &agentgateway.AIBackend{
+						LLM: &agentgateway.LLMProvider{
+							Bedrock: &agentgateway.BedrockConfig{
+								Region: "us-east-1",
+								Auth: &agentgateway.AwsAuth{
+									SecretRef: &corev1.LocalObjectReference{
+										Name: "bedrock-secret",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+			inputs: []any{
+				createMockSecret("test-ns", "bedrock-secret", map[string]string{
+					"accessKey":    "secret-accessKey",
+					"secretKey":    "secret-secretKey",
+					"sessionToken": "secret-sessionToken",
+				}),
+			},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {