1 files changed, 93 insertions, 6 deletions

diff --git a/lib/gitlab/middleware/multipart.rb b/lib/gitlab/middleware/multipart.rb
index c0b671abd44..8e6ac7610f2 100644
--- a/lib/gitlab/middleware/multipart.rb
+++ b/lib/gitlab/middleware/multipart.rb
@@ -29,6 +29,8 @@ module Gitlab
   module Middleware
     class Multipart
       RACK_ENV_KEY = 'HTTP_GITLAB_WORKHORSE_MULTIPART_FIELDS'
+      JWT_PARAM_SUFFIX = '.gitlab-workhorse-upload'
+      JWT_PARAM_FIXED_KEY = 'upload'
 
       class Handler
         def initialize(env, message)
@@ -57,7 +59,8 @@ module Gitlab
 
           yield
         ensure
-          @open_files.each(&:close)
+          @open_files.compact
+                     .each(&:close)
         end
 
         # This function calls itself recursively
@@ -122,15 +125,89 @@ module Gitlab
 
         def allowed_paths
           [
+            Dir.tmpdir,
             ::FileUploader.root,
-            Gitlab.config.uploads.storage_path,
-            JobArtifactUploader.workhorse_upload_path,
-            LfsObjectUploader.workhorse_upload_path,
+            ::Gitlab.config.uploads.storage_path,
+            ::JobArtifactUploader.workhorse_upload_path,
+            ::LfsObjectUploader.workhorse_upload_path,
             File.join(Rails.root, 'public/uploads/tmp')
           ] + package_allowed_paths
         end
       end
 
+      # TODO this class is meant to replace Handler when the feature flag
+      # upload_middleware_jwt_params_handler is removed
+      class HandlerForJWTParams < Handler
+        def with_open_files
+          @rewritten_fields.keys.each do |field|
+            parsed_field = Rack::Utils.parse_nested_query(field)
+            raise "unexpected field: #{field.inspect}" unless parsed_field.count == 1
+
+            key, value = parsed_field.first
+            if value.nil? # we have a top level param, eg. field = 'foo' and not 'foo[bar]'
+              raise "invalid field: #{field.inspect}" if field != key
+
+              value = open_file(extract_upload_params_from(@request.params, with_prefix: key))
+              @open_files << value
+            else
+              value = decorate_params_value(value, @request.params[key])
+            end
+
+            update_param(key, value)
+          end
+
+          yield
+        ensure
+          @open_files.compact
+                     .each(&:close)
+        end
+
+        # This function calls itself recursively
+        def decorate_params_value(hash_path, value_hash)
+          unless hash_path.is_a?(Hash) && hash_path.count == 1
+            raise "invalid path: #{hash_path.inspect}"
+          end
+
+          path_key, path_value = hash_path.first
+
+          unless value_hash.is_a?(Hash) && value_hash[path_key]
+            raise "invalid value hash: #{value_hash.inspect}"
+          end
+
+          case path_value
+          when nil
+            value_hash[path_key] = open_file(extract_upload_params_from(value_hash[path_key]))
+            @open_files << value_hash[path_key]
+            value_hash
+          when Hash
+            decorate_params_value(path_value, value_hash[path_key])
+            value_hash
+          else
+            raise "unexpected path value: #{path_value.inspect}"
+          end
+        end
+
+        def open_file(params)
+          ::UploadedFile.from_params_without_field(params, allowed_paths)
+        end
+
+        private
+
+        def extract_upload_params_from(params, with_prefix: '')
+          param_key = "#{with_prefix}#{JWT_PARAM_SUFFIX}"
+          jwt_token = params[param_key]
+          raise "Empty JWT param: #{param_key}" if jwt_token.blank?
+
+          payload = Gitlab::Workhorse.decode_jwt(jwt_token).first
+          raise "Invalid JWT payload: not a Hash" unless payload.is_a?(Hash)
+
+          upload_params = payload.fetch(JWT_PARAM_FIXED_KEY, {})
+          raise "Empty params for: #{param_key}" if upload_params.empty?
+
+          upload_params
+        end
+      end
+
       def initialize(app)
         @app = app
       end
@@ -139,14 +216,24 @@ module Gitlab
         encoded_message = env.delete(RACK_ENV_KEY)
         return @app.call(env) if encoded_message.blank?
 
-        message = Gitlab::Workhorse.decode_jwt(encoded_message)[0]
+        message = ::Gitlab::Workhorse.decode_jwt(encoded_message)[0]
 
-        Handler.new(env, message).with_open_files do
+        handler_class.new(env, message).with_open_files do
           @app.call(env)
         end
       rescue UploadedFile::InvalidPathError => e
         [400, { 'Content-Type' => 'text/plain' }, e.message]
       end
+
+      private
+
+      def handler_class
+        if Feature.enabled?(:upload_middleware_jwt_params_handler)
+          ::Gitlab::Middleware::Multipart::HandlerForJWTParams
+        else
+          ::Gitlab::Middleware::Multipart::Handler
+        end
+      end
     end
   end
 end