Skip to content

Commit e2fb86f

Browse files
author
Matt Muller
committed
Support bearer tokens through environment for Bedrock
1 parent aa4796c commit e2fb86f

File tree

13 files changed

+187
-21
lines changed

13 files changed

+187
-21
lines changed

‎gems/aws-sdk-bedrock/lib/aws-sdk-bedrock/client.rb‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
require 'aws-sdk-core/plugins/telemetry'
3636
require 'aws-sdk-core/plugins/sign'
3737
require 'aws-sdk-core/plugins/protocols/rest_json'
38+
require 'aws-sdk-bedrock/plugins/bearer_authorization'
3839

3940
module Aws::Bedrock
4041
# An API client for Bedrock. To construct a client, you need to configure a `:region` and `:credentials`.
@@ -85,6 +86,7 @@ class Client < Seahorse::Client::Base
8586
add_plugin(Aws::Plugins::Telemetry)
8687
add_plugin(Aws::Plugins::Sign)
8788
add_plugin(Aws::Plugins::Protocols::RestJson)
89+
add_plugin(Aws::Bedrock::Plugins::BearerAuthorization)
8890
add_plugin(Aws::Bedrock::Plugins::Endpoints)
8991

9092
# @overload initialize(options)

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# frozen_string_literal: true
2+
3+
module Aws::Bedrock
4+
module Plugins
5+
# @api private
6+
class BearerAuthorization < Seahorse::Client::Plugin
7+
def after_initialize(client)
8+
return unless (token = ENV['AWS_BEARER_TOKEN_BEDROCK'])
9+
10+
client.config.token_provider ||= Aws::StaticTokenProvider.new(token)
11+
end
12+
13+
class Handler < Seahorse::Client::Handler
14+
def call(context)
15+
# This also sets the preferred auth scheme even if the code token has precedence.
16+
context[:auth_scheme] = { 'name' => 'bearer' } if ENV['AWS_BEARER_TOKEN_BEDROCK']
17+
with_metric { @handler.call(context) }
18+
end
19+
20+
private
21+
22+
def with_metric(&block)
23+
if ENV['AWS_BEARER_TOKEN_BEDROCK']
24+
Aws::Plugins::UserAgent.metric('BEARER_SERVICE_ENV_VARS', &block)
25+
else
26+
block.call
27+
end
28+
end
29+
end
30+
31+
handle(Handler)
32+
end
33+
end
34+
end

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

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# frozen_string_literal: true
2+
3+
require_relative 'spec_helper'
4+
5+
module Aws
6+
module Bedrock
7+
describe Client do
8+
it 'uses a bearer token from the environment' do
9+
ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
10+
client = Client.new(stub_responses: true, token_provider: nil)
11+
expect(client.config.token_provider.token.token).to eq('bedrock-token')
12+
resp = client.list_imported_models
13+
expect(resp.context.http_request.headers['Authorization']).to eq('Bearer bedrock-token')
14+
end
15+
16+
it 'does not use a token for a different service' do
17+
ENV['AWS_BEARER_TOKEN_FOO'] = 'foo-token'
18+
client = Client.new(stub_responses: true, token_provider: nil)
19+
expect(client.config.token_provider).to be_nil
20+
resp = client.list_imported_models
21+
expect(resp.context.http_request.headers['Authorization']).to_not eq('Bearer foo-token')
22+
end
23+
24+
it 'still prefers bearer token when given an auth scheme preference' do
25+
ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
26+
ENV['AWS_AUTH_SCHEME_PREFERENCE'] = 'sigv4,httpBearerAuth'
27+
client = Client.new(stub_responses: true, token_provider: nil)
28+
resp = client.list_imported_models
29+
expect(resp.context.http_request.headers['Authorization']).to eq('Bearer bedrock-token')
30+
end
31+
32+
it 'uses explicit config over the environment token' do
33+
ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
34+
client = Client.new(
35+
stub_responses: true,
36+
token_provider: Aws::StaticTokenProvider.new('explicit-code-token')
37+
)
38+
resp = client.list_imported_models
39+
expect(resp.context.http_request.headers['Authorization']).to eq('Bearer explicit-code-token')
40+
end
41+
end
42+
end
43+
end

‎gems/aws-sdk-bedrockruntime/lib/aws-sdk-bedrockruntime/async_client.rb‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
require 'aws-sdk-core/plugins/telemetry'
3232
require 'aws-sdk-core/plugins/sign'
3333
require 'aws-sdk-core/plugins/protocols/rest_json'
34+
require 'aws-sdk-bedrockruntime/plugins/bearer_authorization'
3435
require 'aws-sdk-core/plugins/event_stream_configuration'
3536

3637
Aws::Plugins::GlobalConfiguration.add_identifier(:bedrockruntime)
@@ -77,6 +78,7 @@ class AsyncClient < Seahorse::Client::AsyncBase
7778
add_plugin(Aws::Plugins::Telemetry)
7879
add_plugin(Aws::Plugins::Sign)
7980
add_plugin(Aws::Plugins::Protocols::RestJson)
81+
add_plugin(Aws::BedrockRuntime::Plugins::BearerAuthorization)
8082
add_plugin(Aws::Plugins::EventStreamConfiguration)
8183
add_plugin(Aws::BedrockRuntime::Plugins::Endpoints)
8284

‎gems/aws-sdk-bedrockruntime/lib/aws-sdk-bedrockruntime/client.rb‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
require 'aws-sdk-core/plugins/telemetry'
3636
require 'aws-sdk-core/plugins/sign'
3737
require 'aws-sdk-core/plugins/protocols/rest_json'
38+
require 'aws-sdk-bedrockruntime/plugins/bearer_authorization'
3839
require 'aws-sdk-core/plugins/event_stream_configuration'
3940

