lib/http/protocol/http1/connection.rb in http-protocol-0.16.0 vs lib/http/protocol/http1/connection.rb in http-protocol-0.17.0

@@ -30,22 +30,32 @@
 				
 				def initialize(stream, persistent = true)
 					@stream = stream
 					
 					@persistent = persistent
+					@upgrade = nil
 					
 					@count = 0
 				end
 				
 				attr :stream
 				
 				# Whether the connection is persistent.
 				attr :persistent
 				
+				# Whether the connection has been upgraded, and to what.
+				attr :upgrade
+				
 				# The number of requests processed.
 				attr :count
 				
+				def upgrade?(headers)
+					if upgrade = headers[UPGRADE]
+						return upgrade.first
+					end
+				end
+				
 				def persistent?(version, headers)
 					if version == HTTP10
 						if connection = headers[CONNECTION]
 							return connection.include?(KEEP_ALIVE)
 						else
@@ -58,22 +68,33 @@
 							return true
 						end
 					end
 				end
 				
+				def upgrade!(protocol)
+					@upgrade = protocol
+					@persistent = false
+					
+					return @stream
+				end
+				
 				# Write the appropriate header for connection persistence.
-				def write_persistent_header(version)
-					if version == HTTP10
-						@stream.write("connection: keep-alive\r\n") if @persistent
+				def write_connection_header(version)
+					if @upgrade
+						@stream.write("connection: upgrade\r\nupgrade: #{@upgrade}\r\n")
 					else
-						@stream.write("connection: close\r\n") unless @persistent
+						if version == HTTP10
+							@stream.write("connection: keep-alive\r\n") if @persistent
+						else
+							@stream.write("connection: close\r\n") unless @persistent
+						end
 					end
 				end
 				
 				# Effectively close the connection and return the underlying IO.
 				# @return [IO] the underlying non-blocking IO.
-				def hijack
+				def hijack!
 					@persistent = false
 					
 					@stream.flush
 					
 					return @stream
@@ -87,18 +108,18 @@
 				def write_request(authority, method, path, version, headers)
 					@stream.write("#{method} #{path} #{version}\r\n")
 					@stream.write("host: #{authority}\r\n")
 					
 					write_headers(headers)
-					write_persistent_header(version)
+					write_connection_header(version)
 				end
 				
 				def write_response(version, status, headers, body = nil, head = false)
 					@stream.write("#{version} #{status}\r\n")
 					
 					write_headers(headers)
-					write_persistent_header(version)
+					write_connection_header(version)
 					write_body(body, version == HTTP11, head)
 				end
 				
 				def write_headers(headers)
 					headers.each do |name, value|
@@ -111,24 +132,19 @@
 						yield line
 					end
 				end
 				
 				def read_line
-					# To support Ruby 2.3, we do the following which is pretty inefficient. Ruby 2.4+ can do the following:
-					# @stream.gets(CRLF, chomp: true) or raise EOFError
-					if line = @stream.gets(CRLF)
-						return line.chomp!(CRLF)
-					else
-						raise EOFError
-					end
+					@stream.gets(CRLF, chomp: true) or raise EOFError
 				end
 				
 				def read_request
 					method, path, version = read_line.split(/\s+/, 3)
 					headers = read_headers
 					
 					@persistent = persistent?(version, headers)
+					@upgrade = upgrade?(headers)
 					
 					body = read_request_body(headers)
 					
 					@count += 1
 					
@@ -181,10 +197,23 @@
 					@stream.read(2)
 					
 					return chunk
 				end
 				
+				def write_upgrade_body
+					@stream.write("\r\n")
+					@stream.flush
+					
+					return @stream unless block_given?
+					
+					begin
+						yield @stream
+					rescue
+						@stream.close_write
+					end
+				end
+				
 				def write_empty_body(body)
 					@stream.write("content-length: 0\r\n\r\n")
 					@stream.flush
 					
 					if body
@@ -262,9 +291,11 @@
 				end
 				
 				def write_body(body, chunked = true, head = false)
 					if body.nil? or body.empty?
 						write_empty_body(body)
+					# elsif body.respond_to?(:call)
+					# 	write_upgrade_body(&body)
 					elsif length = body.length
 						write_fixed_length_body(body, length, head)
 					elsif @persistent and chunked
 						# We specifically ensure that non-persistent connections do not use chunked response, so that hijacking works as expected.
 						write_chunked_body(body, head)