Skip to content
Snippets Groups Projects
Verified Commit bb970999 authored by Nikola Milojevic's avatar Nikola Milojevic :large_blue_circle: Committed by GitLab
Browse files

Merge branch '451998-custom-audience' into 'master'

Cloud Connector: allow specifying JWT audience

See merge request gitlab-org/gitlab!147911



Merged-by: default avatarNikola Milojevic <nmilojevic@gitlab.com>
Approved-by: default avatarRoy Zwambag <rzwambag@gitlab.com>
Approved-by: default avatarNikola Milojevic <nmilojevic@gitlab.com>
Reviewed-by: default avatarRoy Zwambag <rzwambag@gitlab.com>
Co-authored-by: Matthias Käppler's avatarMatthias Kaeppler <mkaeppler@gitlab.com>
parents 54c74bd5 7409472b
No related branches found
No related tags found
1 merge request!147911Cloud Connector: allow specifying JWT audience
Pipeline #1237044587 passed
......@@ -58,10 +58,20 @@ Here we assume your backend service is called `foo` and is already reachable at
We also assume that the backend service exposes the feature using a `/new_feature_endpoint` endpoint.
This allows clients to access the feature at `https://cloud.gitlab.com/foo/new_feature_endpoint`.
Here, the parameters you pass to `access_token` have the following meaning:
- `audience`: The name of the backend service. The token is bound to this backend
using the JWT `aud` claim.
- `scopes`: The list of access scopes carried in this token. They should map to access points
in your backend, which could be HTTP endpoints or RPC calls.
```ruby
include API::Helpers::CloudConnector
token = ::CloudConnector::AccessService.new.access_token(scopes: [:new_feature_scope])
token = ::CloudConnector::AccessService.new.access_token(
audience: 'foo',
scopes: [:new_feature_scope]
)
return unauthorized! if token.nil?
Gitlab::HTTP.post(
......@@ -72,6 +82,11 @@ Gitlab::HTTP.post(
)
```
NOTE:
Any arguments you pass to `access_token` that configure the token returned only take hold for
tokens issued on GitLab.com. For self-managed GitLab instances the token is read as-is from
the database and never modified.
#### CustomersDot
This step is necessary for your feature to work for Self-Managed and GitLab Dedicated deployments.
......
......@@ -7,8 +7,15 @@ class AccessService
# rubocop:disable Gitlab/AvoidGitlabInstanceChecks -- we don't have dedicated SM/.com Cloud Connector features
# or other checks that would allow us to identify where the code is running. We rely on instance checks for now.
# Will be addressed in https://gitlab.com/gitlab-org/gitlab/-/issues/437725
def access_token(scopes:, extra_claims: {})
Gitlab.org_or_com? ? saas_token(scopes, extra_claims) : self_managed_token
#
# audience:
# The JWT aud claim used for GitLab.com tokens (ignored on self-managed)
# scopes:
# The JWT scopes claim used for GitLab.com tokens (ignored on self-managed)
# extra_claims:
# Custom JWT claims used for GitLab.com tokens (ignored on self-managed)
def access_token(audience:, scopes:, extra_claims: {})
Gitlab.org_or_com? ? saas_token(audience, scopes, extra_claims) : self_managed_token
end
# We allow free usage of cloud connected feature on Self-Managed if:
......@@ -35,8 +42,9 @@ def self_managed_token
::CloudConnector::ServiceAccessToken.active.last&.token
end
def saas_token(scopes, extra_claims)
def saas_token(audience, scopes, extra_claims)
Gitlab::CloudConnector::SelfIssuedToken.new(
audience: audience,
subject: Gitlab::CurrentSettings.uuid,
scopes: scopes,
extra_claims: extra_claims
......
......@@ -70,7 +70,7 @@ def saas_headers
documentation: { example: 'namespace/project' }
end
post do
token = ::CloudConnector::AccessService.new.access_token(scopes: [:code_suggestions])
token = Gitlab::Llm::AiGateway::Client.access_token(scopes: [:code_suggestions])
unauthorized! if token.nil?
......
......@@ -68,7 +68,7 @@ def current_namespace
strong_memoize_attr :current_namespace
def ai_gateway_token
::CloudConnector::AccessService.new.access_token(scopes: [:code_suggestions])
Gitlab::Llm::AiGateway::Client.access_token(scopes: [:code_suggestions])
end
strong_memoize_attr :ai_gateway_token
end
......
......@@ -3,7 +3,6 @@
module Gitlab
module CloudConnector
class SelfIssuedToken
JWT_AUDIENCE = 'gitlab-ai-gateway'
NOT_BEFORE_TIME = 5.seconds.to_i.freeze
EXPIRES_IN = 1.hour.to_i.freeze
......@@ -11,9 +10,9 @@ class SelfIssuedToken
attr_reader :issued_at
def initialize(subject:, scopes:, extra_claims: {})
def initialize(audience:, subject:, scopes:, extra_claims: {})
@id = SecureRandom.uuid
@audience = JWT_AUDIENCE
@audience = audience
@subject = subject
@issuer = Doorkeeper::OpenidConnect.configuration.issuer
@issued_at = Time.now.to_i
......
......@@ -16,8 +16,14 @@ class Client
DEFAULT_TYPE = 'prompt'
DEFAULT_SOURCE = 'GitLab EE'
JWT_AUDIENCE = 'gitlab-ai-gateway'
ALLOWED_PAYLOAD_PARAM_KEYS = %i[temperature max_tokens_to_sample stop_sequences].freeze
def self.access_token(scopes:)
::CloudConnector::AccessService.new.access_token(audience: JWT_AUDIENCE, scopes: scopes)
end
def initialize(user, tracking_context: {})
@user = user
@tracking_context = tracking_context
......@@ -82,23 +88,23 @@ def perform_completion_request(prompt:, options:)
end
def enabled?
access_token.present?
chat_access_token.present?
end
def request_headers
{
'X-Gitlab-Host-Name' => Gitlab.config.gitlab.host,
'X-Gitlab-Authentication-Type' => 'oidc',
'Authorization' => "Bearer #{access_token}",
'Authorization' => "Bearer #{chat_access_token}",
'Content-Type' => 'application/json',
'X-Request-ID' => Labkit::Correlation::CorrelationId.current_or_new_id
}.merge(cloud_connector_headers(user))
end
def access_token
::CloudConnector::AccessService.new.access_token(scopes: [:duo_chat])
def chat_access_token
self.class.access_token(scopes: [:duo_chat])
end
strong_memoize_attr :access_token
strong_memoize_attr :chat_access_token
def request_body(prompt:, options: {})
{
......
......@@ -5,7 +5,11 @@
RSpec.describe API::Entities::CodeSuggestionsAccessToken, feature_category: :code_suggestions do
subject { described_class.new(token).as_json }
let_it_be(:token) { Gitlab::CloudConnector::SelfIssuedToken.new(subject: 'ABC-123', scopes: [:code_suggestions]) }
let_it_be(:token) do
Gitlab::CloudConnector::SelfIssuedToken.new(
audience: 'gitlab-ai-gateway', subject: 'ABC-123', scopes: [:code_suggestions]
)
end
it 'exposes correct attributes' do
expect(subject.keys).to contain_exactly(:access_token, :expires_in, :created_at)
......
......@@ -5,7 +5,11 @@
RSpec.describe Gitlab::CloudConnector::SelfIssuedToken, feature_category: :cloud_connector do
let(:extra_claims) { {} }
subject(:token) { described_class.new(subject: 'ABC-123', scopes: [:code_suggestions], extra_claims: extra_claims) }
subject(:token) do
described_class.new(
audience: 'gitlab-ai-gateway', subject: 'ABC-123', scopes: [:code_suggestions], extra_claims: extra_claims
)
end
describe '#payload' do
subject(:payload) { token.payload }
......
......@@ -369,9 +369,7 @@ def request
end
before do
allow_next_instance_of(Gitlab::CloudConnector::SelfIssuedToken) do |instance|
allow(instance).to receive(:encoded).and_return(token)
end
allow(Gitlab::Llm::AiGateway::Client).to receive(:access_token).and_return(token)
end
context 'when user belongs to a namespace with an active code suggestions purchase' do
......
......@@ -103,22 +103,9 @@
context 'with add on' do
before_all { create(:gitlab_subscription_add_on_purchase, namespace: namespace) }
it 'calls ::CloudConnector::AccessService to obtain access token', :aggregate_failures do
expect_next_instance_of(::CloudConnector::AccessService) do |instance|
expect(instance).to receive(:access_token).with(scopes: [:code_suggestions])
.and_return(ai_gateway_token)
end
post_api
expect(response).to have_gitlab_http_status(:ok)
end
context 'when cloud connector access token is missing' do
before do
allow_next_instance_of(::CloudConnector::AccessService) do |instance|
allow(instance).to receive(:access_token).and_return(nil)
end
allow(::Gitlab::Llm::AiGateway::Client).to receive(:access_token).and_return(nil)
end
it 'returns UNAUTHORIZED status' do
......@@ -130,9 +117,7 @@
context 'when cloud connector access token is valid' do
before do
allow_next_instance_of(::CloudConnector::AccessService) do |instance|
allow(instance).to receive(:access_token).and_return(ai_gateway_token)
end
allow(::Gitlab::Llm::AiGateway::Client).to receive(:access_token).and_return(ai_gateway_token)
end
context 'when instance has uuid available' do
......@@ -184,9 +169,7 @@
end
before do
allow_next_instance_of(::CloudConnector::AccessService) do |instance|
allow(instance).to receive(:access_token).and_return(ai_gateway_token)
end
allow(::Gitlab::Llm::AiGateway::Client).to receive(:access_token).and_return(ai_gateway_token)
end
context 'with purchase_code_suggestions feature disabled' do
......
......@@ -15,8 +15,9 @@
let_it_be(:cloud_connector_access) { create(:cloud_connector_access, data: data) }
describe '#access_token' do
subject(:access_token) { described_class.new.access_token(scopes: scopes) }
subject(:access_token) { described_class.new.access_token(audience: audience, scopes: scopes) }
let(:audience) { 'gitlab-ai-gateway' }
let(:scopes) { [:code_suggestions, :duo_chat] }
context 'when self-managed' do
......@@ -37,6 +38,7 @@
it 'returns the constructed token' do
expect(Gitlab::CloudConnector::SelfIssuedToken).to receive(:new).with(
audience: audience,
subject: gitlab_instance_id,
scopes: scopes,
extra_claims: {}
......@@ -48,10 +50,13 @@
context 'when passing additional claims' do
let(:extra_claims) { { 'custom_claim' => 'custom_value' } }
subject(:access_token) { described_class.new.access_token(scopes: scopes, extra_claims: extra_claims) }
subject(:access_token) do
described_class.new.access_token(audience: audience, scopes: scopes, extra_claims: extra_claims)
end
it 'includes extra_claims element in token payload' do
expect(Gitlab::CloudConnector::SelfIssuedToken).to receive(:new).with(
audience: audience,
subject: gitlab_instance_id,
scopes: scopes,
extra_claims: extra_claims
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment