Seperate websocket and websocket handler. Fixes #395

spec/config_spec.cr

@@ -27,7 +27,7 @@ describe "Config" do
     config = Kemal.config
     config.add_handler CustomTestHandler.new
     Kemal.config.setup
-    config.handlers.size.should eq(6)
+    config.handlers.size.should eq(7)
   end
 
   it "toggles the shutdown message" do

spec/helpers_spec.cr

@@ -12,7 +12,7 @@ describe "Macros" do
     it "adds a custom handler" do
       add_handler CustomTestHandler.new
       Kemal.config.setup
-      Kemal.config.handlers.size.should eq 6
+      Kemal.config.handlers.size.should eq 7
     end
   end
 

spec/spec_helper.cr

@@ -48,8 +48,7 @@ def create_ws_request_and_return_io(handler, request)
   rescue IO::Error
     # Raises because the IO::Memory is empty
   end
-  response.close
-  io
+  io.rewind
 end
 
 def call_request_on_app(request)
@@ -83,5 +82,5 @@ end
 Spec.after_each do
   Kemal.config.clear
   Kemal::RouteHandler::INSTANCE.http_routes = Radix::Tree(Route).new
-  Kemal::RouteHandler::INSTANCE.ws_routes = Radix::Tree(String).new
+  Kemal::WebSocketHandler::INSTANCE.ws_routes = Radix::Tree(WebSocket).new
 end

spec/websocket_handler_spec.cr

@@ -2,32 +2,47 @@ require "./spec_helper"
 
 describe "Kemal::WebSocketHandler" do
   it "doesn't match on wrong route" do
-    handler = Kemal::WebSocketHandler.new "/" { }
+    handler = Kemal::WebSocketHandler::INSTANCE
+    handler.next = Kemal::CommonExceptionHandler::INSTANCE
+    ws "/" { }
     headers = HTTP::Headers{
       "Upgrade"           => "websocket",
       "Connection"        => "Upgrade",
       "Sec-WebSocket-Key" => "dGhlIHNhbXBsZSBub25jZQ==",
     }
     request = HTTP::Request.new("GET", "/asd", headers)
-    io_with_context = create_request_and_return_io(handler, request)
-    client_response = HTTP::Client::Response.from_io(io_with_context, decompress: false)
+    io = IO::Memory.new
+    response = HTTP::Server::Response.new(io)
+    context = HTTP::Server::Context.new(request, response)
+    begin
+      handler.call context
+    rescue IO::Error
+      # Raises because the IO::Memory is empty
+    end
+    response.close
+    io.rewind
+    client_response = HTTP::Client::Response.from_io(io, decompress: false)
     client_response.status_code.should eq(404)
   end
 
   it "matches on given route" do
-    handler = Kemal::WebSocketHandler.new "/" { }
+    handler = Kemal::WebSocketHandler::INSTANCE
+    ws "/" { |socket, context| socket.send("Match") }
+    ws "/no_match" { |socket, context| socket.send "No Match" }
     headers = HTTP::Headers{
       "Upgrade"           => "websocket",
       "Connection"        => "Upgrade",
       "Sec-WebSocket-Key" => "dGhlIHNhbXBsZSBub25jZQ==",
     }
     request = HTTP::Request.new("GET", "/", headers)
+
     io_with_context = create_ws_request_and_return_io(handler, request)
-    io_with_context.to_s.should eq("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-Websocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\r\n")
+    io_with_context.to_s.should eq("HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-Websocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\r\n\x81\u0005Match")
   end
 
   it "fetches named url parameters" do
