Check repo size for snippets

Add size checker access method to snippet.
Move size check logic to FOSS level.

app/models/snippet.rb

@@ -262,6 +262,15 @@ class Snippet < ApplicationRecord
     @repository ||= Repository.new(full_path, self, shard: repository_storage, disk_path: disk_path, repo_type: Gitlab::GlRepository::SNIPPET)
   end
 
+  def repository_size_checker
+    strong_memoize(:repository_size_checker) do
+      ::Gitlab::RepositorySizeChecker.new(
+        current_size_proc: -> { repository._uncached_size.megabytes },
+        limit: Gitlab::CurrentSettings.snippet_size_limit
+      )
+    end
+  end
+
   def storage
     @storage ||= Storage::Hashed.new(self, prefix: Storage::Hashed::SNIPPET_REPOSITORY_PATH_PREFIX)
   end

ee/lib/ee/gitlab/git_access.rb

@@ -96,12 +96,6 @@ module EE
         actor == :geo
       end
 
-      def check_size_before_push!
-        if check_size_limit? && size_checker.above_size_limit?
-          raise ::Gitlab::GitAccess::ForbiddenError, size_checker.error_message.push_error
-        end
-      end
-
       def check_if_license_blocks_changes!
         if ::License.block_changes?
           message = ::LicenseHelper.license_message(signed_in: true, is_admin: (user && user.admin?))
@@ -109,65 +103,12 @@ module EE
         end
       end
 
-      def check_push_size!
-        return unless check_size_limit?
-
-        # If there are worktrees with a HEAD pointing to a non-existent object,
-        # calls to `git rev-list --all` will fail in git 2.15+. This should also
-        # clear stale lock files.
-        project.repository.clean_stale_repository_files
-
-        # Use #check_repository_disk_size to get correct push size whenever a lot of changes
-        # gets pushed at the same time containing the same blobs. This is only
-        # doable if GIT_OBJECT_DIRECTORY_RELATIVE env var is set and happens
-        # when git push comes from CLI (not via UI and API).
-        #
-        # Fallback to determining push size using the changes_list so we can still
-        # determine the push size if env var isn't set (e.g. changes are made
-        # via UI and API).
-        if check_quarantine_size?
-          check_repository_disk_size
-        else
-          check_changes_size
-        end
-      end
-
-      def check_quarantine_size?
-        git_env = ::Gitlab::Git::HookEnv.all(repository.gl_repository)
-
-        git_env['GIT_OBJECT_DIRECTORY_RELATIVE'].present?
-      end
-
-      def check_repository_disk_size
-        check_size_against_limit(project.repository.object_directory_size)
-      end
-
-      def check_changes_size
-        changes_size = 0
-
-        changes_list.each do |change|
-          changes_size += repository.new_blobs(change[:newrev]).sum(&:size) # rubocop: disable CodeReuse/ActiveRecord
-
-          check_size_against_limit(changes_size)
-        end
-      end
-
-      def check_size_against_limit(size)
-        if size_checker.changes_will_exceed_size_limit?(size)
-          raise ::Gitlab::GitAccess::ForbiddenError, size_checker.error_message.new_changes_error
-        end
-      end
-
+      override :check_size_limit?
       def check_size_limit?
         strong_memoize(:check_size_limit) do
-          size_checker.enabled? &&
-            changes_list.any? { |change| !::Gitlab::Git.blank_ref?(change[:newrev]) }
+          size_checker.enabled? && super
         end
       end
-
-      def size_checker
-        project.repository_size_checker
-      end
     end
   end
 end

lib/gitlab/git_access.rb

@@ -45,6 +45,8 @@ module Gitlab
 
     attr_reader :actor, :project, :protocol, :authentication_abilities, :namespace_path, :repository_path, :redirected_path, :auth_result_type, :changes, :logger
 
+    alias_method :container, :project
+
     def initialize(actor, project, protocol, authentication_abilities:, namespace_path: nil, repository_path: nil, redirected_path: nil, auth_result_type: nil)
       @actor    = actor
       @project  = project
@@ -429,7 +431,72 @@ module Gitlab
     end
 
     def repository
