Rack::ETag now supports SHA-1 and SHA-2 in addition to MD5

lib/rack/contrib/etag.rb

@@ -1,20 +1,61 @@
-require 'digest/md5'
-
 module Rack
   # Automatically sets the ETag header on all String bodies
+  # if none is set.
+  #
+  # By default, uses an MD5 hash to generate the ETag.
+  #
+  # @param [#call] app the underlying Rack application. Required.
+  # @param [Symbol] digest the digest to use. Optional. Options are [:md5, :sha1, :sha256, :sha384, :sha512].
   class ETag
-    def initialize(app)
+    def initialize(app, digest = :md5)
       @app = app
+      load_digest(digest)
+      @digest_method = self.method("#{digest}_hash")
     end
 
     def call(env)
       status, headers, body = @app.call(env)
 
       if !headers.has_key?('ETag') && body.is_a?(String)
-        headers['ETag'] = %("#{Digest::MD5.hexdigest(body)}")
+        headers['ETag'] = %("#{@digest_method.call(body)}")
       end
 
       [status, headers, body]
     end
+
+    private
+
+    def load_digest(digest)
+      case digest
+      when :md5
+        require 'digest/md5'
+      when :sha1
+        require 'digest/sha1'
+      when :sha256, :sha384, :sha512
+        require 'digest/sha2'
+      else
+        raise ArgumentError.new("Digest #{digest} is not supported.")
+      end
+    end
+
+    def md5_hash(body)
+      Digest::MD5.hexdigest(body)
+    end
+
+    def sha1_hash(body)
+      Digest::SHA1.hexdigest(body)
+    end
+
+    def sha256_hash(body)
+      Digest::SHA2.hexdigest(body)
+    end
+
+    def sha384_hash(body)
+      (Digest::SHA2.new(384) << body).to_s
+    end
+
+    def sha512_hash(body)
+      (Digest::SHA2.new(512) << body).to_s
+    end
   end
 end

test/spec_rack_etag.rb

@@ -1,12 +1,38 @@
 require 'test/spec'
 require 'rack/mock'
 require 'rack/contrib/etag'
+require 'digest/md5'
+require 'digest/sha1'
+require 'digest/sha2'
 
 context "Rack::ETag" do
-  specify "sets ETag if none is set" do
-    app = lambda { |env| [200, {'Content-Type' => 'text/plain'}, "Hello, World!"] }
-    response = Rack::ETag.new(app).call({})
-    response[1]['ETag'].should.equal "\"65a8e27d8879283831b664bd8b7f0ad4\""
+
+  body   = 'Hello, World!'
+
+  context('if no ETag is set on a String body') do
+    before(:each) do
+      @app = lambda { |env| [200, {'Content-Type' => 'text/plain'}, body] }
+    end
+    specify 'sets ETag' do
+      response = Rack::ETag.new(@app).call({})
+      response[1]['ETag'].should.equal "\"#{Digest::MD5.hexdigest(body)}\""
+    end
+    specify 'uses SHA-1 if specified' do
+      response = Rack::ETag.new(@app, :sha1).call({})
+      response[1]['ETag'].should.equal "\"#{Digest::SHA1.hexdigest(body)}\""
+    end
+    specify 'uses SHA-256 if specified' do
+      response = Rack::ETag.new(@app, :sha256).call({})
+      response[1]['ETag'].should.equal "\"#{Digest::SHA2.hexdigest(body)}\""
+    end
+    specify 'uses SHA-384 if specified' do
+      response = Rack::ETag.new(@app, :sha384).call({})
+      response[1]['ETag'].should.equal "\"#{(Digest::SHA2.new(384) << body).to_s}\""
+    end
+    specify 'uses SHA-512 if specified' do
+      response = Rack::ETag.new(@app, :sha512).call({})
+      response[1]['ETag'].should.equal "\"#{(Digest::SHA2.new(512) << body).to_s}\""
+    end
   end
 
   specify "does not change ETag if it is already set" do
@@ -15,7 +41,7 @@
     response[1]['ETag'].should.equal "\"abc\""
   end
 
-  specify "does not set ETag if steaming body" do
+  specify "does not set ETag if streaming body" do
     app = lambda { |env| [200, {'Content-Type' => 'text/plain'}, ["Hello", "World"]] }
     response = Rack::ETag.new(app).call({})
     response[1]['ETag'].should.equal nil