From 3b691870070505d07af7aac964ddc19170da1f3c Mon Sep 17 00:00:00 2001
From: Arijit Das <arijit@status.im>
Date: Tue, 1 Dec 2020 18:13:59 +0530
Subject: [PATCH] Implement web socket handshake

---
 .editorconfig   |  5 +++
 Readme.md       | 21 ++++++++++++
 lint.nims       | 14 ++++++++
 test/server.nim |  8 +++++
 ws.nimble       | 11 +++++++
 ws/ws.nim       | 86 +++++++++++++++++++++++++++++++++++++++++++++++++
 6 files changed, 145 insertions(+)
 create mode 100644 .editorconfig
 create mode 100644 Readme.md
 create mode 100644 lint.nims
 create mode 100644 test/server.nim
 create mode 100644 ws.nimble
 create mode 100644 ws/ws.nim

diff --git a/.editorconfig b/.editorconfig
new file mode 100644
index 0000000000..1963f8debf
--- /dev/null
+++ b/.editorconfig
@@ -0,0 +1,5 @@
+[*]
+indent_style = space
+insert_final_newline = true
+indent_size = 2
+trim_trailing_whitespace = true
\ No newline at end of file
diff --git a/Readme.md b/Readme.md
new file mode 100644
index 0000000000..bd3ae88716
--- /dev/null
+++ b/Readme.md
@@ -0,0 +1,21 @@
+Websocket for Nim
+
+We're working towards an implementation of the
+[Websocket](https://tools.ietf.org/html/rfc6455) protocol for
+[Nim](https://nim-lang.org/). This is very much a work in progress, and not yet
+in a usable state.
+
+Building and testing
+--------------------
+
+Install dependencies:
+
+```bash
+nimble install -d
+```
+
+Run tests:
+
+```bash
+nimble test
+```
diff --git a/lint.nims b/lint.nims
new file mode 100644
index 0000000000..08b8c830d8
--- /dev/null
+++ b/lint.nims
@@ -0,0 +1,14 @@
+#!/usr/bin/env nim
+import std/strutils
+
+proc lintFile*(file: string) =
+  if file.endsWith(".nim"):
+    exec "nimpretty " & file
+
+proc lintDir*(dir: string) =
+  for file in listFiles(dir):
+    lintFile(file)
+  for subdir in listDirs(dir):
+    lintDir(subdir)
+
+lintDir(projectDir())
\ No newline at end of file
diff --git a/test/server.nim b/test/server.nim
new file mode 100644
index 0000000000..a0b83b19de
--- /dev/null
+++ b/test/server.nim
@@ -0,0 +1,8 @@
+import ../ws, chronos, asynchttpserver
+
+proc cb(req: Request) {.async.} =
+  var ws = await newWebSocket(req)
+  ws.close()
+
+var server = newAsyncHttpServer()
+waitFor server.serve(Port(9001), cb)
diff --git a/ws.nimble b/ws.nimble
new file mode 100644
index 0000000000..44459ee141
--- /dev/null
+++ b/ws.nimble
@@ -0,0 +1,11 @@
+packageName = "ws"
+version = "0.1.0"
+author = "Status Research & Development GmbH"
+description = "WS protocol implementation"
+license = "MIT"
+
+requires "nim >= 1.2.6"
+requires "chronos >= 2.5.2 & < 3.0.0"
+
+task lint, "format source files according to the official style guide":
+  exec "./lint.nims"
\ No newline at end of file
diff --git a/ws/ws.nim b/ws/ws.nim
new file mode 100644
index 0000000000..44e390146b
--- /dev/null
+++ b/ws/ws.nim
@@ -0,0 +1,86 @@
+import chronos, asyncdispatch, asynchttpserver, base64, nativesockets
+
+type HeaderVerificationError* {.pure.} = enum
+  None
+    ## No error.
+  UnsupportedVersion
+    ## The Sec-Websocket-Version header gave an unsupported version.
+    ## The only currently supported version is 13.
+  NoKey
+    ## No Sec-Websocket-Key was provided.
+  ProtocolAdvertised
+    ## A protocol was advertised but the server gave no protocol.
+  NoProtocolsSupported
+    ## None of the advertised protocols match the server protocol.
+  NoProtocolAdvertised
+    ## Server asked for a protocol but no protocol was advertised.
+
+proc `$`*(error: HeaderVerificationError): string =
+  const errorTable: array[HeaderVerificationError, string] = [
+    "no error",
+    "the only supported sec-websocket-version is 13",
+    "no sec-websocket-key provided",
+    "server does not support protocol negotation",
+    "no advertised protocol supported",
+    "no protocol advertised"
+  ]
+  result = errorTable[error]
+
+type
+  ReadyState* = enum
+    Connecting = 0 # The connection is not yet open.
+    Open = 1       # The connection is open and ready to communicate.
+    Closing = 2    # The connection is in the process of closing.
+    Closed = 3     # The connection is closed or couldn't be opened.
+
+  WebSocket* = ref object
+    tcpSocket*: AsyncSocket
+    version*: int
+    key*: string
+    protocol*: string
+    readyState*: ReadyState
+    masked*: bool # send masked packets
+
+  WebSocketError* = object of IOError
+
+proc handshake*(ws: WebSocket, headers: HttpHeaders): Future[error: HeaderVerificationError] {.async.} =
+  ws.version = parseInt(headers["Sec-WebSocket-Version"])
+  ws.key = headers["Sec-WebSocket-Key"].strip()
+  if headers.hasKey("Sec-WebSocket-Protocol"):
+    let wantProtocol = headers["Sec-WebSocket-Protocol"].strip()
+    if ws.protocol != wantProtocol:
+      return NoProtocolsSupported
+
+  let
+    sh = secureHash(ws.key & "258EAFA5-E914-47DA-95CA-C5AB0DC85B11")
+    acceptKey = base64.encode(decodeBase16($sh))
+
+  var response = "HTTP/1.1 101 Web Socket Protocol Handshake\c\L"
+  response.add("Sec-WebSocket-Accept: " & acceptKey & "\c\L")
+  response.add("Connection: Upgrade\c\L")
+  response.add("Upgrade: webSocket\c\L")
+
+  if ws.protocol != "":
+    response.add("Sec-WebSocket-Protocol: " & ws.protocol & "\c\L")
+  response.add "\c\L"
+
+  await ws.tcpSocket.send(response)
+  ws.readyState = Open
+
+proc newWebSocket*(req: Request, protocol: string = ""): Future[tuple[ws: AsyncWebSocket, error: HeaderVerificationError]] {.async.} =
+  if not req.headers.hasKey("Sec-WebSocket-Version"):
+    return ("", UnsupportedVersion)
+
+  var ws = WebSocket()
+  ws.masked = false
+  ws.tcpSocket = req.client
+  ws.protocol = protocol
+  let (ws, error) = await ws.handshake(req.headers)
+  return ws, error
+
+proc close*(ws: WebSocket) =
+  ws.readyState = Closed
+  proc close() {.async.} =
+    await ws.send("", Close)
+    ws.tcpSocket.close()
+  asyncCheck close()