4041
module Aws::BedrockRuntime
@@ -86,6 +87,7 @@ class Client < Seahorse::Client::Base
8687
add_plugin(Aws::Plugins::Telemetry)
8788
add_plugin(Aws::Plugins::Sign)
8889
add_plugin(Aws::Plugins::Protocols::RestJson)
90+
add_plugin(Aws::BedrockRuntime::Plugins::BearerAuthorization)
8991
add_plugin(Aws::Plugins::EventStreamConfiguration)
9092
add_plugin(Aws::BedrockRuntime::Plugins::Endpoints)
9193

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

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# frozen_string_literal: true
2+
3+
module Aws::BedrockRuntime
4+
module Plugins
5+
# @api private
6+
class BearerAuthorization < Seahorse::Client::Plugin
7+
def after_initialize(client)
8+
return unless (token = ENV['AWS_BEARER_TOKEN_BEDROCK'])
9+
10+
client.config.token_provider ||= Aws::StaticTokenProvider.new(token)
11+
end
12+
13+
class Handler < Seahorse::Client::Handler
14+
def call(context)
15+
# This also sets the preferred auth scheme even if the code token has precedence.
16+
context[:auth_scheme] = { 'name' => 'bearer' } if ENV['AWS_BEARER_TOKEN_BEDROCK']
17+
with_metric { @handler.call(context) }
18+
end
19+
20+
private
21+
22+
def with_metric(&block)
23+
if ENV['AWS_BEARER_TOKEN_BEDROCK']
24+
Aws::Plugins::UserAgent.metric('BEARER_SERVICE_ENV_VARS', &block)
25+
else
26+
block.call
27+
end
28+
end
29+
end
30+
31+
handle(Handler)
32+
end
33+
end
34+
end

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

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
# frozen_string_literal: true
2+
3+
require_relative 'spec_helper'
4+
5+
module Aws
6+
module BedrockRuntime
7+
describe Client do
8+
def invoke_model(client)
9+
stub = client.stub_data(:invoke_model, body: 'test')
10+
client.stub_responses(:invoke_model, stub)
11+
client.invoke_model(model_id: 'test')
12+
end
13+
14+
it 'uses a bearer token from the environment' do
15+
ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
16+
client = Client.new(stub_responses: true, token_provider: nil)
17+
expect(client.config.token_provider.token.token).to eq('bedrock-token')
18+
resp = invoke_model(client)
19+
expect(resp.context.http_request.headers['Authorization']).to eq('Bearer bedrock-token')
20+
end
21+
22+
it 'does not use a token for a different service' do
23+
ENV['AWS_BEARER_TOKEN_FOO'] = 'foo-token'
24+
client = Client.new(stub_responses: true, token_provider: nil)
25+
expect(client.config.token_provider).to be_nil
26+
resp = invoke_model(client)
27+
expect(resp.context.http_request.headers['Authorization']).to_not eq('Bearer foo-token')
28+
end
29+
30+
it 'still prefers bearer token when given an auth scheme preference' do
31+
ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
32+
ENV['AWS_AUTH_SCHEME_PREFERENCE'] = 'sigv4,httpBearerAuth'
33+
client = Client.new(stub_responses: true, token_provider: nil)
34+
resp = invoke_model(client)
35+
expect(resp.context.http_request.headers['Authorization']).to eq('Bearer bedrock-token')
36+
end
37+
38+
it 'uses explicit config over the environment token' do
39+
ENV['AWS_BEARER_TOKEN_BEDROCK'] = 'bedrock-token'
40+
client = Client.new(
41+
stub_responses: true,
42+
token_provider: Aws::StaticTokenProvider.new('explicit-code-token')
43+
)
44+
resp = invoke_model(client)
45+
expect(resp.context.http_request.headers['Authorization']).to eq('Bearer explicit-code-token')
46+
end
47+
end
48+
end
49+
end

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

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -99,13 +99,8 @@ class CredentialsConfiguration < Seahorse::Client::Plugin
9999
will be used to search for tokens configured for your profile in shared configuration files.
100100
DOCS
101101
) do |config|
102-
if config.stub_responses
103-
StaticTokenProvider.new('token')
104-
else
105-
TokenProviderChain.new(config).resolve
106-
end
102+
TokenProviderChain.new(config).resolve
107103
end
108-
109104
end
110105
end
111106
end

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

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -77,16 +77,12 @@ def initialize
7777

7878
def sign(context)
7979
if context.http_request.endpoint.scheme != 'https'
80-
raise ArgumentError,
81-
'Unable to use bearer authorization on non https endpoint.'
80+
raise ArgumentError, 'Unable to use bearer authorization on non https endpoint.'
8281
end
83-
8482
token_provider = context.config.token_provider
85-
8683
raise Errors::MissingBearerTokenError unless token_provider&.set?
8784

88-
context.http_request.headers['Authorization'] =
89-
"Bearer #{token_provider.token.token}"
85+
context.http_request.headers['Authorization'] = "Bearer #{token_provider.token.token}"
9086
end
9187

9288
def presign_url(*args)

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,12 @@ class StubResponses < Seahorse::Client::Plugin
2929
end
3030
end
3131

32+
option(:token_provider) do |config|
33+
if config.stub_responses
34+
StaticTokenProvider.new('stubbed-token')
35+
end
36+
end
37+
3238
option(:stubs) { {} }
3339
option(:stubs_mutex) { Mutex.new }
3440
option(:api_requests) { [] }

0 commit comments

Comments
 (0)