Convenient cache and content type methods for `Rack::Response`.

3 files changed, 56 insertions, 0 deletions

diff --git a/lib/rack.rb b/lib/rack.rb
index 7600c40f..cab2bb8b 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 408a38fc..fbbcb03e 100644
--- a/lib/rack/response.rb
+++ b/lib/rack/response.rb
@@ -178,10 +178,16 @@ module Rack
         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 @@ module Rack
         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 b2ba59a8..1dfafcdb 100644
--- a/test/spec_response.rb
+++ b/test/spec_response.rb
@@ -30,6 +30,14 @@ describe Rack::Response do
     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 @@ describe Rack::Response do
 
     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