add support for error codes on websockets

copied from 541dfc9da5/src/onyx-http/middleware/router/websocket_handler.cr (L36)

src/kemal/websocket.cr

@@ -1,3 +1,5 @@
+require "http"
+
 module Kemal
   # Takes 2 parameters: *path* and a *handler* to specify
   # what action to be done if the route is matched.
@@ -7,8 +9,56 @@ module Kemal
     def initialize(@path : String, &@proc : HTTP::WebSocket, HTTP::Server::Context -> Void)
     end
 
+    def error(code : Int16, message : String)
+      err = WebsocketError.new(code, message)
+      raise err
+    end
+
     def call(context : HTTP::Server::Context)
-      super
+      if websocket_upgrade_request?(context.request)
+        response = context.response
+
+        version = context.request.headers["Sec-WebSocket-Version"]?
+        unless version == ::HTTP::WebSocket::Protocol::VERSION
+          response.headers["Sec-WebSocket-Version"] = ::HTTP::WebSocket::Protocol::VERSION
+          raise UpgradeRequired.new
+        end
+
+        key = context.request.headers["Sec-WebSocket-Key"]?
+        raise BadRequest.new("Sec-WebSocket-Key header is missing") unless key
+
+        accept_code = ::HTTP::WebSocket::Protocol.key_challenge(key)
+
+        response.status_code = 101
+        response.headers["Upgrade"] = "websocket"
+        response.headers["Connection"] = "Upgrade"
+        response.headers["Sec-WebSocket-Accept"] = accept_code
+
+        response.upgrade do |io|
+          socket = ::HTTP::WebSocket.new(io)
+          @proc.call(socket, context)
+          socket.run
+        rescue error : Exception
+          if error.is_a?(WebsocketError)
+            # TODO
+            context.response.status_code = 500
+            code = error.code.to_i16
+            message = error.status_message
+          else
+            context.response.status_code = error.code
+            code = 1011_i16
+            message = "Exception"
+          end
+
+          raw = uninitialized UInt8[2]
+          IO::ByteFormat::BigEndian.encode(code, raw.to_slice)
+          socket.not_nil!.close(String.new(raw.to_slice) + message)
+
+          raise error unless error.is_a?(WebsocketError)
+        end
+      else
+        raise UpgradeRequired.new
+      end
     end
   end
 end

src/kemal/websocket_handler.cr

@@ -1,4 +1,24 @@
 module Kemal
+  class WebsocketError < Exception
+    property code : Int32 = 1011
+    property status_message : String = "websocket error"
+
+    def initialize(@code, @status_message : String)
+    end
+  end
+
+  class UpgradeRequired < WebsocketError
+    def initialize
+      super(426, "Upgrade required")
+    end
+  end
+
+  class BadRequest < WebsocketError
+    def initialize(status_message : String)
+      super(400, status_message)
+    end
+  end
+
   class WebSocketHandler
     include HTTP::Handler