Support bearer tokens through environment for Bedrock

gems/aws-sdk-bedrock/lib/aws-sdk-bedrock/plugins/bearer_authorization.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module Aws::Bedrock
+  module Plugins
+    # @api private
+    class BearerAuthorization < Seahorse::Client::Plugin
+      def after_initialize(client)
+        return unless (token = ENV['AWS_BEARER_TOKEN_BEDROCK'])
+
+        client.config.token_provider ||= Aws::StaticTokenProvider.new(token)
+      end
+
+      class Handler < Seahorse::Client::Handler
+        def call(context)
+          # This also sets the preferred auth scheme even if the code token has precedence.
+          context[:auth_scheme] = { 'name' => 'bearer' } if ENV['AWS_BEARER_TOKEN_BEDROCK']
+          with_metric { @handler.call(context) }
+        end
+
+        private
+
+        def with_metric(&block)
+          if ENV['AWS_BEARER_TOKEN_BEDROCK']
+            Aws::Plugins::UserAgent.metric('BEARER_SERVICE_ENV_VARS', &block)
+          else
+            block.call
+          end
+        end
+      end
+
+      handle(Handler)
+    end
+  end
+end

gems/aws-sdk-bedrock/spec/client_spec.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require_relative 'spec_helper'
+
+module Aws
+  module Bedrock
+    describe Client do
+      it 'uses a bearer token from the environment' do
+        ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
+        client = Client.new(stub_responses: true, token_provider: nil)
+        expect(client.config.token_provider.token.token).to eq('bedrock-token')
+        resp = client.list_imported_models
+        expect(resp.context.http_request.headers['Authorization']).to eq('Bearer bedrock-token')
+      end
+
+      it 'does not use a token for a different service' do
+        ENV['AWS_BEARER_TOKEN_FOO'] = 'foo-token'
+        client = Client.new(stub_responses: true, token_provider: nil)
+        expect(client.config.token_provider).to be_nil
+        resp = client.list_imported_models
+        expect(resp.context.http_request.headers['Authorization']).to_not eq('Bearer foo-token')
+      end
+
+      it 'still prefers bearer token when given an auth scheme preference' do
+        ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
+        ENV['AWS_AUTH_SCHEME_PREFERENCE'] = 'sigv4,httpBearerAuth'
+        client = Client.new(stub_responses: true, token_provider: nil)
+        resp = client.list_imported_models
+        expect(resp.context.http_request.headers['Authorization']).to eq('Bearer bedrock-token')
+      end
+
+      it 'uses explicit config over the environment token' do
+        ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
+        client = Client.new(
+          stub_responses: true,
+          token_provider: Aws::StaticTokenProvider.new('explicit-code-token')
+        )
+        resp = client.list_imported_models
+        expect(resp.context.http_request.headers['Authorization']).to eq('Bearer explicit-code-token')
+      end
+    end
+  end
+end

gems/aws-sdk-bedrockruntime/lib/aws-sdk-bedrockruntime/plugins/bearer_authorization.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module Aws::BedrockRuntime
+  module Plugins
+    # @api private
+    class BearerAuthorization < Seahorse::Client::Plugin
+      def after_initialize(client)
+        return unless (token = ENV['AWS_BEARER_TOKEN_BEDROCK'])
+
+        client.config.token_provider ||= Aws::StaticTokenProvider.new(token)
+      end
+
+      class Handler < Seahorse::Client::Handler
+        def call(context)
+          # This also sets the preferred auth scheme even if the code token has precedence.
+          context[:auth_scheme] = { 'name' => 'bearer' } if ENV['AWS_BEARER_TOKEN_BEDROCK']
+          with_metric { @handler.call(context) }
+        end
+
+        private
+
+        def with_metric(&block)
+          if ENV['AWS_BEARER_TOKEN_BEDROCK']
+            Aws::Plugins::UserAgent.metric('BEARER_SERVICE_ENV_VARS', &block)
+          else
+            block.call
+          end
+        end
+      end
+
+      handle(Handler)
+    end
+  end
+end

