Build direct upload token based on attachment name

actionview/lib/action_view/helpers/form_tag_helper.rb

@@ -956,7 +956,13 @@ def convert_direct_upload_option_to_url(name, options)
               attachment_reflection = options[:object].class.reflect_on_attachment(name)
 
               if service_name = attachment_reflection.options[:service_name]
-                options["data-direct-upload-service"] = ActiveStorage.verifier.generate(service_name)
+                class_with_attachment = "#{options[:object].class.name}##{name}"
+                options["data-direct-upload-attachment-name"] = class_with_attachment
+                options["data-direct-upload-token"] = ActiveStorage::DirectUploadToken.generate_direct_upload_token(
+                  class_with_attachment,
+                  service_name,
+                  session
+                )
               end
             end
           end

activestorage/app/controllers/active_storage/direct_uploads_controller.rb

@@ -4,18 +4,20 @@
 # When the client-side upload is completed, the signed_blob_id can be submitted as part of the form to reference
 # the blob that was created up front.
 class ActiveStorage::DirectUploadsController < ActiveStorage::BaseController
+  include ActiveStorage::DirectUploadToken
+
   def create
     blob = ActiveStorage::Blob.create_before_direct_upload!(**blob_args.merge(service_name: verified_service_name))
     render json: direct_upload_json(blob)
   end
 
   private
     def blob_args
-      params.require(:blob).permit(:filename, :byte_size, :checksum, :content_type, :service_name, metadata: {}).to_h.symbolize_keys
+      params.require(:blob).permit(:filename, :byte_size, :checksum, :content_type, metadata: {}).to_h.symbolize_keys
     end
 
     def verified_service_name
-      ActiveStorage.verifier.verify(blob_args[:service_name]) if blob_args[:service_name]
+      verify_direct_upload_token(params[:blob][:direct_upload_token], params[:blob][:attachment_name], session)
     end
 
     def direct_upload_json(blob)

activestorage/app/javascript/activestorage/blob_record.js

