Revでhttpプロキシその２

リファクタリングして、Postメソッドへの対応を実装した。
Proxy用途には、Rev::HttpClientは都合が悪いことが判ったので、素のRev::TCPSocketを使うように変えた。
keep-aliveとssl対応をしてからコンテンツフィルターとキャッシングに使いたいのだがいけるかな。

#!/usr/bin/ruby
require 'rev'
require 'logger'
require 'stringio'
require 'strscan'
require 'uri'
require 'socket'


def logger
  $logger||=Logger.new(STDOUT)
end


module Proxy
  class HTTPRequest
    attr_reader :method, :uri, :major, :minor, :headers, :body
    def initialize(method, uri, major, minor, headers)
      @method=method.upcase.to_sym
      @uri=uri
      @major=major
      @minor=minor
      @headers=headers
      @body=''
    end

    def format
      http="#@method #@uri HTTP/#@major.#@minor\r\n"
      # keep-aliveしない
      @headers['connection']='close'
      @headers.each do |k, v|
        case k
        when /^proxy-/, 'keep-alive'
          # proxyとkeep-aliveのヘッダとを落とす
        else
          http << "#{k}: #{v}\r\n"
        end
      end
      http << "\r\n"
    end

    def body_is_end?
      if @headers.key? 'content-length' then
        if @headers['content-length'].to_i==@body.bytesize then
          return true
        else
          return false
        end
      else
        if @body.end_with? "\r\n" then
          return true
        else
          return false
        end
      end
    end
  end


  class RequestBuilder
    attr_reader :request
    def initialize
      @lines=['']
      @request=nil
    end

    def push(data)
      last_byte=nil
      data.each_byte do |b|
        if @request then
          # request body
          @request.body << b
        else
          # request header
          case b
          when 0x0d
          when 0x0a
            if last_byte==0x0d then
              if @lines.last=='' then
                # request complete
                build_header
              else
                # next line
                @lines << ''
              end
            end
          else
            @lines.last << b
          end
          last_byte=b
        end
      end
    end

    def build_header
      # HTTP Request 1st line
      method, url, major, minor=@lines.shift.match(
        %r|(\S+)\s+(\S+)\s+HTTP/(\d+)\.(\d+)|).captures
      # HTTP request headers
      headers={}
      @lines.each{|l|
        if l.size>0 then
          k, v=l.split(/: */, 2)
          headers[k.downcase]=v
        end
      }
      # create HTTP request
      @request=HTTPRequest.new(method, URI.parse(url), 
                      major.to_i, minor.to_i, headers)
    end
  end


  class UpStream < Rev::TCPSocket
    def initialize(uri, down_stream)
      super TCPSocket.new(uri.host, uri.port)
      @down_stream=down_stream
    end

    def on_read(data)
      logger.info "UpStream::on_read #{data.bytesize} bytes"
      @down_stream.write data
    end

    def on_close
      logger.info "UpStream::on_close"
      @down_stream.close_on_write_complete
    end
  end

  class DownStream < Rev::TCPSocket
    def initialize(socket)
      super
      @requests=[]
    end

    def on_read(data)
      logger.info "DownStream::on_read #{data.bytesize} bytes"
      @builder||=RequestBuilder.new
      @builder.push data

      if @builder.request then
        case @builder.request.method
        when :GET, :HEAD
          proxy_request
        when :POST
          if @builder.request.body_is_end? then
            proxy_request
          else
            logger.error @builder.request
          end
        else
          logger.error "unknown method: #{@builder.header.method}"
        end
      end
    end

    def proxy_request
      request=@builder.request
      @requests << request
      @builder=nil

      uri=request.uri
      c=Proxy::UpStream.new(uri, self)
      c.write request.format
      if request.method==:POST then
        c.write request.body
      end
      logger.info request.format
      c.attach(Rev::Loop.default)
    end

    def on_write_compilete
      if @should_close then
        close
      end
    end

    def close_on_write_complete
      if output_buffer_size==0 then
        close
      else
        @should_close=true
      end
    end
  end
end


def start(listen=0, port=8080)
  logger.info("start.")
  logger.info("listen: #{listen}.")
  logger.info("port: #{port}.")

  server = Rev::TCPServer.new(listen, port, Proxy::DownStream)
  server.attach(Rev::Loop.default)

  Rev::Loop.default.run
end


if $0==__FILE__ then
  start
end