gems/aws-sdk-bedrockruntime/spec/client_spec.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require_relative 'spec_helper'
+
+module Aws
+  module BedrockRuntime
+    describe Client do
+      def invoke_model(client)
+        stub = client.stub_data(:invoke_model, body: 'test')
+        client.stub_responses(:invoke_model, stub)
+        client.invoke_model(model_id: 'test')
+      end
+
+      it 'uses a bearer token from the environment' do
+        ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
+        client = Client.new(stub_responses: true, token_provider: nil)
+        expect(client.config.token_provider.token.token).to eq('bedrock-token')
+        resp = invoke_model(client)
+        expect(resp.context.http_request.headers['Authorization']).to eq('Bearer bedrock-token')
+      end
+
+      it 'does not use a token for a different service' do
+        ENV['AWS_BEARER_TOKEN_FOO'] = 'foo-token'
+        client = Client.new(stub_responses: true, token_provider: nil)
+        expect(client.config.token_provider).to be_nil
+        resp = invoke_model(client)
+        expect(resp.context.http_request.headers['Authorization']).to_not eq('Bearer foo-token')
+      end
+
+      it 'still prefers bearer token when given an auth scheme preference' do
+        ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
+        ENV['AWS_AUTH_SCHEME_PREFERENCE'] = 'sigv4,httpBearerAuth'
+        client = Client.new(stub_responses: true, token_provider: nil)
+        resp = invoke_model(client)
+        expect(resp.context.http_request.headers['Authorization']).to eq('Bearer bedrock-token')
+      end
+
+      it 'uses explicit config over the environment token' do
+        ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
+        client = Client.new(
+          stub_responses: true,
+          token_provider: Aws::StaticTokenProvider.new('explicit-code-token')
+        )
+        resp = invoke_model(client)
+        expect(resp.context.http_request.headers['Authorization']).to eq('Bearer explicit-code-token')
+      end
+    end
+  end
+end

gems/aws-sdk-core/lib/aws-sdk-core/plugins/credentials_configuration.rb

@@ -99,13 +99,8 @@ class CredentialsConfiguration < Seahorse::Client::Plugin
 will be used to search for tokens configured for your profile in shared configuration files.
       DOCS
       ) do |config|
-        if config.stub_responses
-          StaticTokenProvider.new('token')
-        else
-          TokenProviderChain.new(config).resolve
-        end
+        TokenProviderChain.new(config).resolve
       end
-
     end
   end
 end

gems/aws-sdk-core/lib/aws-sdk-core/plugins/sign.rb

@@ -77,16 +77,12 @@ def initialize
 
         def sign(context)
           if context.http_request.endpoint.scheme != 'https'
-            raise ArgumentError,
-                  'Unable to use bearer authorization on non https endpoint.'
+            raise ArgumentError, 'Unable to use bearer authorization on non https endpoint.'
           end
-
           token_provider = context.config.token_provider
-
           raise Errors::MissingBearerTokenError unless token_provider&.set?
 
-          context.http_request.headers['Authorization'] =
-            "Bearer #{token_provider.token.token}"
+          context.http_request.headers['Authorization'] = "Bearer #{token_provider.token.token}"
         end
 
         def presign_url(*args)

gems/aws-sdk-core/lib/aws-sdk-core/plugins/stub_responses.rb

@@ -29,6 +29,12 @@ class StubResponses < Seahorse::Client::Plugin
         end
       end
 
+      option(:token_provider) do |config|
+        if config.stub_responses
+          StaticTokenProvider.new('stubbed-token')
+        end
+      end
+
       option(:stubs) { {} }
       option(:stubs_mutex) { Mutex.new }
       option(:api_requests) { [] }

gems/aws-sdk-core/lib/aws-sdk-core/plugins/user_agent.rb

@@ -54,7 +54,8 @@ class UserAgent < Seahorse::Client::Plugin
           "CREDENTIALS_HTTP" : "z",
           "CREDENTIALS_IMDS" : "0",
           "SSO_LOGIN_DEVICE" : "1",
-          "SSO_LOGIN_AUTH" : "2"
+          "SSO_LOGIN_AUTH" : "2",
+          "BEARER_SERVICE_ENV_VARS": "3"
         }
       METRICS
 

gems/aws-sdk-core/lib/aws-sdk-core/token_provider_chain.rb

@@ -27,17 +27,13 @@ def providers
 
     def static_profile_sso_token(options)
       if Aws.shared_config.config_enabled? && options[:config] && options[:config].profile
-        Aws.shared_config.sso_token_from_config(
-          profile: options[:config].profile
-        )
+        Aws.shared_config.sso_token_from_config(profile: options[:config].profile)
       end
     end
 
-
     def sso_token(options)
-      profile_name = determine_profile_name(options)
       if Aws.shared_config.config_enabled?
-        Aws.shared_config.sso_token_from_config(profile: profile_name)
+        Aws.shared_config.sso_token_from_config(profile: determine_profile_name(options))
       end
     rescue Errors::NoSuchProfileError
       nil

services.json

@@ -127,7 +127,10 @@
     "models": "batch/2016-08-10"
   },
   "Bedrock": {
-    "models": "bedrock/2023-04-20"
+    "models": "bedrock/2023-04-20",
+    "addPlugins": [
+      "Aws::Bedrock::Plugins::BearerAuthorization"
+    ]
   },
   "BedrockAgent": {
     "models": "bedrock-agent/2023-06-05"
@@ -142,7 +145,10 @@
     "models": "bedrock-data-automation-runtime/2024-06-13"
   },
   "BedrockRuntime": {
-    "models": "bedrock-runtime/2023-09-30"
+    "models": "bedrock-runtime/2023-09-30",
+    "addPlugins": [
+      "Aws::BedrockRuntime::Plugins::BearerAuthorization"
+    ]
   },
   "Billing": {
     "models": "billing/2023-09-07"