Better support for connection upgrade and bi-directional streaming. (#101)

Author: ioquatix

lib/webrick/httpresponse.rb

@@ -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.
@@ -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+
 
@@ -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")
@@ -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
 
@@ -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
 

test/webrick/test_httpresponse.rb

@@ -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'

test/webrick/test_httpserver.rb

@@ -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|