module Conjur
  module WebServer
    require 'rack/streaming_proxy'
    
    class APIProxy < Rack::StreamingProxy::Proxy

      class Request < Rack::StreamingProxy::Request
        def initialize env
          path = env["PATH_INFO"]

          path =~ /^\/([^\/]+)(.*)/
          app = $1
          path_remainder = $2

          new_url = case app
          when 'authn', 'authz', 'audit', 'pubkeys'
            [ Conjur.configuration.send("#{app}_url"), path_remainder ].join
          else
            [ Conjur.configuration.send("core_url"), path ].join
          end
          if query = env["QUERY_STRING"]
            new_url = [ new_url, query ].join('?')
          end

          super new_url.to_s, Rack::Request.new(env)
        end
      end

      Rack::StreamingProxy::Proxy.set_default_configuration

      def initialize
        super nil
      end

      def call env
        request = Request.new(env)
        request.http_request['Authorization'] = authorization_header
        response = Rack::StreamingProxy::Session.new(request).start
        rewrite_response(env, response.status, response.headers, response)
      end

      def rewrite_response(*args)
        env, status, headers, body = args
        
        source_request = Rack::Request.new(env)

        headers = Hash[*headers.flat_map { |k, v| [capitalize_header(k), v] }]
        headers.delete 'Transfer-Encoding' # let Puma handle chunking
    
        # Rewrite location
        if location = headers["Location"]
          headers["Location"] = location.gsub(Conjur.configuration.service_url, "http://#{source_request.host}:#{source_request.port}")
        end
        
        [ status, headers, body ]
      end
      
      protected
      
      def authorization_header
        require 'conjur/authn'
        require 'base64'
        token = Conjur::Authn.authenticate
        "Token token=\"#{Base64.strict_encode64(token.to_json)}\""
      end
      
      def perform_request(env)
        triplet = super(env)
        [ env ] + triplet
      end

      private

      def capitalize_header hdr
        hdr.split('-').map(&:capitalize).join('-')
      end
    end
  end
end

# Rack::StreamingProxy by default doesn't handle closing, which leads to stale
# (but still running) connections. Handle the close by quitting the child
# (it will close connection with upstream automatically).
class Rack::StreamingProxy::Response
  def initialize piper
    @piper = piper
    @client_http_version = '1.0'
    receive
  end

  def close
    @piper.signal :QUIT
  end
end