-      project.repository
+      container&.repository
+    end
+
+    def check_size_before_push!
+      if check_size_limit? && size_checker.above_size_limit?
+        raise ForbiddenError, size_checker.error_message.push_error
+      end
+    end
+
+    def check_push_size!
+      return unless check_size_limit?
+
+      # If there are worktrees with a HEAD pointing to a non-existent object,
+      # calls to `git rev-list --all` will fail in git 2.15+. This should also
+      # clear stale lock files.
+      repository.clean_stale_repository_files
+
+      # Use #check_repository_disk_size to get correct push size whenever a lot of changes
+      # gets pushed at the same time containing the same blobs. This is only
+      # doable if GIT_OBJECT_DIRECTORY_RELATIVE env var is set and happens
+      # when git push comes from CLI (not via UI and API).
+      #
+      # Fallback to determining push size using the changes_list so we can still
+      # determine the push size if env var isn't set (e.g. changes are made
+      # via UI and API).
+      if check_quarantine_size?
+        check_repository_disk_size
+      else
+        check_changes_size
+      end
+    end
+
+    def check_quarantine_size?
+      git_env = ::Gitlab::Git::HookEnv.all(repository.gl_repository)
+
+      git_env['GIT_OBJECT_DIRECTORY_RELATIVE'].present?
+    end
+
+    def check_repository_disk_size
+      check_size_against_limit(repository.object_directory_size)
+    end
+
+    def check_changes_size
+      changes_size = 0
+
+      changes_list.each do |change|
+        changes_size += repository.new_blobs(change[:newrev]).sum(&:size) # rubocop: disable CodeReuse/ActiveRecord
+
+        check_size_against_limit(changes_size)
+      end
+    end
+
+    def check_size_against_limit(size)
+      if size_checker.changes_will_exceed_size_limit?(size)
+        raise ForbiddenError, size_checker.error_message.new_changes_error
+      end
+    end
+
+    def check_size_limit?
+      strong_memoize(:check_size_limit) do
+        changes_list.any? { |change| !Gitlab::Git.blank_ref?(change[:newrev]) }
+      end
+    end
+
+    def size_checker
+      container.repository_size_checker
     end
   end
 end

lib/gitlab/git_access_snippet.rb

@@ -14,6 +14,8 @@ module Gitlab
 
     attr_reader :snippet
 
+    alias_method :container, :snippet
+
     def initialize(actor, snippet, protocol, **kwargs)
       @snippet = snippet
 
@@ -53,11 +55,6 @@ module Gitlab
       check_change_access!
     end
 
-    override :repository
-    def repository
-      snippet&.repository
-    end
-
     def check_snippet_accessibility!
       if snippet.blank?
         raise NotFoundError, ERROR_MESSAGES[:snippet_not_found]
@@ -89,11 +86,15 @@ module Gitlab
         raise ForbiddenError, ERROR_MESSAGES[:update_snippet]
       end
 
+      check_size_before_push!
+
       changes_list.each do |change|
         # If user does not have access to make at least one change, cancel all
         # push by allowing the exception to bubble up
         check_single_change_access(change)
       end
+
+      check_push_size!
     end
 
     def check_single_change_access(change)

lib/gitlab/repository_size_checker.rb

@@ -5,12 +5,14 @@ module Gitlab
   class RepositorySizeChecker
     attr_reader :limit
 
+    # @param current_size_proc [Proc] returns repository size in bytes
     def initialize(current_size_proc:, limit:, enabled: true)
       @current_size_proc = current_size_proc
       @limit = limit
       @enabled = enabled && limit != 0
     end
 
+    # @return [Integer] bytes
     def current_size
       @current_size ||= @current_size_proc.call
     end

spec/lib/gitlab/git_access_snippet_spec.rb

@@ -11,6 +11,7 @@ describe Gitlab::GitAccessSnippet do
   let_it_be(:user) { create(:user) }
   let_it_be(:project) { create(:project, :public) }
   let_it_be(:snippet) { create(:project_snippet, :public, :repository, project: project) }
+  let(:repository) { snippet.repository }
 
   let(:actor) { user }
   let(:protocol) { 'ssh' }
@@ -211,6 +212,84 @@ describe Gitlab::GitAccessSnippet do
     end
   end
 
