require 'digest/md5'
require 'digest/sha1'

module Rack
  module Contrib
    #
    # Add Content-MD5 and Content-SHA1 headers, containing the hashes of the
    # body's response.
    #
    # Best installed as your very last middleware piece to ensure the contents
    # doesn't become invalidated by another middleware piece.
    #
    # Ex:
    #     use Rack::Contrib::ContentHash, md5: true, sha1: true
    #
    class ContentHash
      VERSION = '0.0.1'

      # Initialize the middleware, pass a Hash of algorithms in the second
      # parameter. Ex: .new app, md5: true, sha1: false
      def initialize app, algorithms = {}
        @app = app
        @algo= algorithms
      end

      # Hash the body of your response, and append the headers to the Hash
      def call env
        status, headers, body = @app.call env

        headers['Content-MD5'] = Digest::MD5.hexdigest(body) if @algo[:md5]
        headers['Content-SHA1'] = Digest::SHA1.hexdigest(body) if @algo[:sha1]

        [status, headers, body]
      end
    end
  end
end