Configure per-service request headers for direct uploads (#83)

* Configure per-service request headers for direct uploads

* Fix header hashes

app/controllers/active_storage/direct_uploads_controller.rb

@@ -4,11 +4,18 @@
 class ActiveStorage::DirectUploadsController < ActionController::Base
   def create
     blob = ActiveStorage::Blob.create_before_direct_upload!(blob_args)
-    render json: { upload_to_url: blob.service_url_for_direct_upload, signed_blob_id: blob.signed_id }
+    render json: direct_upload_json(blob)
   end
 
   private
     def blob_args
       params.require(:blob).permit(:filename, :byte_size, :checksum, :content_type, :metadata).to_h.symbolize_keys
     end
+
+    def direct_upload_json(blob)
+      blob.as_json(methods: :signed_id).merge(direct_upload: {
+        url: blob.service_url_for_direct_upload,
+        headers: blob.service_headers_for_direct_upload
+      })
+    end
 end

app/javascript/activestorage/blob_record.js

@@ -30,8 +30,10 @@ export class BlobRecord {
   requestDidLoad(event) {
     const { status, response } = this.xhr
     if (status >= 200 && status < 300) {
-      this.attributes.signed_id = response.signed_blob_id
-      this.uploadURL = response.upload_to_url
+      const { direct_upload } = response
+      delete response.direct_upload
+      this.attributes = response
+      this.directUploadData = direct_upload
       this.callback(null, this.toJSON())
     } else {
       this.requestDidError(event)

app/javascript/activestorage/blob_upload.js

@@ -3,10 +3,13 @@ export class BlobUpload {
     this.blob = blob
     this.file = blob.file
 
+    const { url, headers } = blob.directUploadData
+
     this.xhr = new XMLHttpRequest
-    this.xhr.open("PUT", blob.uploadURL, true)
-    this.xhr.setRequestHeader("Content-Type", blob.attributes.content_type)
-    this.xhr.setRequestHeader("Content-MD5", blob.attributes.checksum)
+    this.xhr.open("PUT", url, true)
+    for (const key in headers) {
+      this.xhr.setRequestHeader(key, headers[key])
+    }
     this.xhr.addEventListener("load", event => this.requestDidLoad(event))
     this.xhr.addEventListener("error", event => this.requestDidError(event))
   }

app/models/active_storage/blob.rb

@@ -63,14 +63,14 @@ class ActiveStorage::Blob < ActiveRecord::Base
     end
   end
 
-  
+
   # Returns a signed ID for this blob that's suitable for reference on the client-side without fear of tampering.
   # It uses the framework-wide verifier on `ActiveStorage.verifier`, but with a dedicated purpose.
   def signed_id
     ActiveStorage.verifier.generate(id, purpose: :blob_id)
   end
 
-  # Returns the key pointing to the file on the service that's associated with this blob. The key is in the 
+  # Returns the key pointing to the file on the service that's associated with this blob. The key is in the
   # standard secure-token format from Rails. So it'll look like: XTAPjJCJiuDrLk3TmwyJGpUo. This key is not intended
   # to be revealed directly to the user. Always refer to blobs using the signed_id or a verified form of the key.
   def key
@@ -130,6 +130,10 @@ class ActiveStorage::Blob < ActiveRecord::Base
     service.url_for_direct_upload key, expires_in: expires_in, content_type: content_type, content_length: byte_size, checksum: checksum
   end
 
+  # Returns a Hash of headers for `service_url_for_direct_upload` requests.
+  def service_headers_for_direct_upload
+    service.headers_for_direct_upload key, filename: filename, content_type: content_type, content_length: byte_size, checksum: checksum
+  end
 
   # Uploads the `io` to the service on the `key` for this blob. Blobs are intended to be immutable, so you shouldn't be
   # using this method after a file has already been uploaded to fit with a blob. If you want to create a derivative blob,
@@ -176,7 +180,6 @@ class ActiveStorage::Blob < ActiveRecord::Base
     ActiveStorage::PurgeJob.perform_later(self)
   end
 
-
   private
     def compute_checksum_in_chunks(io)
       Digest::MD5.new.tap do |checksum|

app/models/active_storage/filename.rb

@@ -35,6 +35,14 @@ class ActiveStorage::Filename
     sanitized.to_s
   end
 
+  def as_json(*)
+    to_s
+  end
+
+  def to_json
+    to_s
+  end
+
   def <=>(other)
     to_s.downcase <=> other.to_s.downcase
   end

lib/active_storage/service.rb

@@ -87,13 +87,18 @@ class ActiveStorage::Service
   end
 
   # Returns a signed, temporary URL that a direct upload file can be PUT to on the `key`.
-  # The URL will be valid for the amount of seconds specified in `expires_in`. 
+  # The URL will be valid for the amount of seconds specified in `expires_in`.
   # You most also provide the `content_type`, `content_length`, and `checksum` of the file
   # that will be uploaded. All these attributes will be validated by the service upon upload.
   def url_for_direct_upload(key, expires_in:, content_type:, content_length:, checksum:)
     raise NotImplementedError
   end
 
+  # Returns a Hash of headers for `url_for_direct_upload` requests.
+  def headers_for_direct_upload(key, filename:, content_type:, content_length:, checksum:)
+    {}
+  end
+
   private
     def instrument(operation, key, payload = {}, &block)
       ActiveSupport::Notifications.instrument(

lib/active_storage/service/disk_service.rb

@@ -98,6 +98,10 @@ class ActiveStorage::Service::DiskService < ActiveStorage::Service
     end
   end
 
+  def headers_for_direct_upload(key, content_type:, **)
+    { "Content-Type" => content_type }
+  end
+
   private
     def path_for(key)
       File.join root, folder_for(key), key

lib/active_storage/service/gcs_service.rb

@@ -68,6 +68,10 @@ class ActiveStorage::Service::GCSService < ActiveStorage::Service
     end
   end
 
+  def headers_for_direct_upload(key, content_type:, checksum:, **)
+    { "Content-Type" => content_type, "Content-MD5" => checksum }
+  end
+
   private
     def file_for(key)
       bucket.file(key)

lib/active_storage/service/s3_service.rb

@@ -72,6 +72,10 @@ class ActiveStorage::Service::S3Service < ActiveStorage::Service
     end
   end
 
+  def headers_for_direct_upload(key, content_type:, checksum:, **)
+    { "Content-Type" => content_type, "Content-MD5" => checksum }
+  end
+
   private
     def object_for(key)
       bucket.object(key)

test/controllers/direct_uploads_controller_test.rb

@@ -13,12 +13,19 @@ if SERVICE_CONFIGURATIONS[:s3]
     end
 
     test "creating new direct upload" do
+      checksum = Digest::MD5.base64digest("Hello")
+
       post rails_direct_uploads_url, params: { blob: {
-        filename: "hello.txt", byte_size: 6, checksum: Digest::MD5.base64digest("Hello"), content_type: "text/plain" } }
+        filename: "hello.txt", byte_size: 6, checksum: checksum, content_type: "text/plain" } }
 
       response.parsed_body.tap do |details|
-        assert_match(/#{SERVICE_CONFIGURATIONS[:s3][:bucket]}\.s3.(\S+)?amazonaws\.com/, details["upload_to_url"])
-        assert_equal "hello.txt", ActiveStorage::Blob.find_signed(details["signed_blob_id"]).filename.to_s
+        assert_equal ActiveStorage::Blob.find(details["id"]), ActiveStorage::Blob.find_signed(details["signed_id"])
+        assert_equal "hello.txt", details["filename"]
+        assert_equal 6, details["byte_size"]
+        assert_equal checksum, details["checksum"]
+        assert_equal "text/plain", details["content_type"]
+        assert_match /#{SERVICE_CONFIGURATIONS[:s3][:bucket]}\.s3.(\S+)?amazonaws\.com/, details["direct_upload"]["url"]
+        assert_equal({ "Content-Type" => "text/plain", "Content-MD5" => checksum }, details["direct_upload"]["headers"])
       end
     end
   end
@@ -40,12 +47,19 @@ if SERVICE_CONFIGURATIONS[:gcs]
     end
 
     test "creating new direct upload" do
+      checksum = Digest::MD5.base64digest("Hello")
+
       post rails_direct_uploads_url, params: { blob: {
-        filename: "hello.txt", byte_size: 6, checksum: Digest::MD5.base64digest("Hello"), content_type: "text/plain" } }
+        filename: "hello.txt", byte_size: 6, checksum: checksum, content_type: "text/plain" } }
 
       @response.parsed_body.tap do |details|
-        assert_match %r{storage\.googleapis\.com/#{@config[:bucket]}}, details["upload_to_url"]
-        assert_equal "hello.txt", ActiveStorage::Blob.find_signed(details["signed_blob_id"]).filename.to_s
+        assert_equal ActiveStorage::Blob.find(details["id"]), ActiveStorage::Blob.find_signed(details["signed_id"])
+        assert_equal "hello.txt", details["filename"]
+        assert_equal 6, details["byte_size"]
+        assert_equal checksum, details["checksum"]
+        assert_equal "text/plain", details["content_type"]
+        assert_match %r{storage\.googleapis\.com/#{@config[:bucket]}}, details["direct_upload"]["url"]
+        assert_equal({ "Content-Type" => "text/plain", "Content-MD5" => checksum }, details["direct_upload"]["headers"])
       end
     end
   end
@@ -55,12 +69,19 @@ end
 
 class ActiveStorage::DiskDirectUploadsControllerTest < ActionDispatch::IntegrationTest
   test "creating new direct upload" do
+    checksum = Digest::MD5.base64digest("Hello")
+
     post rails_direct_uploads_url, params: { blob: {
-      filename: "hello.txt", byte_size: 6, checksum: Digest::MD5.base64digest("Hello"), content_type: "text/plain" } }
+      filename: "hello.txt", byte_size: 6, checksum: checksum, content_type: "text/plain" } }
 
     @response.parsed_body.tap do |details|
-      assert_match /rails\/active_storage\/disk/, details["upload_to_url"]
-      assert_equal "hello.txt", ActiveStorage::Blob.find_signed(details["signed_blob_id"]).filename.to_s
+      assert_equal ActiveStorage::Blob.find(details["id"]), ActiveStorage::Blob.find_signed(details["signed_id"])
+      assert_equal "hello.txt", details["filename"]
+      assert_equal 6, details["byte_size"]
+      assert_equal checksum, details["checksum"]
+      assert_equal "text/plain", details["content_type"]
+      assert_match /rails\/active_storage\/disk/, details["direct_upload"]["url"]
+      assert_equal({ "Content-Type" => "text/plain" }, details["direct_upload"]["headers"])
     end
   end
 end