+  describe 'repository size restrictions' do
+    let(:snippet) { create(:personal_snippet, :public, :repository) }
+    let(:actor) { snippet.author }
+
+    let(:oldrev) { TestEnv::BRANCH_SHA["snippet/single-file"] }
+    let(:newrev) { TestEnv::BRANCH_SHA["snippet/edit-file"] }
+    let(:ref) { "refs/heads/snippet/edit-file" }
+    let(:changes) { "#{oldrev} #{newrev} #{ref}" }
+
+    shared_examples_for 'a push to repository already over the limit' do
+      it 'errs' do
+        expect(snippet.repository_size_checker).to receive(:above_size_limit?).and_return(true)
+
+        expect do
+          push_access_check
+        end.to raise_error(described_class::ForbiddenError, /Your push has been rejected/)
+      end
+    end
+
+    shared_examples_for 'a push to repository below the limit' do
+      it 'does not err' do
+        expect(snippet.repository_size_checker).to receive(:above_size_limit?).and_return(false)
+        expect(snippet.repository_size_checker)
+          .to receive(:changes_will_exceed_size_limit?)
+            .with(change_size)
+            .and_return(false)
+
+        expect { push_access_check }.not_to raise_error
+      end
+    end
+
+    shared_examples_for 'a push to repository to make it over the limit' do
+      it 'errs' do
+        expect(snippet.repository_size_checker).to receive(:above_size_limit?).and_return(false)
+        expect(snippet.repository_size_checker)
+          .to receive(:changes_will_exceed_size_limit?)
+            .with(change_size)
+            .and_return(true)
+
+        expect do
+          push_access_check
+        end.to raise_error(described_class::ForbiddenError, /Your push to this repository would cause it to exceed the size limit/)
+      end
+    end
+
+    context 'when GIT_OBJECT_DIRECTORY_RELATIVE env var is set' do
+      let(:change_size) { 100 }
+
+      before do
+        allow(Gitlab::Git::HookEnv)
+          .to receive(:all)
+            .with(repository.gl_repository)
+            .and_return({ 'GIT_OBJECT_DIRECTORY_RELATIVE' => 'objects' })
+
+        # Stub the object directory size to "simulate" quarantine size
+        allow(repository).to receive(:object_directory_size).and_return(change_size)
+      end
+
+      it_behaves_like 'a push to repository already over the limit'
+      it_behaves_like 'a push to repository below the limit'
+      it_behaves_like 'a push to repository to make it over the limit'
+    end
+
+    context 'when GIT_OBJECT_DIRECTORY_RELATIVE env var is not set' do
+      let(:change_size) { 200 }
+
+      before do
+        allow(snippet.repository).to receive(:new_blobs).and_return(
+          [double(:blob, size: change_size)]
+        )
+      end
+
+      it_behaves_like 'a push to repository already over the limit'
+      it_behaves_like 'a push to repository below the limit'
+      it_behaves_like 'a push to repository to make it over the limit'
+    end
+  end
+
   private
 
   def raise_snippet_not_found

spec/models/snippet_spec.rb

@@ -696,6 +696,23 @@ describe Snippet do
     end
   end
 
+  describe '#repository_size_checker' do
+    subject { build(:personal_snippet) }
+
+    let(:checker) { subject.repository_size_checker }
+    let(:current_size) { 60 }
+
+    before do
+      allow(subject.repository).to receive(:_uncached_size).and_return(current_size)
+    end
+
+    it 'sets up size checker', :aggregate_failures do
+      expect(checker.current_size).to eq(current_size.megabytes)
+      expect(checker.limit).to eq(Gitlab::CurrentSettings.snippet_size_limit)
+      expect(checker.enabled?).to be_truthy
+    end
+  end
+
   describe '#can_cache_field?' do
     using RSpec::Parameterized::TableSyntax
 

spec/support/helpers/test_env.rb

@@ -60,11 +60,11 @@ module TestEnv
     'merge-commit-analyze-before'        => '1adbdef',
     'merge-commit-analyze-side-branch'   => '8a99451',
     'merge-commit-analyze-after'         => '646ece5',
-    'snippet/single-file'                => '43e4080',
-    'snippet/multiple-files'             => 'b80faa8',
-    'snippet/rename-and-edit-file'       => '220a1e4',
-    'snippet/edit-file'                  => 'c2f074f',
-    'snippet/no-files'                   => '671aaa8',
+    'snippet/single-file'                => '43e4080aaa14fc7d4b77ee1f5c9d067d5a7df10e',
+    'snippet/multiple-files'             => 'b80faa8c5b2b62f6489a0d84755580e927e1189b',
+    'snippet/rename-and-edit-file'       => '220a1e4b4dff37feea0625a7947a4c60fbe78365',
+    'snippet/edit-file'                  => 'c2f074f4f26929c92795a75775af79a6ed6d8430',
+    'snippet/no-files'                   => '671aaa842a4875e5f30082d1ab6feda345fdb94d',
     '2-mb-file'                          => 'bf12d25',
     'before-create-delete-modify-move'   => '845009f',
     'between-create-delete-modify-move'  => '3f5f443',