From 06ced7790bfa70fa07ecebc272854593c7cd828f Mon Sep 17 00:00:00 2001
From: Sdogruyol <dogruyolserdar@gmail.com>
Date: Tue, 15 Dec 2015 23:11:21 +0200
Subject: [PATCH] Adding websocket support :)

---
 spec/kemal_ws_handler_spec.cr  | 61 ++++++++++++++++++++++++++++++++++
 src/kemal/config.cr            |  5 +++
 src/kemal/dsl.cr               |  4 +++
 src/kemal/websocket_handler.cr | 13 ++++++++
 4 files changed, 83 insertions(+)
 create mode 100644 spec/kemal_ws_handler_spec.cr
 create mode 100644 src/kemal/websocket_handler.cr

diff --git a/spec/kemal_ws_handler_spec.cr b/spec/kemal_ws_handler_spec.cr
new file mode 100644
index 0000000..67bd515
--- /dev/null
+++ b/spec/kemal_ws_handler_spec.cr
@@ -0,0 +1,61 @@
+require "./spec_helper"
+
+describe "Kemal::WebsocketHandler" do
+
+  it "doesn't match on wrong route" do
+    handler = Kemal::WebsocketHandler.new "/" { }
+    headers = HTTP::Headers{
+      "Upgrade":           "websocket",
+      "Connection":        "Upgrade",
+      "Sec-WebSocket-Key": "dGhlIHNhbXBsZSBub25jZQ==",
+    }
+    request = HTTP::Request.new("GET", "/asd", headers)
+    response = handler.call request
+    response.status_code.should eq(404)
+  end
+
+  it "matches on given route" do
+    handler = Kemal::WebsocketHandler.new "/" { }
+    headers = HTTP::Headers{
+      "Upgrade":           "websocket",
+      "Connection":        "Upgrade",
+      "Sec-WebSocket-Key": "dGhlIHNhbXBsZSBub25jZQ==",
+    }
+    request = HTTP::Request.new("GET", "/", headers)
+    response = handler.call request
+    response.status_code.should eq(101)
+    response.headers["Upgrade"].should eq("websocket")
+    response.headers["Connection"].should eq("Upgrade")
+    response.headers["Sec-WebSocket-Accept"].should eq("s3pPLMBiTxaQ9kYGzzhZRbK+xOo=")
+    response.upgrade_handler.should_not be_nil
+  end
+
+  it "doesn't mix http and ws on same route" do
+    kemal = Kemal::Handler.new
+    kemal.add_route "GET", "/" do |env|
+      "hello #{env.params["message"]}"
+    end
+
+    ws_handler = Kemal::WebsocketHandler.new "/" { }
+    headers = HTTP::Headers{
+      "Upgrade":           "websocket",
+      "Connection":        "Upgrade",
+      "Sec-WebSocket-Key": "dGhlIHNhbXBsZSBub25jZQ==",
+    }
+
+    # HTTP Request
+    request = HTTP::Request.new("GET", "/?message=world")
+    response = kemal.call(request)
+    response.body.should eq("hello world")
+
+    # Websocket request
+    request = HTTP::Request.new("GET", "/", headers)
+    response = ws_handler.call request
+    response.status_code.should eq(101)
+    response.headers["Upgrade"].should eq("websocket")
+    response.headers["Connection"].should eq("Upgrade")
+    response.headers["Sec-WebSocket-Accept"].should eq("s3pPLMBiTxaQ9kYGzzhZRbK+xOo=")
+    response.upgrade_handler.should_not be_nil
+  end
+
+end
diff --git a/src/kemal/config.cr b/src/kemal/config.cr
index 84477d2..e5fcca5 100644
--- a/src/kemal/config.cr
+++ b/src/kemal/config.cr
@@ -4,6 +4,7 @@ module Kemal
   class Config
     INSTANCE = Config.new
     HANDLERS = [] of HTTP::Handler
+    WS_HANDLERS = [] of HTTP::Handler
     property ssl, port, env, workers, public_folder
 
     def initialize
@@ -26,6 +27,10 @@ module Kemal
       HANDLERS << handler
     end
 
+    def add_ws_handler(handler : HTTP::WebSocketHandler)
+      HANDLERS << handler
+    end
+
     # Reads configuration from config.yml. Currently it only supports the public_folder
     # option.
     # config.yml
diff --git a/src/kemal/dsl.cr b/src/kemal/dsl.cr
index 111613c..3e9685f 100644
--- a/src/kemal/dsl.cr
+++ b/src/kemal/dsl.cr
@@ -5,3 +5,7 @@ HTTP_METHODS = %w(get post put patch delete)
    Kemal::Handler::INSTANCE.add_route({{method}}.upcase, path, &block)
   end
 {% end %}
+
+def ws(path, &block : HTTP::WebSocketHandler::WebSocketSession -> _)
+  Kemal::WebsocketHandler.new path, &block
+end
diff --git a/src/kemal/websocket_handler.cr b/src/kemal/websocket_handler.cr
new file mode 100644
index 0000000..dd646b2
--- /dev/null
+++ b/src/kemal/websocket_handler.cr
@@ -0,0 +1,13 @@
+class Kemal::WebsocketHandler < HTTP::WebSocketHandler
+  getter handler
+
+  def initialize(@path, &@proc : WebSocketSession ->)
+    @handler = @proc
+    Kemal.config.add_ws_handler self
+  end
+
+  def call(request)
+    return call_next(request) unless request.path.not_nil! == @path
+    super
+  end
+end