Merge pull request rails#41391 from santib/fix-disk-storage-protocol

Fix rails#41388 by preserving protocol and port when generating routes

Author: kaspth

activestorage/CHANGELOG.md

@@ -1,3 +1,8 @@
+*   Deprecate `ActiveStorage::Current.host` in favor of `ActiveStorage::Current.url_options` which accepts
+    a host, protocol and port.
+
+    *Santiago Bartesaghi*
+
 *   Allow using [IAM](https://cloud.google.com/storage/docs/access-control/signed-urls) when signing URLs with GCS.
 
     ```yaml

activestorage/app/controllers/concerns/active_storage/set_current.rb

@@ -1,15 +1,15 @@
 # frozen_string_literal: true
 
-# Sets the <tt>ActiveStorage::Current.host</tt> attribute, which the disk service uses to generate URLs.
+# Sets the <tt>ActiveStorage::Current.url_options</tt> attribute, which the disk service uses to generate URLs.
 # Include this concern in custom controllers that call ActiveStorage::Blob#url,
 # ActiveStorage::Variant#url, or ActiveStorage::Preview#url so the disk service can
-# generate URLs using the same host, protocol, and base path as the current request.
+# generate URLs using the same host, protocol, and port as the current request.
 module ActiveStorage::SetCurrent
   extend ActiveSupport::Concern
 
   included do
     before_action do
-      ActiveStorage::Current.host = request.base_url
+      ActiveStorage::Current.url_options = { protocol: request.protocol, host: request.host, port: request.port }
     end
   end
 end

activestorage/app/models/active_storage/current.rb

@@ -1,5 +1,15 @@
 # frozen_string_literal: true
 
 class ActiveStorage::Current < ActiveSupport::CurrentAttributes #:nodoc:
-  attribute :host
+  attribute :url_options
+
+  def host=(host)
+    ActiveSupport::Deprecation.warn("ActiveStorage::Current.host= is deprecated, instead use ActiveStorage::Current.url_options=")
+    self.url_options = { host: host }
+  end
+
+  def host
+    ActiveSupport::Deprecation.warn("ActiveStorage::Current.host is deprecated, instead use ActiveStorage::Current.url_options")
+    self.url_options&.dig(:host)
+  end
 end

activestorage/lib/active_storage/service/disk_service.rb

@@ -86,11 +86,9 @@ def url_for_direct_upload(key, expires_in:, content_type:, content_length:, chec
           purpose: :blob_token
         )
 
-        generated_url = url_helpers.update_rails_disk_service_url(verified_token_with_expiration, host: current_host)
-
-        payload[:url] = generated_url
-
-        generated_url
+        url_helpers.update_rails_disk_service_url(verified_token_with_expiration, url_options).tap do |generated_url|
+          payload[:url] = generated_url
+        end
       end
     end
 
@@ -124,18 +122,11 @@ def generate_url(key, expires_in:, filename:, content_type:, disposition:)
           purpose: :blob_key
         )
 
-        if current_host.blank?
-          raise ArgumentError, "Cannot generate URL for #{filename} using Disk service, please set ActiveStorage::Current.host."
+        if url_options.blank?
+          raise ArgumentError, "Cannot generate URL for #{filename} using Disk service, please set ActiveStorage::Current.url_options."
         end
 
-        current_uri = URI.parse(current_host)
-
-        url_helpers.rails_disk_service_url(verified_key_with_expiration,
-          protocol: current_uri.scheme,
-          host: current_uri.host,
-          port: current_uri.port,
-          filename: filename
-        )
+        url_helpers.rails_disk_service_url(verified_key_with_expiration, filename: filename, **url_options)
       end
 
 
@@ -168,8 +159,8 @@ def url_helpers
         @url_helpers ||= Rails.application.routes.url_helpers
       end
 
-      def current_host
-        ActiveStorage::Current.host
+      def url_options
+        ActiveStorage::Current.url_options
       end
   end
 end

activestorage/test/service/disk_service_test.rb

@@ -12,6 +12,22 @@ class ActiveStorage::Service::DiskServiceTest < ActiveSupport::TestCase
     assert_equal :tmp, @service.name
   end
 
+  test "url_for_direct_upload" do
+    original_url_options = Rails.application.routes.default_url_options.dup
+    Rails.application.routes.default_url_options.merge!(protocol: "http", host: "test.example.com", port: 3001)
+
+    key      = SecureRandom.base58(24)
+    data     = "Something else entirely!"
+    checksum = Digest::MD5.base64digest(data)
+
+    begin
+      assert_match(/^https:\/\/example.com\/rails\/active_storage\/disk\/.*$/,
+        @service.url_for_direct_upload(key, expires_in: 5.minutes, content_type: "text/plain", content_length: data.size, checksum: checksum))
+    ensure
+      Rails.application.routes.default_url_options = original_url_options
+    end
+  end
+
   test "URL generation" do
     original_url_options = Rails.application.routes.default_url_options.dup
     Rails.application.routes.default_url_options.merge!(protocol: "http", host: "test.example.com", port: 3001)
@@ -23,14 +39,28 @@ class ActiveStorage::Service::DiskServiceTest < ActiveSupport::TestCase
     end
   end
 
-  test "URL generation without ActiveStorage::Current.host set" do
-    ActiveStorage::Current.host = nil
+  test "URL generation without ActiveStorage::Current.url_options set" do
+    ActiveStorage::Current.url_options = nil
 
     error = assert_raises ArgumentError do
       @service.url(@key, expires_in: 5.minutes, disposition: :inline, filename: ActiveStorage::Filename.new("avatar.png"), content_type: "image/png")
     end
 
-    assert_equal("Cannot generate URL for avatar.png using Disk service, please set ActiveStorage::Current.host.", error.message)
+    assert_equal("Cannot generate URL for avatar.png using Disk service, please set ActiveStorage::Current.url_options.", error.message)
+  end
+
+  test "URL generation keeps working with ActiveStorage::Current.host set" do
+    ActiveStorage::Current.url_options = nil
+    ActiveStorage::Current.host = "https://example.com"
+
+    original_url_options = Rails.application.routes.default_url_options.dup
+    Rails.application.routes.default_url_options.merge!(protocol: "http", host: "test.example.com", port: 3001)
+    begin
+      assert_match(/^http:\/\/example.com:3001\/rails\/active_storage\/disk\/.*\/avatar\.png$/,
+        @service.url(@key, expires_in: 5.minutes, disposition: :inline, filename: ActiveStorage::Filename.new("avatar.png"), content_type: "image/png"))
+    ensure
+      Rails.application.routes.default_url_options = original_url_options
+    end
   end
 
   test "headers_for_direct_upload generation" do

activestorage/test/test_helper.rb

@@ -53,7 +53,7 @@ class ActiveSupport::TestCase
   self.fixture_path = File.expand_path("fixtures", __dir__)
 
   setup do
-    ActiveStorage::Current.host = "https://example.com"
+    ActiveStorage::Current.url_options = { protocol: "https://", host: "example.com", port: nil }
   end
 
   teardown do