1 files changed, 38 insertions, 13 deletions

diff --git a/lib/gitlab/github_gists_import/importer/gist_importer.rb b/lib/gitlab/github_gists_import/importer/gist_importer.rb
index 4018f425e7c..71dfe5e2aa5 100644
--- a/lib/gitlab/github_gists_import/importer/gist_importer.rb
+++ b/lib/gitlab/github_gists_import/importer/gist_importer.rb
@@ -4,10 +4,13 @@ module Gitlab
   module GithubGistsImport
     module Importer
       class GistImporter
-        attr_reader :gist, :user
+        attr_reader :gist, :user, :snippet
 
         FileCountLimitError = Class.new(StandardError)
+        RepoSizeLimitError = Class.new(StandardError)
+        SnippetRepositoryError = Class.new(StandardError)
         FILE_COUNT_LIMIT_MESSAGE = 'Snippet maximum file count exceeded'
+        REPO_SIZE_LIMIT_MESSAGE = 'Snippet repository size exceeded'
 
         # gist - An instance of `Gitlab::GithubGistsImport::Representation::Gist`.
         def initialize(gist, user_id)
@@ -16,12 +19,15 @@ module Gitlab
         end
 
         def execute
-          snippet = build_snippet
-          import_repository(snippet) if snippet.save!
+          validate_gist!
 
-          return ServiceResponse.success unless max_snippet_files_count_exceeded?(snippet)
+          @snippet = build_snippet
+          import_repository if snippet.save!
+          validate_repository!
 
-          fail_and_track(snippet)
+          ServiceResponse.success
+        rescue FileCountLimitError, RepoSizeLimitError, SnippetRepositoryError => exception
+          fail_and_track(snippet, exception)
         end
 
         private
@@ -40,13 +46,13 @@ module Gitlab
           PersonalSnippet.new(attrs)
         end
 
-        def import_repository(snippet)
+        def import_repository
           resolved_address = get_resolved_address
 
           snippet.create_repository
           snippet.repository.fetch_as_mirror(gist.git_pull_url, forced: true, resolved_address: resolved_address)
         rescue StandardError
-          remove_snippet_and_repository(snippet)
+          remove_snippet_and_repository
 
           raise
         end
@@ -61,11 +67,19 @@ module Gitlab
           host.present? ? validated_pull_url.host.to_s : ''
         end
 
-        def max_snippet_files_count_exceeded?(snippet)
-          snippet.all_files.size > Snippet.max_file_limit
+        def check_gist_files_count!
+          return if gist.files.count <= Snippet.max_file_limit
+
+          raise FileCountLimitError, FILE_COUNT_LIMIT_MESSAGE
         end
 
-        def remove_snippet_and_repository(snippet)
+        def check_gist_repo_size!
+          return if gist.total_files_size <= Gitlab::CurrentSettings.snippet_size_limit
+
+          raise RepoSizeLimitError, REPO_SIZE_LIMIT_MESSAGE
+        end
+
+        def remove_snippet_and_repository
           snippet.repository.remove if snippet.repository_exists?
           snippet.destroy
         end
@@ -74,10 +88,21 @@ module Gitlab
           Gitlab::CurrentSettings.allow_local_requests_from_web_hooks_and_services?
         end
 
-        def fail_and_track(snippet)
-          remove_snippet_and_repository(snippet)
+        def fail_and_track(snippet, exception)
+          remove_snippet_and_repository if snippet
+
+          ServiceResponse.error(message: exception.message).track_exception(as: exception.class)
+        end
+
+        def validate_gist!
+          check_gist_files_count!
+          check_gist_repo_size!
+        end
+
+        def validate_repository!
+          result = Snippets::RepositoryValidationService.new(user, snippet).execute
 
-          ServiceResponse.error(message: FILE_COUNT_LIMIT_MESSAGE).track_exception(as: FileCountLimitError)
+          raise SnippetRepositoryError, result.message if result.error?
         end
       end
     end