Switching to ws instead of socket.io and implementing basic authentication

src/connection.ts

@@ -1,43 +1,91 @@
-import * as io from "socket.io";
-import { Server } from "http";
+import * as WebSocket from "ws";
+import { Server, IncomingMessage } from "http";
 import { DatabaseManager } from "./database/database";
 import Logging from "@hibas123/logging";
 import Query from "./database/query";
 import Session from "./database/session";
 import shortid = require("shortid");
 
+import * as JWT from "jsonwebtoken";
+
+async function verifyJWT(token: string, publicKey: string) {
+   return new Promise<any | undefined>((yes) => {
+      JWT.verify(token, publicKey, (err, decoded) => {
+         if (err)
+            yes(undefined);
+         else
+            yes(decoded);
+      })
+   })
+}
+
+import { URLSearchParams } from "url";
+
 type QueryTypes = "get" | "set" | "push" | "subscribe" | "unsubscribe";
 
 export class ConnectionManager {
-   static server: io.Server;
+   static server: WebSocket.Server;
 
    static bind(server: Server) {
-      this.server = io(server);
+      this.server = new WebSocket.Server({ server });
       this.server.on("connection", this.onConnection.bind(this));
    }
 
-   private static onConnection(socket: io.Socket) {
-      Logging.debug("New Connection:", socket.id);
-      const reqMap = new Map<string, [number, number]>();
-      const stored = new Map<string, Query>();
+   private static async onConnection(socket: WebSocket, req: IncomingMessage) {
+      Logging.log("New Connection:");
+      const sendError = (msg: string) => socket.send(JSON.stringify({ ns: "error_msg", args: [msg] }));
+
+
       const session = new Session();
 
-      const answer = (id: string, data: any, err: boolean = false) => {
-         let time = process.hrtime(reqMap.get(id));
-         reqMap.delete(id);
-         // Logging.debug(`Sending answer for ${id} with data`, data, err ? "as error" : "", "Took", time[1] / 1000, "us");
-         socket.emit("message", id, err, data);
+      let query = new URLSearchParams(req.url.split("?").pop());
+
+      const database = query.get("database");
+      const db = DatabaseManager.getDatabase(database);
+      if (!db) {
+         sendError("Invalid Database!");
+         socket.close();
+         return;
       }
 
-      socket.on("login", (id: string) => {
-         //TODO: implement
-      })
+      const accesskey = query.get("accesskey");
+      if (db.accesskey) {
+         if (!accesskey || accesskey !== db.accesskey) {
+            sendError("Unauthorized!");
+            socket.close();
+            return;
+         }
+      }
 
-      socket.on("query", async (id: string, type: QueryTypes, database: string, path: string[], data: any) => {
-         Logging.debug(`Request with id '${id}' from type '${type}' for database '${database}' and path '${path}' with data`, data)
-         reqMap.set(id, process.hrtime());
+      const authkey = query.get("authkey");
+      if (authkey && db.publickey) {
+         let res = await verifyJWT(authkey, db.publickey);
+         if (!res || !res.uid) {
+            sendError("Invalid JWT");
+            socket.close();
+            return;
+         } else {
+            session.uid = res.uid;
+         }
+      }
+
+      const rootkey = query.get("rootkey");
+      if (rootkey && db.rootkey) {
+         if (rootkey === db.rootkey) {
+            session.root = true;
+            Logging.warning(`Somebody logged into ${database} via rootkey`);
+         }
+      }
+
+      const stored = new Map<string, Query>();
+      const answer = (id: string, data: any, err: boolean = false) => {
+         socket.send(JSON.stringify({ ns: "message", args: [id, err, data] }));
+      }
+
+      const handler = new Map<string, ((...args: any[]) => void)>();
+      handler.set("query", async (id: string, type: QueryTypes, path: string[], data: any) => {
+         Logging.debug(`Request with id '${id}' from type '${type}' and path '${path}' with data`, data)
          try {
-            const db = DatabaseManager.getDatabase(database);
             const perms = db.rules.hasPermission(path, session);
             const noperm = new Error("No permisison!");
             if (!db)
@@ -68,7 +116,7 @@ export class ConnectionManager {
                      let subscriptionID = shortid.generate();
 
                      query.subscribe(data, (data) => {
-                        socket.emit("event", subscriptionID, data);
+                        socket.send(JSON.stringify({ ns: "event", args: [subscriptionID, data] }));
                      });
                      stored.set(id, query);
                      answer(id, subscriptionID);
@@ -86,11 +134,18 @@ export class ConnectionManager {
             Logging.error(err);
             answer(id, err.message, true);
          }
+      })
 
+      socket.on("message", (rawData: string) => {
+         let data: { ns: string, args: any[] } = JSON.parse(rawData);
+         let h = handler.get(data.ns);
+         if (h) {
+            h(...data.args);
+         }
       })
 
       socket.on("disconnect", () => {
-         reqMap.clear();
+         stored.forEach(query => query.unsubscribe());
          stored.clear();
          socket.removeAllListeners();
       })

src/database/rules.ts

@@ -1,4 +1,5 @@
 import Session from "./session";
+import Logging from "@hibas123/nodelogging";
 
 interface IRule<T> {
    ".write"?: T
@@ -64,7 +65,27 @@ export class Rules {
       let rules = this.rules;
 
       for (let segment of path) {
-         rules = rules[segment] || rules["*"];
+         if (segment.startsWith("$") || segment.startsWith(".")) {
+            read = false;
+            write = false;
+            Logging.log("Invalid query path (started with '$' or '.'):", path);
+            break;
+         }
+
+         let k = Object.keys(rules)
+            .filter(e => e.startsWith("$"))
+            .find(e => {
+               switch (e) {
+                  case "$uid":
+                     if (segment === session.uid)
+                        return true;
+                     break;
+               }
+               return false;
+            })
+
+         rules = (k ? rules[k] : undefined) || rules[segment] || rules["*"];
+
          if (rules) {
             if (rules[".read"]) {
                read = rules[".read"]

src/database/session.ts

@@ -1,3 +1,4 @@
 export default class Session {
-
+   root: boolean = false;
+   uid: string = undefined;
 }