Revでhttpプロキシ

以前作っていたフィルタプロキシを今風に作り直そうと思って、
Rackで作ろうと思って調べ始めたのだが、
ストーリミングデータを扱えるような逐次送信のAPIがなさげだったので
なにか他のものを探していたら、
いつの間にかRevを使うことになっていた。

Revはlibevのrubyバインディングだそうで、たぶんlibevとかtwistedとかglibのloopみたいなイベント駆動のフレームワークなんだろう。
というわけで、Revの練習にHTTPプロキシを作ってみた。
まだ到底使えないが、とりあえずプロトタイプが動くところまでできた。

Revのおかげで、TCPServerとselectを使うところがわりとシンプルに書けていて、更に、データの逐次送信をやってもそれほど複雑化しない。
まだ、動きが怪しげなのだが、Webrickのプロキシあたりを参考にしてHTTPヘッダの取り回しを調べないと勘ではこれ以上進まないな。
あと、コネクションをクローズするタイミングがよくわからん。

#!/usr/bin/ruby

require 'rev'
require 'logger'
require 'stringio'
require 'strscan'
require 'uri'


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


class ProxyRequest
  attr_reader :uri

  def initialize
    @lines=['']
  end

  def complete?(data)
    last_byte=nil
    is_complete=false
    data.each_byte do |b|
      case b
      when 0x0d
      when 0x0a
        if last_byte==0x0d then
          if @lines.last=='' then
            is_complete=true
              parse
            break
          else
            @lines << ''
          end
        end
      else
        @lines.last << b
      end
      last_byte=b
    end
    return is_complete
  end

  def parse
    @method, @url, @version=@lines.shift.split
    @uri=URI.parse @url
    @headers={}
    @lines.each{|l|
      if l.size>0 then
        k, v=l.split(/: */, 2)
        @headers[k.downcase.to_sym]=v
      end
    }
    @lines=nil
  end

  def query_hash
    hash={}
    if @uri.query then
      StringScanner.new(@uri.query).scan(/(\w+)=(\w+)/) do |k, v|
        hash[k]=v
      end
    end
    hash
  end
end


class ProxyDestination < Rev::HttpClient
  def set_client(client)
    @client=client
  end

  def on_response_header(header)
    logger.info "on_response_header: #{header.chunked_encoding?}"
    @client.write "HTTP/#{header.http_version} #{header.http_status} #{header.http_reason}"
    header.each do |k, v|
      @client.write "#{k}: #{v}\r\n"
    end
    @client.write "\r\n"
  end

  def on_body_data(data)
    logger.info "on_body_data: #{data.bytesize}"
    @client.write data
  end

  def set_complete_callback(&block)
    @complete_callback=block
  end

  def on_close
    logger.info "on_close"
    @complete_callback.call
  end
end


class ProxyClient < Rev::TCPSocket
  def on_read(data)
    logger.info "ProxyClient::on_read"
    @request||=ProxyRequest.new
    if @request.complete? data then
      uri=@request.uri
      logger.info(uri)

      c=ProxyDestination.connect(uri.host, uri.port).attach(Rev::Loop.default)
      c.set_client(self)
      c.request(:get, uri.path, :query=>@request.query_hash, :head=>@headers)
      c.set_complete_callback do
        logger.info "on_complete"
        @request=nil
        @is_complete=true
        if output_buffer_size==0 then
          close
        end
      end
    end
  end

  def on_write_complete(*args)
    if @is_complete then
      logger.info "complete"
      close
    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, ProxyClient)
  server.attach(Rev::Loop.default)

  Rev::Loop.default.run
end


if $0==__FILE__ then
  s=start
end