Expire caching when a download fail while proxying in ActiveStorage

Fixes rails#51284

Both Proxy controllers in ActiveStorage set the caching headers early
before streaming. In some instances, it is possible for the file
download to fail before we send the first byte (response not yet
committed). In those instance we invalidate the cache and return a
better response status before closing the stream.

activestorage/app/controllers/active_storage/blobs/proxy_controller.rb

@@ -15,12 +15,10 @@ def show
     if request.headers["Range"].present?
       send_blob_byte_range_data @blob, request.headers["Range"]
     else
-      http_cache_forever public: true do
-        response.headers["Accept-Ranges"] = "bytes"
-        response.headers["Content-Length"] = @blob.byte_size.to_s
+      response.headers["Accept-Ranges"] = "bytes"
+      response.headers["Content-Length"] = @blob.byte_size.to_s
 
-        send_blob_stream @blob, disposition: params[:disposition]
-      end
+      send_blob_stream @blob, disposition: params[:disposition]
     end
   end
 end

activestorage/app/controllers/active_storage/representations/proxy_controller.rb

@@ -11,8 +11,6 @@ class ActiveStorage::Representations::ProxyController < ActiveStorage::Represent
   include ActiveStorage::DisableSession
 
   def show
-    http_cache_forever public: true do
-      send_blob_stream @representation, disposition: params[:disposition]
-    end
+    send_blob_stream @representation, disposition: params[:disposition]
   end
 end

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

@@ -54,12 +54,21 @@ def send_blob_byte_range_data(blob, range_header, disposition: nil)
     # Stream the blob from storage directly to the response. The disposition can be controlled by setting +disposition+.
     # The content type and filename is set directly from the +blob+.
     def send_blob_stream(blob, disposition: nil) # :doc:
-      send_stream(
+      http_cache_forever public: true do
+        send_stream(
           filename: blob.filename.sanitized,
           disposition: blob.forced_disposition_for_serving || disposition || DEFAULT_BLOB_STREAMING_DISPOSITION,
           type: blob.content_type_for_serving) do |stream|
-        blob.download do |chunk|
-          stream.write chunk
+          blob.download do |chunk|
+            stream.write chunk
+          end
+        rescue ActiveStorage::FileNotFoundError
+          expires_now
+          head :not_found
+        rescue
+          expires_now
+          head :internal_server_error
+          raise
         end
       end
     end

activestorage/test/controllers/blobs/proxy_controller_test.rb

@@ -16,6 +16,32 @@ class ActiveStorage::Blobs::ProxyControllerTest < ActionDispatch::IntegrationTes
     assert_equal "max-age=3155695200, public", response.headers["Cache-Control"]
   end
 
+  test "invalidates cache and returns a 404 if the file is not found on download" do
+    blob = create_file_blob(filename: "racecar.jpg")
+    mock_download = lambda do |_|
+      raise ActiveStorage::FileNotFoundError.new "File still uploading!"
+    end
+    blob.service.stub(:download, mock_download) do
+      get rails_storage_proxy_url(blob)
+    end
+    assert_response :not_found
+    assert_equal "no-cache", response.headers["Cache-Control"]
+  end
+
+
+  test "invalidates cache and returns a 500 if an error is raised on download" do
+    blob = create_file_blob(filename: "racecar.jpg")
+    mock_download = lambda do |_|
+      raise StandardError.new "Something is not cool!"
+    end
+    blob.service.stub(:download, mock_download) do
+      get rails_storage_proxy_url(blob)
+    end
+    assert_response :internal_server_error
+    assert_equal "no-cache", response.headers["Cache-Control"]
+  end
+
+
   test "forcing Content-Type to binary" do
     get rails_storage_proxy_url(create_blob(content_type: "text/html"))
     assert_equal "application/octet-stream", response.headers["Content-Type"]

activestorage/test/controllers/representations/proxy_controller_test.rb

@@ -2,6 +2,7 @@
 
 require "test_helper"
 require "database/setup"
+require "minitest/mock"
 
 class ActiveStorage::Representations::ProxyControllerWithVariantsTest < ActionDispatch::IntegrationTest
   setup do
@@ -84,6 +85,39 @@ class ActiveStorage::Representations::ProxyControllerWithVariantsWithStrictLoadi
     assert_match(/^inline/, response.headers["Content-Disposition"])
     assert_equal @blob.variant(@transformations).download, response.body
   end
+
+
+  test "invalidates cache and returns a 404 if the file is not found on download" do
+    #This mock requires a pre-processed variant as processing the variant will call to download
+    mock_download = lambda do |_|
+      raise ActiveStorage::FileNotFoundError.new "File still uploading!"
+    end
+
+    @blob.service.stub(:download, mock_download) do
+      get rails_blob_representation_proxy_url(
+        filename: @blob.filename,
+        signed_blob_id: @blob.signed_id,
+        variation_key: ActiveStorage::Variation.encode(@transformations))
+    end
+    assert_response :not_found
+    assert_equal "no-cache", response.headers["Cache-Control"]
+  end
+
+  test "invalidates cache and returns a 500 if the an error is raised on download" do
+    #This mock requires a pre-processed variant as processing the variant will call to download
+    mock_download = lambda do |_|
+      raise StandardError.new "Something is not cool!"
+    end
+
+    @blob.service.stub(:download, mock_download) do
+      get rails_blob_representation_proxy_url(
+        filename: @blob.filename,
+        signed_blob_id: @blob.signed_id,
+        variation_key: ActiveStorage::Variation.encode(@transformations))
+    end
+    assert_response :internal_server_error
+    assert_equal "no-cache", response.headers["Cache-Control"]
+  end
 end
 
 class ActiveStorage::Representations::ProxyControllerWithPreviewsTest < ActionDispatch::IntegrationTest