require 'rack'
require 'rack/request'

module Rack
  class HSTS
    def initialize(app, options = {})
      @app, @options = app, { :expires => 600, :port => 443 }.merge(options)
      @hsts = { 'Strict-Transport-Security' => "max-age=#{@options[:expires]}" }
    end

    def call(env)
      if env['rack.url_scheme'] == 'http'
        env['rack.url_scheme'] = 'https'
        env['SERVER_PORT'] = "#{@options[:port]}"
        [301, { 'Location' => Request.new(env).url }, []]
      else
        status, headers, body = @app.call(env)
        headers = @hsts.merge(headers) if env['rack.url_scheme'] == 'https'
        [status, headers, body]
      end
    end
  end
end