WEBrick::HTTPProxyServerでデータを逐次送信させる

HTTPProxyServerはデータをすべて取得してから一度に返す動きをしていたのでこれを逐次で返すように改造してみた。
ニコ串のソースを見てやり方はすぐ判ったのだが、これをWEBrickの枠に当てはめるのにだいぶ苦労した。
なんか動作が怪しいところもあるがyoutubeとかでデータがちょっとずつ来るのは確認できた。

#!/usr/bin/env ruby
require 'webrick'
require 'webrick/httpproxy'
require 'net/http'

module WEBrick
  # 返送用のソケットにアクセスできるようにする
  class HTTPRequest
    attr_reader :socket
  end

  class SequencialProxy < HTTPProxyServer
    def initialize(config={}, default=Config::HTTP)
      super(config, default)
      @http_version="1.0"
    end

    def proxy_service(req, res)
      # Proxy Authentication
      proxy_auth(req, res)      
      begin
        self.send("do_#{req.request_method}", req, res)
      rescue NoMethodError
        raise HTTPStatus::MethodNotAllowed,
          "unsupported method `#{req.request_method}'."
      rescue => err
p err.backtrace
        logger.debug("#{err.class}: #{err.message}")
        raise HTTPStatus::ServiceUnavailable, err.message
      end

      # Process contents
      if handler = @config[:ProxyContentHandler]
        handler.call(req, res)
      end
    end

    # とりあえずゲットだけ
    def do_GET(req, res)
      uri = req.request_uri
      path = uri.path.dup
      path << "?" << uri.query if uri.query
      header = setup_proxy_header(req, res)

      # send request
      # Net::HTTPだと逐次にできないのでTCPSocketでやる
      server=Net::BufferedIO.new(
        TCPSocket.new(req.request_uri.host, req.request_uri.port))
      server.writeline "GET #{path} HTTP/1.0"
      header.each do |k, v|
        server.writeline "#{k}: #{v}"
      end
      server.writeline ""

      # 逐次で返すため先にヘッダを処理する
      # read header 
      response=Net::HTTPResponse.read_new(server)

      # Convert Net::HTTP::HTTPResponse to WEBrick::HTTPResponse
      res.status = response.code.to_i
      choose_header(response, res)
      set_cookie(response, res)
      set_via(res)

      # Persistent connection requirements are mysterious for me.
      # So I will close the connection in every response.
      res['proxy-connection'] = "close"
      res['connection'] = "close"

      # header
      res.send_header(req.socket)

      # read body
      res.body = ''
      while true  do
        begin
          block=''
          server.read(4096, block)
        rescue EOFError => e
          break
        ensure
          res.body << block
          # 逐次
          req.socket << block

          if block=='' then
            break
          end
        end
      end

      # 既にデータは返したので元の場所では何もしない
      def res.send_message(socket)
        # dummy
      end
    end
  end
end

# プロキシサーバオブジェクトを作る
s = WEBrick::SequencialProxy.new(
  :BindAddress => '127.0.0.1',
  :Port => 18080,
  :ProxyVia => false,

  # WEBrick::GenericServer callback
  :StartCallback => 
  ->{
    #p "StartCallback"
  },

  :StopCallback =>
  ->{
    #p "StopCallback"
  },

    :AcceptCallback =>
  -> sock {
    #p "AcceptCallback", sock
  },

  # WEBrick::HTTPServer callback
  :RequestCallback =>
  -> req, res {
    #p "RequestCallback"
  },
)

Signal.trap('INT') do
  s.shutdown
end

s.start

現状だとProxyContentHandlerが無視されるのでまだ改造が必要だがmimeを見て動作を振り分ける方がよさげ。

参考

http://jp.rubyist.net/magazine/?0002-WEBrickProxy
http://d.hatena.ne.jp/wocota/20090419/1240131284