From bad4766f1f1087880a9b4e55967364c60132634b Mon Sep 17 00:00:00 2001 From: Max Woolf Date: Fri, 28 Jul 2023 11:26:31 +0100 Subject: [PATCH 1/3] Self-generate JWT for SaaS completions For SaaS installations the /completions API endpoint now self-generates a JWT for authentication to the model gateway. The existing /tokens endpoint remains unchanged to preserve SM functionality and backwards compatibility with existing plugins. EE: true Changelog: changed --- ee/lib/api/code_suggestions.rb | 8 +++++++- ee/spec/requests/api/code_suggestions_spec.rb | 3 +++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb index b2f8be7baac52c..c9766943f0ed88 100644 --- a/ee/lib/api/code_suggestions.rb +++ b/ee/lib/api/code_suggestions.rb @@ -35,9 +35,15 @@ def active_code_suggestions_purchase?(project_id) def model_gateway_headers(headers) telemetry_headers = headers.select { |k| /\Ax-gitlab-cs-/i.match?(k) } + token = if Gitlab.org_or_com? + Gitlab::CodeSuggestions::AccessToken.new(current_user, gitlab_realm: gitlab_realm).encoded + else + headers['X-Gitlab-Oidc-Token'] + end + { 'X-Gitlab-Authentication-Type' => 'oidc', - 'Authorization' => "Bearer #{headers['X-Gitlab-Oidc-Token']}", + 'Authorization' => "Bearer #{token}", 'Content-Type' => 'application/json', 'User-Agent' => headers["User-Agent"] # Forward the User-Agent on to the model gateway }.merge(telemetry_headers).transform_values { |v| Array(v) } diff --git a/ee/spec/requests/api/code_suggestions_spec.rb b/ee/spec/requests/api/code_suggestions_spec.rb index f7b2c85e29c857..62d601a3396f9e 100644 --- a/ee/spec/requests/api/code_suggestions_spec.rb +++ b/ee/spec/requests/api/code_suggestions_spec.rb @@ -213,6 +213,9 @@ before do stub_env('CODE_SUGGESTIONS_BASE_URL', nil) + allow_next_instance_of(Gitlab::CodeSuggestions::AccessToken) do |token| + allow(token).to receive(:encoded).and_return('JWTTOKEN') + end end it 'delegates downstream service call to Workhorse with auth token from the DB' do -- GitLab From 73d49c3f5c9a0512c84f7248bdcc971c5ed14fc8 Mon Sep 17 00:00:00 2001 From: Max Woolf Date: Fri, 28 Jul 2023 15:12:32 +0100 Subject: [PATCH 2/3] Backend reviewer suggestions --- ee/lib/api/code_suggestions.rb | 17 ++++++------- ee/spec/requests/api/code_suggestions_spec.rb | 25 +++++++++++++++++-- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb index c9766943f0ed88..8eb7daa242d0c5 100644 --- a/ee/lib/api/code_suggestions.rb +++ b/ee/lib/api/code_suggestions.rb @@ -35,15 +35,9 @@ def active_code_suggestions_purchase?(project_id) def model_gateway_headers(headers) telemetry_headers = headers.select { |k| /\Ax-gitlab-cs-/i.match?(k) } - token = if Gitlab.org_or_com? - Gitlab::CodeSuggestions::AccessToken.new(current_user, gitlab_realm: gitlab_realm).encoded - else - headers['X-Gitlab-Oidc-Token'] - end - { 'X-Gitlab-Authentication-Type' => 'oidc', - 'Authorization' => "Bearer #{token}", + 'Authorization' => "Bearer #{headers['X-Gitlab-Oidc-Token']}", 'Content-Type' => 'application/json', 'User-Agent' => headers["User-Agent"] # Forward the User-Agent on to the model gateway }.merge(telemetry_headers).transform_values { |v| Array(v) } @@ -105,10 +99,15 @@ def completions_endpoint resources :completions do post do if Gitlab.org_or_com? - not_found! unless ::Feature.enabled?(:code_suggestions_completion_api, current_user) + forbidden! unless ::Feature.enabled?(:code_suggestions_completion_api, current_user) not_found! unless active_code_suggestions_purchase?(params['project_id']) + + headers['X-Gitlab-Oidc-Token'] = Gitlab::CodeSuggestions::AccessToken.new( + current_user, + gitlab_realm: gitlab_realm + ).encoded else - not_found! unless ::Feature.enabled?(:self_managed_code_suggestions_completion_api) + forbidden! unless ::Feature.enabled?(:self_managed_code_suggestions_completion_api) code_suggestions_token = ::Ai::ServiceAccessToken.code_suggestions.active.last unauthorized! if code_suggestions_token.nil? diff --git a/ee/spec/requests/api/code_suggestions_spec.rb b/ee/spec/requests/api/code_suggestions_spec.rb index 62d601a3396f9e..5cef8d0555747a 100644 --- a/ee/spec/requests/api/code_suggestions_spec.rb +++ b/ee/spec/requests/api/code_suggestions_spec.rb @@ -54,6 +54,15 @@ end end + shared_examples 'a forbidden response' do + include_examples 'a response', 'unauthorized' do + let(:result) { :forbidden } + let(:body) do + { "message" => "403 Forbidden" } + end + end + end + shared_examples 'a not found response' do include_examples 'a response', 'not found' do let(:result) { :not_found } @@ -218,6 +227,18 @@ end end + it 'finds the token using the correct method' do + if Gitlab.org_or_com? + expect_next_instance_of(Gitlab::CodeSuggestions::AccessToken) do |token| + expect(token).to receive(:encoded).once + end + else + expect(Ai::ServiceAccessToken).to receive(:code_suggestions).once + end + + post_api + end + it 'delegates downstream service call to Workhorse with auth token from the DB' do post_api @@ -330,7 +351,7 @@ stub_feature_flags(code_suggestions_completion_api: false) end - include_examples 'a not found response' + include_examples 'a forbidden response' end context 'when purchase_code_suggestions feature flag is disabled' do @@ -381,7 +402,7 @@ stub_feature_flags(self_managed_code_suggestions_completion_api: false) end - include_examples 'a not found response' + include_examples 'a forbidden response' end end end -- GitLab From ca0dc83b1ae2f74efc10d341d2b4f550a155f9f5 Mon Sep 17 00:00:00 2001 From: Max Woolf Date: Tue, 1 Aug 2023 10:52:17 +0100 Subject: [PATCH 3/3] Backend maintainer suggestions --- ee/lib/api/code_suggestions.rb | 10 ++--- ee/spec/requests/api/code_suggestions_spec.rb | 38 +++++++------------ 2 files changed, 18 insertions(+), 30 deletions(-) diff --git a/ee/lib/api/code_suggestions.rb b/ee/lib/api/code_suggestions.rb index 8eb7daa242d0c5..4b51fb7d9ecb54 100644 --- a/ee/lib/api/code_suggestions.rb +++ b/ee/lib/api/code_suggestions.rb @@ -32,12 +32,12 @@ def active_code_suggestions_purchase?(project_id) end end - def model_gateway_headers(headers) + def model_gateway_headers(headers, gateway_token) telemetry_headers = headers.select { |k| /\Ax-gitlab-cs-/i.match?(k) } { 'X-Gitlab-Authentication-Type' => 'oidc', - 'Authorization' => "Bearer #{headers['X-Gitlab-Oidc-Token']}", + 'Authorization' => "Bearer #{gateway_token}", 'Content-Type' => 'application/json', 'User-Agent' => headers["User-Agent"] # Forward the User-Agent on to the model gateway }.merge(telemetry_headers).transform_values { |v| Array(v) } @@ -102,7 +102,7 @@ def completions_endpoint forbidden! unless ::Feature.enabled?(:code_suggestions_completion_api, current_user) not_found! unless active_code_suggestions_purchase?(params['project_id']) - headers['X-Gitlab-Oidc-Token'] = Gitlab::CodeSuggestions::AccessToken.new( + token = Gitlab::CodeSuggestions::AccessToken.new( current_user, gitlab_realm: gitlab_realm ).encoded @@ -112,14 +112,14 @@ def completions_endpoint code_suggestions_token = ::Ai::ServiceAccessToken.code_suggestions.active.last unauthorized! if code_suggestions_token.nil? - headers['X-Gitlab-Oidc-Token'] = code_suggestions_token.token + token = code_suggestions_token.token end workhorse_headers = Gitlab::Workhorse.send_url( completions_endpoint, body: params.except(:private_token).to_json, - headers: model_gateway_headers(headers), + headers: model_gateway_headers(headers, token), method: "POST" ) diff --git a/ee/spec/requests/api/code_suggestions_spec.rb b/ee/spec/requests/api/code_suggestions_spec.rb index 5cef8d0555747a..36dc9a78b760e0 100644 --- a/ee/spec/requests/api/code_suggestions_spec.rb +++ b/ee/spec/requests/api/code_suggestions_spec.rb @@ -178,7 +178,6 @@ end describe 'POST /code_suggestions/completions' do - let_it_be(:token) { 'JWTTOKEN' } let(:access_code_suggestions) { true } let(:body) do @@ -199,6 +198,7 @@ end before do + allow(Gitlab).to receive(:org_or_com?).and_return(is_saas) allow(Ability).to receive(:allowed?).and_call_original allow(Ability).to receive(:allowed?).with(current_user, :access_code_suggestions, :global) .and_return(access_code_suggestions) @@ -222,24 +222,9 @@ before do stub_env('CODE_SUGGESTIONS_BASE_URL', nil) - allow_next_instance_of(Gitlab::CodeSuggestions::AccessToken) do |token| - allow(token).to receive(:encoded).and_return('JWTTOKEN') - end end - it 'finds the token using the correct method' do - if Gitlab.org_or_com? - expect_next_instance_of(Gitlab::CodeSuggestions::AccessToken) do |token| - expect(token).to receive(:encoded).once - end - else - expect(Ai::ServiceAccessToken).to receive(:code_suggestions).once - end - - post_api - end - - it 'delegates downstream service call to Workhorse with auth token from the DB' do + it 'delegates downstream service call to Workhorse with correct auth token' do post_api expect(response.status).to be(200) @@ -312,9 +297,8 @@ end context 'when the instance is Gitlab.org_or_com' do - before do - allow(Gitlab).to receive(:org_or_com?).and_return(true) - end + let(:is_saas) { true } + let_it_be(:token) { 'generated-jwt' } let(:headers) do { @@ -325,6 +309,12 @@ } end + before do + allow_next_instance_of(Gitlab::CodeSuggestions::AccessToken) do |instance| + allow(instance).to receive(:encoded).and_return(token) + end + end + context 'when project does not have active code suggestions purchase' do let(:current_user) { create(:user) } @@ -366,9 +356,9 @@ end context 'when the instance is Gitlab self-managed' do - before do - allow(Gitlab).to receive(:org_or_com?).and_return(false) - end + let(:is_saas) { false } + let_it_be(:token) { 'stored-token' } + let_it_be(:service_access_token) { create(:service_access_token, :code_suggestions, :active, token: token) } let(:headers) do { @@ -378,8 +368,6 @@ } end - let_it_be(:service_access_token) { create(:service_access_token, :code_suggestions, :active, token: token) } - it_behaves_like 'code completions endpoint' context 'when there is no active code suggestions token' do -- GitLab