From d07cce737c51d7bc8054069ba3243a0636c16bd4 Mon Sep 17 00:00:00 2001
From: Matt-Yorkley <9029026+Matt-Yorkley@users.noreply.github.com>
Date: Sat, 19 Aug 2023 18:01:41 +0100
Subject: [PATCH] Allow customisation of transferable headers

Headers specified in the :transfer_headers option are passed (transferred) from responses from the backend through to responses sent to the client (included cached entries).
---
 lib/rack/cache/context.rb |  2 +-
 lib/rack/cache/options.rb |  8 ++++++++
 test/context_test.rb      | 24 ++++++++++++++++++++++++
 3 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/lib/rack/cache/context.rb b/lib/rack/cache/context.rb
index 58996b7..583f4bf 100644
--- a/lib/rack/cache/context.rb
+++ b/lib/rack/cache/context.rb
@@ -237,7 +237,7 @@ def validate(entry)
 
         entry = entry.dup
         entry.headers.delete('date')
-        %w[Date expires cache-control etag last-modified].each do |name|
+        %w[Date expires cache-control etag last-modified].concat(transfer_headers).each do |name|
           next unless value = response.headers[name]
           entry.headers[name] = value
         end
diff --git a/lib/rack/cache/options.rb b/lib/rack/cache/options.rb
index 0007605..d2d5112 100644
--- a/lib/rack/cache/options.rb
+++ b/lib/rack/cache/options.rb
@@ -86,6 +86,13 @@ def option_name(key)
     # Default: ['set-cookie']
     option_accessor :ignore_headers
 
+    # Set of response headers that are transferred from the backend response
+    # onto the cache entry when validating a cached entity (after receiving a
+    # 304 response from the backend) before sending it to the client.
+    #
+    # Default: []
+    option_accessor :transfer_headers
+
     # Set of request headers that trigger "private" cache-control behavior
     # on responses that don't explicitly state whether the response is
     # public or private via a cache-control directive. Applications that use
@@ -149,6 +156,7 @@ def initialize_options(options={})
         'rack-cache.entitystore'      => 'heap:/',
         'rack-cache.default_ttl'      => 0,
         'rack-cache.ignore_headers'   => ['set-cookie'],
+        'rack-cache.transfer_headers' => [],
         'rack-cache.private_headers'  => ['Authorization', 'Cookie'],
         'rack-cache.allow_reload'     => false,
         'rack-cache.allow_revalidate' => false,
diff --git a/test/context_test.rb b/test/context_test.rb
index acf570a..15470d9 100644
--- a/test/context_test.rb
+++ b/test/context_test.rb
@@ -776,6 +776,30 @@
     cache.trace.wont_include :miss
   end
 
+  it 'includes specified additional headers from the backend in responses to the client' do
+    count = 0
+    respond_with do |req,res|
+      count += 1
+      res['ETAG'] = '"12345"'
+      res['unique-snowflake'] = '"so-special"'
+      res.status = (count == 1) ? 200 : 304
+    end
+
+    get '/', 'rack-cache.transfer_headers' => ['unique-snowflake']
+    assert app.called?
+    assert response.ok?
+    response.headers.must_include 'unique-snowflake'
+    response['unique-snowflake'].must_equal '"so-special"'
+    cache.trace.must_include :miss
+
+    get '/', 'rack-cache.transfer_headers' => ['unique-snowflake']
+    assert app.called?
+    assert response.ok?
+    response.headers.must_include 'unique-snowflake'
+    response['unique-snowflake'].must_equal '"so-special"'
+    cache.trace.must_include :valid
+  end
+
   it 'replaces cached responses when validation results in non-304 response' do
     timestamp = Time.now.httpdate
     count = 0