-    handler = Kemal::WebSocketHandler.new "/:id" { |s, c| c.params.url["id"] }
+    handler = Kemal::WebSocketHandler::INSTANCE
+    ws "/:id" { |s, c| c.params.url["id"] }
     headers = HTTP::Headers{
       "Upgrade"           => "websocket",
       "Connection"        => "Upgrade",

src/kemal/config.cr

@@ -109,9 +109,9 @@ module Kemal
         setup_static_file_handler
         setup_custom_handlers
         setup_filter_handlers
-        setup_websocket_handlers
         @default_handlers_setup = true
         @router_included = true
+        HANDLERS.insert(HANDLERS.size, Kemal::WebSocketHandler::INSTANCE)
         HANDLERS.insert(HANDLERS.size, Kemal::RouteHandler::INSTANCE)
       end
     end

src/kemal/dsl.cr

@@ -16,7 +16,7 @@ FILTER_METHODS = %w(get post put patch delete options all)
 
 def ws(path : String, &block : HTTP::WebSocket, HTTP::Server::Context -> Void)
   raise Kemal::Exceptions::InvalidPathStartException.new("ws", path) unless Kemal::Utils.path_starts_with_slash?(path)
-  Kemal::WebSocketHandler.new path, &block
+  Kemal::WebSocketHandler::INSTANCE.add_ws_route path, &block
 end
 
 def error(status_code : Int32, &block : HTTP::Server::Context, Exception -> _)

src/kemal/ext/context.cr

@@ -13,13 +13,7 @@ class HTTP::Server
     end
 
     def params
-      connection_type = @request.headers.fetch("Connection", nil)
-      @request.url_params ||= unless connection_type == "Upgrade"
-        route_lookup.params
-      else
-        ws_route_lookup.params
-      end
-
+      @request.url_params ||= route_lookup.params
       @params ||= if @request.param_parser
                     @request.param_parser.not_nil!
                   else
@@ -36,6 +30,10 @@ class HTTP::Server
       route_lookup.payload
     end
 
+    def websocket
+      ws_route_lookup.payload
+    end
+
     def route_lookup
       Kemal::RouteHandler::INSTANCE.lookup_route(@request.override_method.as(String), @request.path)
     end
@@ -45,7 +43,7 @@ class HTTP::Server
     end
 
     def ws_route_lookup
-      Kemal::RouteHandler::INSTANCE.lookup_ws_route(@request.path)
+      Kemal::WebSocketHandler::INSTANCE.lookup_ws_route(@request.path)
     end
 
     def ws_route_defined?

src/kemal/route_handler.cr

@@ -26,19 +26,11 @@ module Kemal
       add_to_http_radix_tree("HEAD", path, Route.new("HEAD", path) { |ctx| "" }) if method == "GET"
     end
 
-    def add_ws_route(path : String)
-      add_to_ws_radix_tree path
-    end
-
     # Check if a route is defined and returns the lookup
     def lookup_route(verb : String, path : String)
       @http_routes.find radix_path(verb, path)
     end
 
-    def lookup_ws_route(path : String)
-      @ws_routes.find "/ws#{path}"
-    end
-
     # Processes the route if it's a match. Otherwise renders 404.
     private def process_request(context)
       raise Kemal::Exceptions::RouteNotFound.new(context) unless context.route_defined?
@@ -59,10 +51,5 @@ module Kemal
       node = radix_path method, path
       @http_routes.add node, route
     end
-
-    private def add_to_ws_radix_tree(path)
-      node = radix_path "ws", path
-      @ws_routes.add node, node
-    end
   end
 end

src/kemal/websocket.cr

@@ -0,0 +1,15 @@
+module Kemal
+  # Route is the main building block of Kemal.
+  # It takes 3 parameters: Method, path and a block to specify
+  # what action to be done if the route is matched.
+  class WebSocket < HTTP::WebSocketHandler
+    getter proc
+
+    def initialize(@path : String, &@proc : HTTP::WebSocket, HTTP::Server::Context -> Void)
+    end
+
+    def call(context : HTTP::Server::Context)
+      super
+    end
+  end
+end

src/kemal/websocket_handler.cr

@@ -1,15 +1,38 @@
 module Kemal
   # Kemal::WebSocketHandler is used for building a WebSocket route.
   # For each WebSocket route a new handler is created and registered to global handlers.
-  class WebSocketHandler < HTTP::WebSocketHandler
-    def initialize(@path : String, &@proc : HTTP::WebSocket, HTTP::Server::Context -> Void)
-      Kemal.config.add_handler self
-      Kemal::RouteHandler::INSTANCE.add_ws_route @path
+  class WebSocketHandler
+    include HTTP::Handler
+    INSTANCE = new
+    property ws_routes
+
+    def initialize
+      @ws_routes = Radix::Tree(WebSocket).new
     end
 
     def call(context : HTTP::Server::Context)
       return call_next(context) unless context.ws_route_defined?
-      super
+      context.request.url_params ||= context.ws_route_lookup.params
+      content = context.websocket.call(context)
+      context.response.print(content)
+      context
+    end
+
+    def lookup_ws_route(path : String)
+      @ws_routes.find "/ws#{path}"
+    end
+
+    def add_ws_route(path : String, &handler : HTTP::WebSocket, HTTP::Server::Context -> Void)
+      add_to_ws_radix_tree path, WebSocket.new(path, &handler)
+    end
+
+    private def add_to_ws_radix_tree(path, websocket)
+      node = radix_path "ws", path
+      @ws_routes.add node, websocket
+    end
+
+    private def radix_path(method, path)
+      "/#{method.downcase}#{path}"
     end
   end
 end