@@ -1,15 +1,16 @@
 import { getMetaValue } from "./helpers"
 
 export class BlobRecord {
-  constructor(file, checksum, url, serviceName) {
+  constructor(file, checksum, url, directUploadToken, attachmentName) {
     this.file = file
 
     this.attributes = {
       filename: file.name,
       content_type: file.type || "application/octet-stream",
       byte_size: file.size,
       checksum: checksum,
-      service_name: serviceName
+      direct_upload_token: directUploadToken,
+      attachment_name: attachmentName
     }
 
     this.xhr = new XMLHttpRequest

activestorage/app/javascript/activestorage/direct_upload.js

@@ -5,11 +5,12 @@ import { BlobUpload } from "./blob_upload"
 let id = 0
 
 export class DirectUpload {
-  constructor(file, url, serviceName, delegate) {
+  constructor(file, url, serviceName, attachmentName, delegate) {
     this.id = ++id
     this.file = file
     this.url = url
     this.serviceName = serviceName
+    this.attachmentName = attachmentName
     this.delegate = delegate
   }
 
@@ -20,7 +21,7 @@ export class DirectUpload {
         return
       }
 
-      const blob = new BlobRecord(this.file, checksum, this.url, this.serviceName)
+      const blob = new BlobRecord(this.file, checksum, this.url, this.serviceName, this.attachmentName)
       notify(this.delegate, "directUploadWillCreateBlobWithXHR", blob.xhr)
 
       blob.create(error => {

activestorage/app/javascript/activestorage/direct_upload_controller.js

@@ -5,7 +5,7 @@ export class DirectUploadController {
   constructor(input, file) {
     this.input = input
     this.file = file
-    this.directUpload = new DirectUpload(this.file, this.url, this.serviceName, this)
+    this.directUpload = new DirectUpload(this.file, this.url, this.directUploadToken, this.attachmentName, this)
     this.dispatch("initialize")
   }
 
@@ -41,8 +41,12 @@ export class DirectUploadController {
     return this.input.getAttribute("data-direct-upload-url")
   }
 
-  get serviceName() {
-    return this.input.getAttribute("data-direct-upload-service")
+  get directUploadToken() {
+    return this.input.getAttribute("data-direct-upload-token")
+  }
+
+  get attachmentName() {
+    return this.input.getAttribute("data-direct-upload-attachment-name")
   }
 
   dispatch(name, detail = {}) {

activestorage/lib/active_storage.rb

@@ -41,6 +41,7 @@ module ActiveStorage
   autoload :Service
   autoload :Previewer
   autoload :Analyzer
+  autoload :DirectUploadToken
 
   mattr_accessor :logger
   mattr_accessor :verifier

activestorage/lib/active_storage/direct_upload_token.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module ActiveStorage
+  module DirectUploadToken
+    SEPARATOR = "."
+
+    module_function
+
+    def generate_direct_upload_token(attachment_name, service_name, session)
+      token = per_attachment_direct_upload_token(session, attachment_name)
+      encode_csrf_token([service_name, token].join(SEPARATOR))
+    end
+
+    def verify_direct_upload_token(token, attachment_name, session)
+      return if token.nil?
+
+      service_name, csrf_token = decode_csrf_token(token).split(SEPARATOR)
+      return service_name if valid_per_attachment_token?(csrf_token, attachment_name, session)
+
+      raise ActiveStorage::WrongDirectUploadTokenError
+    end
+
+    def valid_per_attachment_token?(token, attachment_name, session) # :doc:
+      correct_token = per_attachment_direct_upload_token(session, attachment_name)
+      ActiveSupport::SecurityUtils.fixed_length_secure_compare(token, correct_token)
+    rescue ArgumentError
+      raise ActiveStorage::WrongDirectUploadTokenError
+    end
+
+    def per_attachment_direct_upload_token(session, attachment_name) # :doc:
+      direct_upload_token_hmac(session, "direct_upload##{attachment_name}")
+    end
+
+    def direct_upload_token_hmac(session, identifier) # :doc:
+      OpenSSL::HMAC.digest(
+        OpenSSL::Digest::SHA256.new,
+        real_direct_upload_token(session),
+        identifier
+      )
+    end
+
+    DIRECT_UPLOAD_TOKEN_LENGTH = 32
+
+    def real_direct_upload_token(session) # :doc:
+      session[:_direct_upload_token] ||= SecureRandom.urlsafe_base64(DIRECT_UPLOAD_TOKEN_LENGTH, padding: false)
+      encode_csrf_token(session[:_direct_upload_token])
+    end
+
+    def decode_csrf_token(encoded_token) # :nodoc:
+      Base64.urlsafe_decode64(encoded_token)
+    end
+
+    def encode_csrf_token(raw_token) # :nodoc:
+      Base64.urlsafe_encode64(raw_token)
+    end
+  end
+end

activestorage/lib/active_storage/errors.rb

@@ -23,4 +23,7 @@ class IntegrityError < Error; end
   # Raised when ActiveStorage::Blob#download is called on a blob where the
   # backing file is no longer present in its service.
   class FileNotFoundError < Error; end
+
+  # Raised when direct upload fails because of wrong token
+  class WrongDirectUploadTokenError < Error; end
 end

activestorage/test/controllers/direct_uploads_controller_test.rb

@@ -138,10 +138,9 @@ class ActiveStorage::DiskDirectUploadsControllerTest < ActionDispatch::Integrati
       "platform": "my_platform",
       "library_ID": "12345"
     }
-    service_name = ActiveStorage.verifier.generate("local")
 
     post rails_direct_uploads_url, params: { blob: {
-      filename: "hello.txt", byte_size: 6, checksum: checksum, content_type: "text/plain", metadata: metadata, service_name: service_name } }
+      filename: "hello.txt", byte_size: 6, checksum: checksum, content_type: "text/plain", metadata: metadata } }
 
     @response.parsed_body.tap do |details|
       assert_equal ActiveStorage::Blob.find(details["id"]), ActiveStorage::Blob.find_signed!(details["signed_id"])
@@ -150,13 +149,12 @@ class ActiveStorage::DiskDirectUploadsControllerTest < ActionDispatch::Integrati
       assert_equal checksum, details["checksum"]
       assert_equal metadata, details["metadata"].transform_keys(&:to_sym)
       assert_equal "text/plain", details["content_type"]
-      assert_equal "local", details["service_name"]
       assert_match(/rails\/active_storage\/disk/, details["direct_upload"]["url"])
       assert_equal({ "Content-Type" => "text/plain" }, details["direct_upload"]["headers"])
     end
   end
 
-  test "creating new direct upload does not include root in json" do
+  test "handling direct upload with custom service name" do
     checksum = Digest::MD5.base64digest("Hello")
     metadata = {
       "foo": "bar",
@@ -166,14 +164,18 @@ class ActiveStorage::DiskDirectUploadsControllerTest < ActionDispatch::Integrati
       "library_ID": "12345"
     }
 
-    set_include_root_in_json(true) do
-      post rails_direct_uploads_url, params: { blob: {
-        filename: "hello.txt", byte_size: 6, checksum: checksum, content_type: "text/plain", metadata: metadata } }
-    end
+    post rails_direct_uploads_url, params: { blob: {
+      filename: "hello.txt", byte_size: 6, checksum: checksum, content_type: "text/plain", metadata: metadata } }
 
     @response.parsed_body.tap do |details|
-      assert_nil details["blob"]
-      assert_not_nil details["id"]
+      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 metadata, details["metadata"].transform_keys(&:to_sym)
+      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
 

activestorage/test/direct_upload_token_test.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class DirectUploadTokenTest < ActionController::TestCase
+  setup do
+    @session = {}
+    @service_name = "local"
+    @avatar_attachment_name = "User#avatar"
+  end
+
+  def test_validates_correct_token
+    token = ActiveStorage::DirectUploadToken.generate_direct_upload_token(@avatar_attachment_name, @service_name, @session)
+    verified_service_name = ActiveStorage::DirectUploadToken.verify_direct_upload_token(token, @avatar_attachment_name, @session)
+    assert_equal verified_service_name, @service_name
+  end
+
+  def test_not_validates_token_when_session_is_empty
+    token = ActiveStorage::DirectUploadToken.generate_direct_upload_token(@avatar_attachment_name, @service_name, {})
+
+    assert_raises(ActiveStorage::WrongDirectUploadTokenError) do
+      ActiveStorage::DirectUploadToken.verify_direct_upload_token(token, @avatar_attachment_name, @session)
+    end
+  end
+
+  def test_not_validates_token_from_different_attachment
+    background_attachment_name = "User#background"
+    token = ActiveStorage::DirectUploadToken.generate_direct_upload_token(background_attachment_name, @service_name, @session)
+
+    assert_raises(ActiveStorage::WrongDirectUploadTokenError) do
+      ActiveStorage::DirectUploadToken.verify_direct_upload_token(token, @avatar_attachment_name, @session)
+    end
+  end
+
+  def test_not_validates_token_from_different_session
+    token = ActiveStorage::DirectUploadToken.generate_direct_upload_token(@avatar_attachment_name, @service_name, @session)
+
+    another_session = {}
+    ActiveStorage::DirectUploadToken.generate_direct_upload_token(@avatar_attachment_name, @service_name, another_session)
+
+    assert_raises(ActiveStorage::WrongDirectUploadTokenError) do
+      ActiveStorage::DirectUploadToken.verify_direct_upload_token(token, @avatar_attachment_name, another_session)
+    end
+  end
+end