1 files changed, 33 insertions, 16 deletions

diff --git a/app/services/pages/zip_directory_service.rb b/app/services/pages/zip_directory_service.rb
index a27ad5fda46..ba7a8571e88 100644
--- a/app/services/pages/zip_directory_service.rb
+++ b/app/services/pages/zip_directory_service.rb
@@ -2,37 +2,43 @@
 
 module Pages
   class ZipDirectoryService
-    InvalidArchiveError = Class.new(RuntimeError)
-    InvalidEntryError = Class.new(RuntimeError)
+    include BaseServiceUtility
+    include Gitlab::Utils::StrongMemoize
+
+    # used only to track exceptions in Sentry
+    InvalidEntryError = Class.new(StandardError)
 
     PUBLIC_DIR = 'public'
 
     def initialize(input_dir)
-      @input_dir = File.realpath(input_dir)
-      @output_file = File.join(@input_dir, "@migrated.zip") # '@' to avoid any name collision with groups or projects
+      @input_dir = input_dir
     end
 
     def execute
-      FileUtils.rm_f(@output_file)
+      return error("Can not find valid public dir in #{@input_dir}") unless valid_path?(public_dir)
+
+      output_file = File.join(real_dir, "@migrated.zip") # '@' to avoid any name collision with groups or projects
+
+      FileUtils.rm_f(output_file)
 
-      count = 0
-      ::Zip::File.open(@output_file, ::Zip::File::CREATE) do |zipfile|
+      entries_count = 0
+      ::Zip::File.open(output_file, ::Zip::File::CREATE) do |zipfile|
         write_entry(zipfile, PUBLIC_DIR)
-        count = zipfile.entries.count
+        entries_count = zipfile.entries.count
       end
 
-      [@output_file, count]
+      success(archive_path: output_file, entries_count: entries_count)
+    rescue => e
+      FileUtils.rm_f(output_file) if output_file
+      raise e
     end
 
     private
 
     def write_entry(zipfile, zipfile_path)
-      disk_file_path = File.join(@input_dir, zipfile_path)
+      disk_file_path = File.join(real_dir, zipfile_path)
 
       unless valid_path?(disk_file_path)
-        # archive without public directory is completelly unusable
-        raise InvalidArchiveError if zipfile_path == PUBLIC_DIR
-
         # archive with invalid entry will just have this entry missing
         raise InvalidEntryError
       end
@@ -71,13 +77,24 @@ module Pages
     def valid_path?(disk_file_path)
       realpath = File.realpath(disk_file_path)
 
-      realpath == File.join(@input_dir, PUBLIC_DIR) ||
-        realpath.start_with?(File.join(@input_dir, PUBLIC_DIR + "/"))
+      realpath == public_dir || realpath.start_with?(public_dir + "/")
     # happens if target of symlink isn't there
     rescue => e
-      Gitlab::ErrorTracking.track_exception(e, input_dir: @input_dir, disk_file_path: disk_file_path)
+      Gitlab::ErrorTracking.track_exception(e, input_dir: real_dir, disk_file_path: disk_file_path)
 
       false
     end
+
+    def real_dir
+      strong_memoize(:real_dir) do
+        File.realpath(@input_dir) rescue nil
+      end
+    end
+
+    def public_dir
+      strong_memoize(:public_dir) do
+        File.join(real_dir, PUBLIC_DIR) rescue nil
+      end
+    end
   end
 end