Skip to content

Commit

Permalink
Better support for connection upgrade and bi-directional streaming. (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
ioquatix authored Jan 27, 2023
1 parent 6cb9bf6 commit 1c9d2f4
Show file tree
Hide file tree
Showing 3 changed files with 64 additions and 8 deletions.
36 changes: 30 additions & 6 deletions lib/webrick/httpresponse.rb
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,11 @@ class InvalidHeader < StandardError

attr_reader :sent_size

##
# Set the response body proc as an streaming/upgrade response.

attr_accessor :upgrade

##
# Creates a new HTTP response object. WEBrick::Config::HTTP is the
# default configuration.
Expand Down Expand Up @@ -217,6 +222,16 @@ def keep_alive?
@keep_alive
end

##
# Sets the response to be a streaming/upgrade response.
# This will disable keep-alive and chunked transfer encoding.

def upgrade!(protocol)
@upgrade = protocol
@keep_alive = false
@chunked = false
end

##
# Sends the response on +socket+

Expand All @@ -242,6 +257,14 @@ def setup_header() # :nodoc:
@header['server'] ||= @config[:ServerSoftware]
@header['date'] ||= Time.now.httpdate

if @upgrade
@header['connection'] = 'upgrade'
@header['upgrade'] = @upgrade
@keep_alive = false

return
end

# HTTP/0.9 features
if @request_http_version < "1.0"
@http_version = HTTPVersion.new("0.9")
Expand All @@ -268,11 +291,10 @@ def setup_header() # :nodoc:
elsif %r{^multipart/byteranges} =~ @header['content-type']
@header.delete('content-length')
elsif @header['content-length'].nil?
if @body.respond_to? :readpartial
elsif @body.respond_to? :call
make_body_tempfile
if @body.respond_to?(:bytesize)
@header['content-length'] = @body.bytesize.to_s
else
@header['content-length'] = (@body ? @body.bytesize : 0).to_s
@header['connection'] = 'close'
end
end

Expand Down Expand Up @@ -517,14 +539,16 @@ def send_body_proc(socket)
@body.call(ChunkedWrapper.new(socket, self))
socket.write("0#{CRLF}#{CRLF}")
else
size = @header['content-length'].to_i
if @bodytempfile
@bodytempfile.rewind
IO.copy_stream(@bodytempfile, socket)
else
@body.call(socket)
end
@sent_size = size

if content_length = @header['content-length']
@sent_size = content_length.to_i
end
end
end

Expand Down
33 changes: 33 additions & 0 deletions test/webrick/test_httpresponse.rb
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,39 @@ def test_send_body_proc_chunked
assert_equal 0, logger.messages.length
end

def test_send_body_proc_upgrade
@res.body = Proc.new { |out| out.write('hello'); out.close }
@res.upgrade!("text")

IO.pipe do |r, w|
@res.send_response(w)
w.close
assert_match /Connection: upgrade\r\nUpgrade: text\r\n\r\nhello/, r.read
end
assert_empty logger.messages
end

def test_send_body_proc_stream
@res.body = Proc.new do |socket|
chunk = socket.read
socket.write(chunk)
socket.close
end

UNIXSocket.pair do |s1, s2|
thread = Thread.new do
@res.send_response(s1)
end

s2.write("hello")
s2.close_write
chunk = s2.read
assert_match /Connection: close\r\n\r\nhello/, chunk
s2.close
end
assert_empty logger.messages
end

def test_set_error
status = 400
message = 'missing attribute'
Expand Down
3 changes: 1 addition & 2 deletions test/webrick/test_httpserver.rb
Original file line number Diff line number Diff line change
Expand Up @@ -377,8 +377,7 @@ def test_response_io_without_chunked_set
:ServerName => "localhost"
}
log_tester = lambda {|log, access_log|
assert_equal(1, log.length)
assert_match(/WARN Could not determine content-length of response body./, log[0])
assert_empty log
}
TestWEBrick.start_httpserver(config, log_tester){|server, addr, port, log|
server.mount_proc("/", lambda { |req, res|
Expand Down

0 comments on commit 1c9d2f4

Please sign in to comment.