diff --git a/lib/webrick/httpresponse.rb b/lib/webrick/httpresponse.rb index eae14d6597..2da1625fba 100644 --- a/lib/webrick/httpresponse.rb +++ b/lib/webrick/httpresponse.rb @@ -303,6 +303,8 @@ module WEBrick def send_body(socket) # :nodoc: if @body.respond_to? :readpartial then send_body_io(socket) + elsif @body.respond_to?(:call) then + send_body_proc(socket) else send_body_string(socket) end @@ -440,6 +442,39 @@ module WEBrick end end + def send_body_proc(socket) + if @request_method == "HEAD" + # do nothing + elsif chunked? + @body.call(ChunkedWrapper.new(socket, self)) + _write_data(socket, "0#{CRLF}#{CRLF}") + else + size = @header['content-length'].to_i + @body.call(socket) + @sent_size = size + end + end + + class ChunkedWrapper + def initialize(socket, resp) + @socket = socket + @resp = resp + end + + def write(buf) + return if buf.empty? + socket = @socket + @resp.instance_eval { + size = buf.bytesize + data = "#{size.to_s(16)}#{CRLF}#{buf}#{CRLF}" + _write_data(socket, data) + data.clear + @sent_size += size + } + end + alias :<< :write + end + def _send_file(output, input, offset, size) while offset > 0 sz = @buffer_size < size ? @buffer_size : size diff --git a/test/webrick/test_httpresponse.rb b/test/webrick/test_httpresponse.rb index d97a65b5db..b06cf925a2 100644 --- a/test/webrick/test_httpresponse.rb +++ b/test/webrick/test_httpresponse.rb @@ -147,6 +147,29 @@ module WEBrick assert_equal 0, logger.messages.length end + def test_send_body_proc + @res.body = Proc.new { |out| out.write('hello') } + IO.pipe do |r, w| + @res.send_body(w) + w.close + r.binmode + assert_equal 'hello', r.read + end + assert_equal 0, logger.messages.length + end + + def test_send_body_proc_chunked + @res.body = Proc.new { |out| out.write('hello') } + @res.chunked = true + IO.pipe do |r, w| + @res.send_body(w) + w.close + r.binmode + assert_equal "5\r\nhello\r\n0\r\n\r\n", r.read + end + assert_equal 0, logger.messages.length + end + def test_set_error status = 400 message = 'missing attribute'