From dbced5bfaa06904ef89ad8a7d29d4605c4041062 Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuel.williams@oriontransfer.co.nz>
Date: Tue, 8 Sep 2015 17:25:56 +1200
Subject: [PATCH] Convenient cache and content type methods for
 `Rack::Response`.

---
 lib/rack.rb           |  1 +
 lib/rack/response.rb  | 22 ++++++++++++++++++++++
 test/spec_response.rb | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 56 insertions(+)

diff --git a/lib/rack.rb b/lib/rack.rb
index 7600c40f6..cab2bb8b8 100644
--- a/lib/rack.rb
+++ b/lib/rack.rb
@@ -27,6 +27,7 @@ module Rack
   SERVER_ADDR       = 'SERVER_ADDR'
   SERVER_PORT       = 'SERVER_PORT'
   CACHE_CONTROL     = 'Cache-Control'
+  EXPIRES           = 'Expires'
   CONTENT_LENGTH    = 'Content-Length'
   CONTENT_TYPE      = 'Content-Type'
   SET_COOKIE        = 'Set-Cookie'
diff --git a/lib/rack/response.rb b/lib/rack/response.rb
index 408a38fc8..fbbcb03ef 100644
--- a/lib/rack/response.rb
+++ b/lib/rack/response.rb
@@ -178,10 +178,16 @@ def add_header(key, v)
         end
       end
 
+      # Get the content type of the response.
       def content_type
         get_header CONTENT_TYPE
       end
 
+      # Set the content type of the response.
+      def content_type=(content_type)
+        set_header CONTENT_TYPE, content_type
+      end
+
       def media_type
         MediaType.type(content_type)
       end
@@ -228,6 +234,22 @@ def cache_control=(v)
         set_header CACHE_CONTROL, v
       end
 
+      # Specifies that the content shouldn't be cached. Overrides `cache!` if already called.
+      def do_not_cache!
+        set_header CACHE_CONTROL, "no-cache, must-revalidate"
+        set_header EXPIRES, Time.now.httpdate
+      end
+
+      # Specify that the content should be cached.
+      # @param duration [Integer] The number of seconds until the cache expires.
+      # @option directive [String] The cache control directive, one of "public", "private", "no-cache" or "no-store".
+      def cache!(duration = 3600, directive: "public")
+        unless headers[CACHE_CONTROL] =~ /no-cache/
+          set_header CACHE_CONTROL, "#{directive}, max-age=#{duration}"
+          set_header EXPIRES, (Time.now + duration).httpdate
+        end
+      end
+
       def etag
         get_header ETAG
       end
diff --git a/test/spec_response.rb b/test/spec_response.rb
index b2ba59a82..1dfafcdb5 100644
--- a/test/spec_response.rb
+++ b/test/spec_response.rb
@@ -30,6 +30,14 @@
     assert_equal etag, response.to_a[1]['ETag']
   end
 
+  it 'has a content-type method' do
+    response = Rack::Response.new
+    content_type = 'foo'
+    response.content_type = content_type
+    assert_equal content_type, response.content_type
+    assert_equal content_type, response.to_a[1]['Content-Type']
+  end
+
   it "have sensible default values" do
     response = Rack::Response.new
     status, header, body = response.finish
@@ -609,6 +617,31 @@ def obj.each
 
     res.finish.flatten.must_be_kind_of(Array)
   end
+
+  it "should specify not to cache content" do
+    response = Rack::Response.new
+
+    response.cache!(1000)
+    response.do_not_cache!
+
+    expect(response['Cache-Control']).must_equal "no-cache, must-revalidate"
+
+    expires_header = Time.parse(response['Expires'])
+    expect(expires_header).must_be :<=, Time.now
+  end
+
+  it "should specify to cache content" do
+    response = Rack::Response.new
+
+    duration = 120
+    expires = Time.now + 100 # At least this far into the future
+    response.cache!(duration)
+
+    expect(response['Cache-Control']).must_equal "public, max-age=120"
+
+    expires_header = Time.parse(response['Expires'])
+    expect(expires_header).must_be :>=, expires
+  end
 end
 
 describe Rack::Response, 'headers' do