WebSockets

src/lucky/routable.cr

@@ -58,6 +58,21 @@ 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 %}
+
+    {% if block.args.size != 1 %}
+      {% raise "ws takes a block with 1 arg." %}
+    {% end %}
+
+    add_route({{method}}, {{ 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 +100,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(socket : Lucky::WebSocket)
+      # 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.args.first }} = socket
+        {{ block.body }}
+      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.cr

@@ -0,0 +1,8 @@
+module Lucky
+  class WebSocket < HTTP::WebSocketHandler
+    getter proc
+
+    def initialize(@action : Lucky::Action.class, &@proc : HTTP::WebSocket, HTTP::Server::Context -> Void)
+    end
+  end
+end

src/lucky/web_socket_action.cr

@@ -0,0 +1,19 @@
+abstract class Lucky::WebSocketAction < Lucky::Action
+  getter websocket : Lucky::WebSocket
+
+  def initialize(@context : HTTP::Server::Context, @route_params : Hash(String, String))
+    @websocket = Lucky::WebSocket.new(self.class) do |ws|
+      ws.on_message { |message| on_message(message) }
+      ws.on_close { on_close }
+      call(ws)
+    end
+  end
+
+  abstract def call(socket : Lucky::WebSocket)
+
+  def call
+    raise <<-ERROR
+    WebSocketAction must define `call(socket)`
+    ERROR
+  end
+end

src/lucky/web_socket_handler.cr

@@ -0,0 +1,26 @@
+class WebSocketHandler
+  include HTTP::Handler
+
+  def call(context : HTTP::Server::Context)
+    if ws_route_found?(context) && websocket_upgrade_request?(context)
+      websocket_action.payload.new(context, websocket_action.params).websocket.call(context)
+    else
+      call_next(context)
+    end
+  end
+
+  private def ws_route_found?(context)
+    !!websocket_action
+  end
+
+  memoize def websocket_action : LuckyRouter::Match(Lucky::Action.class)?
+    Lucky::Router.find_action(:ws, context.request.path)
+  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