Seperate websocket and websocket handler. Fixes #395

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