-
Notifications
You must be signed in to change notification settings - Fork 618
bring back bed rock auth #13089
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
bring back bed rock auth #13089
Changes from 1 commit
3177084
804d8c2
8e61c84
7af67c2
bcab27c
91bc916
065d8d8
ba9cf57
d44b375
3b733f1
ec09ef7
3835d2a
ee88619
90f4e71
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
yuval-k marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| SecretRef *corev1.LocalObjectReference `json:"secretRef,omitempty"` | ||
|
||
| } | ||
|
|
||
| type AWSGuardrailConfig struct { | ||
|
|
||
Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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...) | ||
|
|
||
|
||
| } | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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", | ||
| }), | ||
| }, | ||
| }, | ||
yuval-k marked this conversation as resolved.
Outdated
Show resolved
Hide resolved
|
||
| } | ||
| for _, tt := range tests { | ||
| t.Run(tt.name, func(t *testing.T) { | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
hmmm not sure this belongs here vs
policies.backend.auth.aws. That would be consistent. On the other hand, this is currently only usable for bedrock. But in the future probably we would support lambda, etc, which would necessitate putting it in a more broad area.I think for consistency and future proofing we need to move this.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
i'll try but it will be slightly annoying to do - as if there is no policy, i'll want to implicitly attach the implicit one inline
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
nvm - agentgw will default to the implicit one, so doing this is not annoying