WebSockets

src/lucky/routable.cr

@@ -58,6 +58,17 @@ module Lucky::Routable
     end
   {% end %}
 
+  # Define a route that responds to a WebSocket request
+  macro ws(path, &block)
+    {% unless path.starts_with?("/") %}
+      {% path.raise "Path must start with a slash. Example: '/#{path}'" %}
+    {% end %}
+
+    add_route(:ws, {{ path }}, {{ @type.name.id }})
+
+    setup_ws_call_method(block)
+  end
+
   # Define a route with a custom HTTP method.
   #
   # Use this method if you need to match a route with a custom HTTP method (verb).
@@ -85,6 +96,35 @@ module Lucky::Routable
     setup_call_method({{ yield }})
   end
 
+  # :nodoc:
+  macro setup_ws_call_method(&block)
+
+    abstract def on_message(message)
+    abstract def on_close
+
+    def call
+      # Ensure clients_desired_format is cached by calling it
+      clients_desired_format
+
+      %pipe_result = run_before_pipes
+
+      %response = if %pipe_result.is_a?(Lucky::Response)
+        %pipe_result
+      else
+        {{ block.body }}
+        send_text_response "", content_type: "plain/text"
+      end
+
+      %pipe_result = run_after_pipes
+
+      if %pipe_result.is_a?(Lucky::Response)
+        %pipe_result
+      else
+        %response
+      end
+    end
+  end
+
   # :nodoc:
   macro setup_call_method(body)
     def call

src/lucky/web_socket_action.cr

@@ -0,0 +1,20 @@
+abstract class Lucky::WebSocketAction < Lucky::Action
+  @socket : HTTP::WebSocket?
+  @handler : HTTP::WebSocketHandler
+
+  def initialize(@context : HTTP::Server::Context, @route_params : Hash(String, String))
+    @handler = HTTP::WebSocketHandler.new do |ws|
+      @socket = ws
+      ws.on_ping { ws.pong("PONG") }
+      call
+    end
+  end
+
+  def perform_websocket_action
+    @handler.call(@context)
+  end
+
+  def socket : HTTP::WebSocket
+    @socket.not_nil!
+  end
+end

src/lucky/web_socket_route_handler.cr

@@ -0,0 +1,26 @@
+module Lucky
+  class WebSocketRouteHandler
+    include HTTP::Handler
+
+    def call(context : HTTP::Server::Context)
+      if websocket_upgrade_request?(context)
+        handler = Lucky::Router.find_action(:ws, context.request.path)
+        if handler
+          Lucky::Log.dexter.debug { {handled_by: handler.payload.to_s} }
+          handler.payload.new(context, handler.params).as(Lucky::WebSocketAction).perform_websocket_action
+        else
+          call_next(context)
+        end
+      else
+        call_next(context)
+      end
+    end
+
+    private def websocket_upgrade_request?(context)
+      return unless upgrade = context.request.headers["Upgrade"]?
+      return unless upgrade.compare("websocket", case_insensitive: true) == 0
+
+      context.request.headers.includes_word?("Connection", "Upgrade")
+    end
+  end
+end