3 files changed, 31 insertions, 4 deletions
diff --git a/lib/safe_zip/entry.rb b/lib/safe_zip/entry.rb
index 52d70e83154..88647b9b1eb 100644
--- a/lib/safe_zip/entry.rb
+++ b/lib/safe_zip/entry.rb
@@ -25,8 +25,8 @@ module SafeZip
     end
 
     def extract
-      # do not extract if file is not part of target directory
-      return false unless matching_target_directory
+      # do not extract if file is not part of target directory or target file
+      return false unless matching_target_directory || matching_target_file
 
       # do not overwrite existing file
       raise SafeZip::Extract::AlreadyExistsError, "File already exists #{zip_entry.name}" if exist?
@@ -44,6 +44,8 @@ module SafeZip
       end
     rescue SafeZip::Extract::Error
       raise
+    rescue Zip::EntrySizeError => e
+      raise SafeZip::Extract::EntrySizeError, e.message
     rescue StandardError => e
       raise SafeZip::Extract::ExtractError, e.message
     end
@@ -84,6 +86,10 @@ module SafeZip
       params.matching_target_directory(path)
     end
 
+    def matching_target_file
+      params.matching_target_file(path)
+    end
+
     def read_symlink
       zip_archive.read(zip_entry)
     end
diff --git a/lib/safe_zip/extract.rb b/lib/safe_zip/extract.rb
index 74df7895afe..b86941e6bea 100644
--- a/lib/safe_zip/extract.rb
+++ b/lib/safe_zip/extract.rb
@@ -6,6 +6,7 @@ module SafeZip
     PermissionDeniedError = Class.new(Error)
     SymlinkSourceDoesNotExistError = Class.new(Error)
     UnsupportedEntryError = Class.new(Error)
+    EntrySizeError = Class.new(Error)
     AlreadyExistsError = Class.new(Error)
     NoMatchingError = Class.new(Error)
     ExtractError = Class.new(Error)
diff --git a/lib/safe_zip/extract_params.rb b/lib/safe_zip/extract_params.rb
index bd3b788bac9..96881ad1abc 100644
--- a/lib/safe_zip/extract_params.rb
+++ b/lib/safe_zip/extract_params.rb
@@ -4,11 +4,13 @@ module SafeZip
   class ExtractParams
     include Gitlab::Utils::StrongMemoize
 
-    attr_reader :directories, :extract_path
+    attr_reader :directories, :files, :extract_path
 
-    def initialize(directories:, to:)
+    def initialize(to:, directories: [], files: [])
       @directories = directories
+      @files = files
       @extract_path = ::File.realpath(to)
+      validate!
     end
 
     def matching_target_directory(path)
@@ -32,5 +34,23 @@ module SafeZip
         end
       end
     end
+
+    def matching_target_file(path)
+      target_files.include?(path)
+    end
+
+    private
+
+    def target_files
+      strong_memoize(:target_files) do
+        files.map do |file|
+          ::File.join(extract_path, file)
+        end
+      end
+    end
+
+    def validate!
+      raise ArgumentError, 'Either directories or files are required' if directories.empty? && files.empty?
+    end
   end
 end