From 5655d6c9d204206537c781876748912f9a6c0b71 Mon Sep 17 00:00:00 2001
From: Eduardo Bonet <ebonet@gitlab.com>
Date: Thu, 8 Jun 2023 10:41:09 +0000
Subject: [PATCH] Adds authorize endpoint for pacakges/ml_models
Adds first endpoint for model registry. Authorization follows the exact
same rules as generic packages, and the tests are the same.
Changelog: added
---
config/open_api.yml | 2 +
lib/api/api.rb | 1 +
lib/api/ml_model_packages.rb | 111 ++++++++++
lib/gitlab/regex.rb | 12 ++
spec/requests/api/ml_model_packages_spec.rb | 200 ++++++++++++++++++
.../api/ml_model_packages_shared_examples.rb | 108 ++++++++++
workhorse/internal/upstream/routes.go | 3 +
workhorse/upload_test.go | 1 +
8 files changed, 438 insertions(+)
create mode 100644 lib/api/ml_model_packages.rb
create mode 100644 spec/requests/api/ml_model_packages_spec.rb
create mode 100644 spec/support/shared_examples/requests/api/ml_model_packages_shared_examples.rb
diff --git a/config/open_api.yml b/config/open_api.yml
index cbf70c24ce1ce..257db9dd69258 100644
--- a/config/open_api.yml
+++ b/config/open_api.yml
@@ -95,6 +95,8 @@ metadata:
description: Operations related to metadata of the GitLab instance
- name: metrics_user_starred_dashboards
description: Operations related to User-starred metrics dashboards
+ - name: ml_model_registry
+ description: Operations related to Model registry
- name: npm_packages
description: Operations related to NPM packages
- name: nuget_packages
diff --git a/lib/api/api.rb b/lib/api/api.rb
index a7acd44e72a2b..5f7faa7eb7d37 100644
--- a/lib/api/api.rb
+++ b/lib/api/api.rb
@@ -255,6 +255,7 @@ class API < ::API::Base
mount ::API::Metadata
mount ::API::Metrics::Dashboard::Annotations
mount ::API::Metrics::UserStarredDashboards
+ mount ::API::MlModelPackages
mount ::API::Namespaces
mount ::API::NpmGroupPackages
mount ::API::NpmInstancePackages
diff --git a/lib/api/ml_model_packages.rb b/lib/api/ml_model_packages.rb
new file mode 100644
index 0000000000000..fec72b03ffd41
--- /dev/null
+++ b/lib/api/ml_model_packages.rb
@@ -0,0 +1,111 @@
+# frozen_string_literal: true
+
+module API
+ class MlModelPackages < ::API::Base
+ include APIGuard
+ include ::API::Helpers::Authentication
+
+ ML_MODEL_PACKAGES_REQUIREMENTS = {
+ package_name: API::NO_SLASH_URL_PART_REGEX,
+ file_name: API::NO_SLASH_URL_PART_REGEX
+ }.freeze
+
+ ALLOWED_STATUSES = %w[default hidden].freeze
+
+ feature_category :mlops
+ urgency :low
+
+ after_validation do
+ require_packages_enabled!
+ authenticate_non_get!
+
+ not_found! unless can?(current_user, :read_model_registry, user_project)
+ end
+
+ authenticate_with do |accept|
+ accept.token_types(:personal_access_token, :deploy_token, :job_token)
+ .sent_through(:http_token)
+ end
+
+ helpers do
+ include ::API::Helpers::PackagesHelpers
+ include ::API::Helpers::Packages::BasicAuthHelpers
+
+ def project
+ authorized_user_project
+ end
+
+ def max_file_size_exceeded?
+ project.actual_limits.exceeded?(:ml_model_max_file_size, params[:file].size)
+ end
+ end
+
+ params do
+ requires :id, types: [String, Integer], desc: 'The ID or URL-encoded path of the project'
+ end
+
+ resource :projects, requirements: API::NAMESPACE_OR_PROJECT_REQUIREMENTS do
+ namespace ':id/packages/ml_models' do
+ params do
+ requires :package_name, type: String, desc: 'Package name', regexp: Gitlab::Regex.ml_model_name_regex,
+ file_path: true
+ requires :package_version, type: String, desc: 'Package version',
+ regexp: Gitlab::Regex.ml_model_version_regex
+ requires :file_name, type: String, desc: 'Package file name',
+ regexp: Gitlab::Regex.ml_model_file_name_regex, file_path: true
+ optional :status, type: String, values: ALLOWED_STATUSES, desc: 'Package status'
+ end
+ namespace ':package_name/*package_version/:file_name', requirements: ML_MODEL_PACKAGES_REQUIREMENTS do
+ desc 'Workhorse authorize model package file' do
+ detail 'Introduced in GitLab 16.1'
+ success code: 200
+ failure [
+ { code: 401, message: 'Unauthorized' },
+ { code: 403, message: 'Forbidden' },
+ { code: 404, message: 'Not Found' }
+ ]
+ tags %w[ml_model_registry]
+ end
+ put 'authorize' do
+ authorize_workhorse!(subject: project, maximum_size: project.actual_limits.ml_model_max_file_size)
+ end
+
+ desc 'Workhorse upload model package file' do
+ detail 'Introduced in GitLab 16.1'
+ success code: 201
+ failure [
+ { code: 401, message: 'Unauthorized' },
+ { code: 403, message: 'Forbidden' },
+ { code: 404, message: 'Not Found' }
+ ]
+ tags %w[ml_model_registry]
+ end
+ params do
+ requires :file,
+ type: ::API::Validations::Types::WorkhorseFile,
+ desc: 'The package file to be published (generated by Multipart middleware)',
+ documentation: { type: 'file' }
+ end
+ put do
+ authorize_upload!(project)
+
+ bad_request!('File is too large') if max_file_size_exceeded?
+
+ create_package_file_params = declared(params).merge(build: current_authenticated_job)
+ package_file = ::Packages::MlModel::CreatePackageFileService
+ .new(project, current_user, create_package_file_params)
+ .execute
+
+ bad_request!('Package creation failed') unless package_file
+
+ created!
+ rescue ObjectStorage::RemoteStoreError => e
+ Gitlab::ErrorTracking.track_exception(e, extra: { file_name: params[:file_name], project_id: project.id })
+
+ forbidden!
+ end
+ end
+ end
+ end
+ end
+end
diff --git a/lib/gitlab/regex.rb b/lib/gitlab/regex.rb
index 79f12ee13f78d..26ca9d2547c86 100644
--- a/lib/gitlab/regex.rb
+++ b/lib/gitlab/regex.rb
@@ -623,6 +623,18 @@ def sep_by_1(separator, part)
def x509_subject_key_identifier_regex
@x509_subject_key_identifier_regex ||= /\A(?:\h{2}:)*\h{2}\z/.freeze
end
+
+ def ml_model_version_regex
+ maven_version_regex
+ end
+
+ def ml_model_name_regex
+ package_name_regex
+ end
+
+ def ml_model_file_name_regex
+ maven_file_name_regex
+ end
end
end
diff --git a/spec/requests/api/ml_model_packages_spec.rb b/spec/requests/api/ml_model_packages_spec.rb
new file mode 100644
index 0000000000000..9c19f522e46cc
--- /dev/null
+++ b/spec/requests/api/ml_model_packages_spec.rb
@@ -0,0 +1,200 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe ::API::MlModelPackages, feature_category: :mlops do
+ include HttpBasicAuthHelpers
+ include PackagesManagerApiSpecHelpers
+ include WorkhorseHelpers
+ using RSpec::Parameterized::TableSyntax
+
+ include_context 'workhorse headers'
+
+ let_it_be(:project, reload: true) { create(:project) }
+ let_it_be(:personal_access_token) { create(:personal_access_token) }
+ let_it_be(:job) { create(:ci_build, :running, user: personal_access_token.user, project: project) }
+ let_it_be(:deploy_token) { create(:deploy_token, read_package_registry: true, write_package_registry: true) }
+ let_it_be(:project_deploy_token) { create(:project_deploy_token, deploy_token: deploy_token, project: project) }
+ let_it_be(:another_project, reload: true) { create(:project) }
+
+ let_it_be(:tokens) do
+ {
+ personal_access_token: personal_access_token.token,
+ deploy_token: deploy_token.token,
+ job_token: job.token
+ }
+ end
+
+ let(:user) { personal_access_token.user }
+ let(:user_role) { :developer }
+ let(:member) { true }
+ let(:ci_build) { create(:ci_build, :running, user: user, project: project) }
+ let(:project_to_enable_ff) { project }
+ let(:headers) { {} }
+
+ shared_context 'ml model authorize permissions table' do # rubocop:disable RSpec/ContextWording
+ # rubocop:disable Metrics/AbcSize
+ # :visibility, :user_role, :member, :token_type, :valid_token, :expected_status
+ def authorize_permissions_table
+ :public | :developer | true | :personal_access_token | true | :success
+ :public | :guest | true | :personal_access_token | true | :forbidden
+ :public | :developer | true | :personal_access_token | false | :unauthorized
+ :public | :guest | true | :personal_access_token | false | :unauthorized
+ :public | :developer | false | :personal_access_token | true | :forbidden
+ :public | :guest | false | :personal_access_token | true | :forbidden
+ :public | :developer | false | :personal_access_token | false | :unauthorized
+ :public | :guest | false | :personal_access_token | false | :unauthorized
+ :public | :anonymous | false | :personal_access_token | true | :unauthorized
+ :private | :developer | true | :personal_access_token | true | :success
+ :private | :guest | true | :personal_access_token | true | :forbidden
+ :private | :developer | true | :personal_access_token | false | :unauthorized
+ :private | :guest | true | :personal_access_token | false | :unauthorized
+ :private | :developer | false | :personal_access_token | true | :not_found
+ :private | :guest | false | :personal_access_token | true | :not_found
+ :private | :developer | false | :personal_access_token | false | :unauthorized
+ :private | :guest | false | :personal_access_token | false | :unauthorized
+ :private | :anonymous | false | :personal_access_token | true | :unauthorized
+ :public | :developer | true | :job_token | true | :success
+ :public | :guest | true | :job_token | true | :forbidden
+ :public | :developer | true | :job_token | false | :unauthorized
+ :public | :guest | true | :job_token | false | :unauthorized
+ :public | :developer | false | :job_token | true | :forbidden
+ :public | :guest | false | :job_token | true | :forbidden
+ :public | :developer | false | :job_token | false | :unauthorized
+ :public | :guest | false | :job_token | false | :unauthorized
+ :private | :developer | true | :job_token | true | :success
+ :private | :guest | true | :job_token | true | :forbidden
+ :private | :developer | true | :job_token | false | :unauthorized
+ :private | :guest | true | :job_token | false | :unauthorized
+ :private | :developer | false | :job_token | true | :not_found
+ :private | :guest | false | :job_token | true | :not_found
+ :private | :developer | false | :job_token | false | :unauthorized
+ :private | :guest | false | :job_token | false | :unauthorized
+ :public | :developer | true | :deploy_token | true | :success
+ :public | :developer | true | :deploy_token | false | :unauthorized
+ :private | :developer | true | :deploy_token | true | :success
+ :private | :developer | true | :deploy_token | false | :unauthorized
+ end
+ # rubocop:enable Metrics/AbcSize
+ end
+
+ before do
+ project.send("add_#{user_role}", user) if member && user_role != :anonymous
+ end
+
+ subject(:api_response) do
+ request
+ response
+ end
+
+ describe 'PUT /api/v4/projects/:id/packages/ml_models/:package_name/:package_version/:file_name/authorize' do
+ include_context 'ml model authorize permissions table'
+
+ let(:token) { tokens[:personal_access_token] }
+ let(:user_headers) { { 'HTTP_AUTHORIZATION' => token } }
+ let(:headers) { user_headers.merge(workhorse_headers) }
+ let(:request) { authorize_upload_file(headers) }
+
+ describe 'user access' do
+ where(:visibility, :user_role, :member, :token_type, :valid_token, :expected_status) do
+ authorize_permissions_table
+ end
+
+ with_them do
+ let(:token) { valid_token ? tokens[token_type] : 'invalid-token123' }
+ let(:user_headers) { user_role == :anonymous ? {} : { 'HTTP_AUTHORIZATION' => token } }
+
+ before do
+ project.update_column(:visibility_level, Gitlab::VisibilityLevel.level_value(visibility.to_s))
+ end
+
+ it { is_expected.to have_gitlab_http_status(expected_status) }
+ end
+
+ it_behaves_like 'Endpoint not found if read_model_registry not available'
+ end
+
+ describe 'application security' do
+ where(:param_name, :param_value) do
+ :package_name | 'my-package/../'
+ :package_name | 'my-package%2f%2e%2e%2f'
+ :file_name | '../.ssh%2fauthorized_keys'
+ :file_name | '%2e%2e%2f.ssh%2fauthorized_keys'
+ end
+
+ with_them do
+ let(:request) { authorize_upload_file(headers, param_name => param_value) }
+
+ it 'rejects malicious request' do
+ is_expected.to have_gitlab_http_status(:bad_request)
+ end
+ end
+ end
+ end
+
+ describe 'PUT /api/v4/projects/:id/packages/ml_models/:package_name/:package_version/:file_name' do
+ include_context 'ml model authorize permissions table'
+
+ let_it_be(:file_name) { 'model.md5' }
+
+ let(:token) { tokens[:personal_access_token] }
+ let(:user_headers) { { 'HTTP_AUTHORIZATION' => token } }
+ let(:headers) { user_headers.merge(workhorse_headers) }
+ let(:params) { { file: temp_file(file_name) } }
+ let(:file_key) { :file }
+ let(:send_rewritten_field) { true }
+
+ let(:request) do
+ upload_file(headers)
+ end
+
+ describe 'success' do
+ it 'creates a new package' do
+ expect { api_response }.to change { Packages::PackageFile.count }.by(1)
+ expect(api_response).to have_gitlab_http_status(:created)
+ end
+ end
+
+ describe 'user access' do
+ where(:visibility, :user_role, :member, :token_type, :valid_token, :expected_status) do
+ authorize_permissions_table
+ end
+
+ with_them do
+ let(:token) { valid_token ? tokens[token_type] : 'invalid-token123' }
+ let(:user_headers) { user_role == :anonymous ? {} : { 'HTTP_AUTHORIZATION' => token } }
+
+ before do
+ project.update_column(:visibility_level, Gitlab::VisibilityLevel.level_value(visibility.to_s))
+ end
+
+ if params[:expected_status] == :success
+ it_behaves_like 'process ml model package upload'
+ else
+ it { is_expected.to have_gitlab_http_status(expected_status) }
+ end
+ end
+
+ it_behaves_like 'Endpoint not found if read_model_registry not available'
+ end
+ end
+
+ def authorize_upload_file(request_headers, package_name: 'mypackage', file_name: 'myfile.tar.gz')
+ url = "/projects/#{project.id}/packages/ml_models/#{package_name}/0.0.1/#{file_name}/authorize"
+
+ put api(url), headers: request_headers
+ end
+
+ def upload_file(request_headers, package_name: 'mypackage')
+ url = "/projects/#{project.id}/packages/ml_models/#{package_name}/0.0.1/#{file_name}"
+
+ workhorse_finalize(
+ api(url),
+ method: :put,
+ file_key: file_key,
+ params: params,
+ headers: request_headers,
+ send_rewritten_field: send_rewritten_field
+ )
+ end
+end
diff --git a/spec/support/shared_examples/requests/api/ml_model_packages_shared_examples.rb b/spec/support/shared_examples/requests/api/ml_model_packages_shared_examples.rb
new file mode 100644
index 0000000000000..81ff004779af2
--- /dev/null
+++ b/spec/support/shared_examples/requests/api/ml_model_packages_shared_examples.rb
@@ -0,0 +1,108 @@
+# frozen_string_literal: true
+
+RSpec.shared_examples 'Endpoint not found if read_model_registry not available' do
+ context 'when read_model_registry disabled for current project' do
+ before do
+ allow(Ability).to receive(:allowed?).and_call_original
+ allow(Ability).to receive(:allowed?)
+ .with(user, :read_model_registry, project)
+ .and_return(false)
+ end
+
+ it "is not found" do
+ is_expected.to have_gitlab_http_status(:not_found)
+ end
+ end
+end
+
+RSpec.shared_examples 'creates model experiments package files' do
+ it 'creates package files', :aggregate_failures do
+ expect { api_response }
+ .to change { project.packages.count }.by(1)
+ .and change { Packages::PackageFile.count }.by(1)
+ expect(api_response).to have_gitlab_http_status(:created)
+
+ package_file = project.packages.last.package_files.reload.last
+ expect(package_file.file_name).to eq(file_name)
+ end
+
+ it 'returns bad request if package creation fails' do
+ allow_next_instance_of(::Packages::MlModel::CreatePackageFileService) do |instance|
+ expect(instance).to receive(:execute).and_return(nil)
+ end
+
+ expect(api_response).to have_gitlab_http_status(:bad_request)
+ end
+
+ context 'when file is too large' do
+ it 'is bad request', :aggregate_failures do
+ allow_next_instance_of(UploadedFile) do |uploaded_file|
+ allow(uploaded_file).to receive(:size).and_return(project.actual_limits.ml_model_max_file_size + 1)
+ end
+
+ expect(api_response).to have_gitlab_http_status(:bad_request)
+ end
+ end
+end
+
+RSpec.shared_examples 'process ml model package upload' do
+ context 'with object storage disabled' do
+ before do
+ stub_package_file_object_storage(enabled: false)
+ end
+
+ context 'without a file from workhorse' do
+ let(:send_rewritten_field) { false }
+
+ it_behaves_like 'returning response status', :bad_request
+ end
+
+ context 'with correct params' do
+ it_behaves_like 'package workhorse uploads'
+ it_behaves_like 'creates model experiments package files'
+ # To be reactivated with https://gitlab.com/gitlab-org/gitlab/-/issues/414270
+ # it_behaves_like 'a package tracking event', '::API::MlModelPackages', 'push_package'
+ end
+ end
+
+ context 'with object storage enabled' do
+ let(:tmp_object) do
+ fog_connection.directories.new(key: 'packages').files.create( # rubocop:disable Rails/SaveBang
+ key: "tmp/uploads/#{file_name}",
+ body: 'content'
+ )
+ end
+
+ let(:fog_file) { fog_to_uploaded_file(tmp_object) }
+ let(:params) { { file: fog_file, 'file.remote_id' => file_name } }
+
+ context 'and direct upload enabled' do
+ let(:fog_connection) do
+ stub_package_file_object_storage(direct_upload: true)
+ end
+
+ it_behaves_like 'creates model experiments package files'
+
+ ['123123', '../../123123'].each do |remote_id|
+ context "with invalid remote_id: #{remote_id}" do
+ let(:params) do
+ {
+ file: fog_file,
+ 'file.remote_id' => remote_id
+ }
+ end
+
+ it { is_expected.to have_gitlab_http_status(:forbidden) }
+ end
+ end
+ end
+
+ context 'and direct upload disabled' do
+ let(:fog_connection) do
+ stub_package_file_object_storage(direct_upload: false)
+ end
+
+ it_behaves_like 'creates model experiments package files'
+ end
+ end
+end
diff --git a/workhorse/internal/upstream/routes.go b/workhorse/internal/upstream/routes.go
index c5af9dc3f006f..dd1725a723e51 100644
--- a/workhorse/internal/upstream/routes.go
+++ b/workhorse/internal/upstream/routes.go
@@ -279,6 +279,9 @@ func configureRoutes(u *upstream) {
// Generic Packages Repository
u.route("PUT", apiProjectPattern+`/packages/generic/`, requestBodyUploader),
+ // Ml Model Packages Repository
+ u.route("PUT", apiProjectPattern+`/packages/ml_models/`, requestBodyUploader),
+
// NuGet Artifact Repository
u.route("PUT", apiProjectPattern+`/packages/nuget/`, mimeMultipartUploader),
diff --git a/workhorse/upload_test.go b/workhorse/upload_test.go
index c05af7317973b..8effe29197930 100644
--- a/workhorse/upload_test.go
+++ b/workhorse/upload_test.go
@@ -580,6 +580,7 @@ func TestPackageFilesUpload(t *testing.T) {
{"PUT", "/api/v4/projects/group%2Fproject/packages/conan/v1/files"},
{"PUT", "/api/v4/projects/group%2Fproject/packages/maven/v1/files"},
{"PUT", "/api/v4/projects/group%2Fproject/packages/generic/mypackage/0.0.1/myfile.tar.gz"},
+ {"PUT", "/api/v4/projects/group%2Fproject/packages/ml_models/mymodel/0.0.1/myfile.tar.gz"},
{"PUT", "/api/v4/projects/group%2Fproject/packages/debian/libsample0_1.2.3~alpha2-1_amd64.deb"},
{"POST", "/api/v4/projects/group%2Fproject/packages/rubygems/api/v1/gems/sample.gem"},
{"POST", "/api/v4/projects/group%2Fproject/packages/rpm/sample-4.23.fc21.x86_64.rpm"},
--
GitLab