HTTP::Server: defer request upgrade (aka: WebSockets)

src/http/server.cr

@@ -511,6 +511,8 @@ class HTTP::Server
     {% end %}
 
     @processor.process(io, io)
+  ensure
+    io.close rescue IO::Error
   end
 
   # This method handles exceptions raised at `Socket#accept?`.

src/http/server/handlers/websocket_handler.cr

@@ -57,8 +57,6 @@ class HTTP::WebSocketHandler
       ws_session = WebSocket.new(io, sync_close: false)
       @proc.call(ws_session, context)
       ws_session.run
-    ensure
-      io.close
     end
   end
 

src/http/server/request_processor.cr

@@ -23,7 +23,6 @@ class HTTP::Server::RequestProcessor
   end
 
   def process(input, output)
-    must_close = true
     response = Response.new(output)
 
     begin
@@ -63,13 +62,15 @@ class HTTP::Server::RequestProcessor
           response.output.close
         end
 
-        if response.upgraded?
-          must_close = false
+        output.flush
+
+        # If there is an upgrade handler, hand over
+        # the connection to it and return
+        if upgrade_handler = response.upgrade_handler
+          upgrade_handler.call(output)
           return
         end
 
-        output.flush
-
         break unless request.keep_alive?
 
         # Don't continue if the handler set `Connection` header to `close`
@@ -93,12 +94,6 @@ class HTTP::Server::RequestProcessor
       end
     rescue IO::Error
       # IO-related error, nothing to do
-    ensure
-      begin
-        input.close if must_close
-      rescue IO::Error
-        # IO-related error, nothing to do
-      end
     end
   end
 end

src/http/server/response.cr

@@ -14,7 +14,6 @@ class HTTP::Server
   #
   # A response can be upgraded with the `upgrade` method. Once invoked, headers
   # are written and the connection `IO` (a socket) is yielded to the given block.
-  # The block must invoke `close` afterwards, the server won't do it in this case.
   # This is useful to implement protocol upgrades, such as websockets.
   class Response < IO
     # The response headers (`HTTP::Headers`). These must be set before writing to the response.
@@ -34,14 +33,16 @@ class HTTP::Server
     # body. If not set, the default value is 200 (OK).
     property status : HTTP::Status
 
+    # :nodoc:
+    property upgrade_handler : (IO ->)?
+
     @cookies : HTTP::Cookies?
 
     # :nodoc:
     def initialize(@io : IO, @version = "HTTP/1.1")
       @headers = Headers.new
       @status = :ok
       @wrote_headers = false
-      @upgraded = false
       @output = output = @original_output = Output.new(@io)
       output.response = self
     end
@@ -53,7 +54,6 @@ class HTTP::Server
       @cookies = nil
       @status = :ok
       @wrote_headers = false
-      @upgraded = false
       @output = @original_output
       @original_output.reset
     end
@@ -97,18 +97,10 @@ class HTTP::Server
     end
 
     # Upgrades this response, writing headers and yieling the connection `IO` (a socket) to the given block.
-    # The block must invoke `close` afterwards, the server won't do it in this case.
     # This is useful to implement protocol upgrades, such as websockets.
-    def upgrade
-      @upgraded = true
+    def upgrade(&block : IO ->)
       write_headers
-      flush
-      yield @io
-    end
-
-    # :nodoc:
-    def upgraded?
-      @upgraded
+      @upgrade_handler = block
     end
 
     # Flushes the output. This method must be implemented if wrapping the response output.