Process workhorse accelerated wiki uploads

Wiki attachments can be workhorse accelerated.
This commit is backward compatible with older workhorse

3 files changed, 52 insertions, 8 deletions

diff --git a/lib/api/validations/types/workhorse_file.rb b/lib/api/validations/types/workhorse_file.rb
new file mode 100644
index 00000000000..18d111f6556
--- /dev/null
+++ b/lib/api/validations/types/workhorse_file.rb
@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+module API
+  module Validations
+    module Types
+      class WorkhorseFile < Virtus::Attribute
+        def coerce(input)
+          # Processing of multipart file objects
+          # is already taken care of by Gitlab::Middleware::Multipart.
+          # Nothing to do here.
+          input
+        end
+
+        def value_coerced?(value)
+          value.is_a?(::UploadedFile)
+        end
+      end
+    end
+  end
+end
diff --git a/lib/api/wikis.rb b/lib/api/wikis.rb
index 5724adb2c40..c5a5488950d 100644
--- a/lib/api/wikis.rb
+++ b/lib/api/wikis.rb
@@ -4,11 +4,21 @@ module API
   class Wikis < Grape::API
     helpers do
       def commit_params(attrs)
-        {
-          file_name: attrs[:file][:filename],
-          file_content: attrs[:file][:tempfile].read,
-          branch_name: attrs[:branch]
-        }
+        # In order to avoid service disruption this can work with an old workhorse without the acceleration
+        # the first branch of this if must be removed when we drop support for non accelerated uploads
+        if attrs[:file].is_a?(Hash)
+          {
+            file_name: attrs[:file][:filename],
+            file_content: attrs[:file][:tempfile].read,
+            branch_name: attrs[:branch]
+          }
+        else
+          {
+            file_name: attrs[:file].original_filename,
+            file_content: attrs[:file].read,
+            branch_name: attrs[:branch]
+          }
+        end
       end
 
       params :common_wiki_page_params do
@@ -106,7 +116,7 @@ module API
         success Entities::WikiAttachment
       end
       params do
-        requires :file, type: ::API::Validations::Types::SafeFile, desc: 'The attachment file to be uploaded'
+        requires :file, types: [::API::Validations::Types::SafeFile, ::API::Validations::Types::WorkhorseFile], desc: 'The attachment file to be uploaded'
         optional :branch, type: String, desc: 'The name of the branch'
       end
       post ":id/wikis/attachments" do
diff --git a/lib/gitlab/middleware/multipart.rb b/lib/gitlab/middleware/multipart.rb
index 86b28e4e20a..0ee9563c227 100644
--- a/lib/gitlab/middleware/multipart.rb
+++ b/lib/gitlab/middleware/multipart.rb
@@ -32,7 +32,7 @@ module Gitlab
 
       class Handler
         def initialize(env, message)
-          @request = ActionDispatch::Request.new(env)
+          @request = Rack::Request.new(env)
           @rewritten_fields = message['rewritten_fields']
           @open_files = []
         end
@@ -50,7 +50,7 @@ module Gitlab
               value = decorate_params_value(value, @request.params[key])
             end
 
-            @request.update_param(key, value)
+            update_param(key, value)
           end
 
           yield
@@ -92,6 +92,20 @@ module Gitlab
 
           ::UploadedFile.from_params(params, key, allowed_paths)
         end
+
+        # update_params ensures that both rails controllers and rack middleware can find
+        # workhorse accelerate files in the request
+        def update_param(key, value)
+          # we make sure we have key in POST otherwise update_params will add it in GET
+          @request.POST[key] ||= value
+
+          # this will force Rack::Request to properly update env keys
+          @request.update_param(key, value)
+
+          # ActionDispatch::Request is based on Rack::Request but it caches params
+          # inside other env keys, here we ensure everything is updated correctly
+          ActionDispatch::Request.new(@request.env).update_param(key, value)
+        end
       end
 
       def initialize(app)