feat: add authentication to yjs server

api/src/yjs-server.ts

@@ -3,8 +3,51 @@ import { WebSocketServer } from "ws";
 import express from "express";
 import http from "http";
 
+import jwt from "jsonwebtoken";
+
 import { setupWSConnection } from "./yjs-setupWS";
 
+import prisma from "./client";
+
+interface TokenInterface {
+  id: string;
+}
+
+/**
+ * Check if user has permission to access document.
+ * @param param0
+ * @returns
+ */
+async function checkPermission({
+  docName,
+  userId,
+}): Promise<"read" | "write" | "none"> {
+  // Docname is socket/he3og11sp3b73oh7k47o
+  // We need to get the actual ID of the pod
+  const repoId = docName.split("/")[1];
+  // Query the DB for the pod
+  const repo = await prisma.repo.findFirst({
+    where: {
+      id: repoId,
+    },
+    include: {
+      collaborators: true,
+      owner: true,
+    },
+  });
+  if (!repo) return "none";
+  if (
+    repo.owner.id === userId ||
+    repo.collaborators.find((collab) => collab.id === userId)
+  ) {
+    return "write";
+  }
+  if (repo.public) {
+    return "read";
+  }
+  return "none";
+}
+
 async function startServer() {
   const expapp = express();
   expapp.use(express.json({ limit: "20mb" }));
@@ -13,16 +56,44 @@ async function startServer() {
 
   wss.on("connection", setupWSConnection);
 
-  http_server.on("upgrade", (request, socket, head) => {
+  http_server.on("upgrade", async (request, socket, head) => {
     // You may check auth of request here..
     // See https://github.com/websockets/ws#client-authentication
-    /**
-     * @param {any} ws
-     */
-    const handleAuth = (ws) => {
-      wss.emit("connection", ws, request);
-    };
-    wss.handleUpgrade(request, socket, head, handleAuth);
+    if (request.url) {
+      const url = new URL(`ws://${request.headers.host}${request.url}`);
+      const docName = request.url.slice(1).split("?")[0];
+      const token = url.searchParams.get("token");
+      if (token) {
+        const decoded = jwt.verify(
+          token,
+          process.env.JWT_SECRET as string
+        ) as TokenInterface;
+        const userId = decoded.id;
+        const permission = await checkPermission({ docName, userId });
+        switch (permission) {
+          case "read":
+            // TODO I should disable editing in the frontend as well.
+            wss.handleUpgrade(request, socket, head, function done(ws) {
+              wss.emit("connection", ws, request, { readOnly: true });
+            });
+            break;
+          case "write":
+            wss.handleUpgrade(request, socket, head, function done(ws) {
+              wss.emit("connection", ws, request, { readOnly: false });
+            });
+            break;
+          case "none":
+            // This should not happen. This should be blocked by frontend code.
+            socket.write("HTTP/1.1 401 Unauthorized\r\n\r\n");
+            socket.destroy();
+            return;
+        }
+        return;
+      }
+    }
+    socket.write("HTTP/1.1 401 Unauthorized\r\n\r\n");
+    socket.destroy();
+    return;
   });
 
   const port = process.env.PORT || 4233;
@@ -32,3 +103,10 @@ async function startServer() {
 }
 
 startServer();
+
+// ts-node-dev might fail to restart. Force the exiting and restarting. Ref:
+// https://github.com/wclr/ts-node-dev/issues/69#issuecomment-493675960
+process.on("SIGTERM", () => {
+  console.log("Received SIGTERM. Exiting...");
+  process.exit();
+});

api/src/yjs-setupWS.ts

@@ -104,20 +104,48 @@ export const getYDoc = (docname, gc = true) =>
     return doc;
   });
 
+/**
+ * Support Read-only mode. Ref: https://discuss.yjs.dev/t/read-only-or-one-way-only-sync/135/4
+ */
+const readSyncMessage = (
+  decoder,
+  encoder,
+  doc,
+  transactionOrigin,
+  readOnly = false
+) => {
+  const messageType = decoding.readVarUint(decoder);
+  switch (messageType) {
+    case syncProtocol.messageYjsSyncStep1:
+      syncProtocol.readSyncStep1(decoder, encoder, doc);
+      break;
+    case syncProtocol.messageYjsSyncStep2:
+      if (!readOnly)
+        syncProtocol.readSyncStep2(decoder, doc, transactionOrigin);
+      break;
+    case syncProtocol.messageYjsUpdate:
+      if (!readOnly) syncProtocol.readUpdate(decoder, doc, transactionOrigin);
+      break;
+    default:
+      throw new Error("Unknown message type");
+  }
+  return messageType;
+};
+
 /**
  * @param {any} conn
  * @param {WSSharedDoc} doc
  * @param {Uint8Array} message
  */
-const messageListener = (conn, doc, message) => {
+const messageListener = (conn, doc, message, readOnly) => {
   try {
     const encoder = encoding.createEncoder();
     const decoder = decoding.createDecoder(message);
     const messageType = decoding.readVarUint(decoder);
     switch (messageType) {
       case messageSync:
         encoding.writeVarUint(encoder, messageSync);
-        syncProtocol.readSyncMessage(decoder, encoder, doc, conn);
+        readSyncMessage(decoder, encoder, doc, conn, readOnly);
 
         // If the `encoder` only contains the type of reply message and no
         // message, there is no need to send the message. When `encoder` only
@@ -204,17 +232,18 @@ const pingTimeout = 30000;
 export const setupWSConnection = (
   conn,
   req,
-  { docName = req.url.slice(1).split("?")[0], gc = true } = {}
+  { docName = req.url.slice(1).split("?")[0], gc = true, readOnly = false } = {}
 ) => {
   conn.binaryType = "arraybuffer";
+  console.log(`setupWSConnection ${docName}, read-only=${readOnly}`);
   // get doc, initialize if it does not exist yet
   const doc = getYDoc(docName, gc);
   doc.conns.set(conn, new Set());
   // listen and reply to events
   conn.on(
     "message",
     /** @param {ArrayBuffer} message */ (message) =>
-      messageListener(conn, doc, new Uint8Array(message))
+      messageListener(conn, doc, new Uint8Array(message), readOnly)
   );
 
   // Check if connection is still alive

ui/src/lib/store/repoStateSlice.tsx

@@ -305,7 +305,12 @@ export const createRepoStateSlice: StateCreator<
         state.provider = new WebsocketProvider(
           serverURL,
           state.repoId,
-          state.ydoc
+          state.ydoc,
+          {
+            params: {
+              token: localStorage.getItem("token") || "",
+            },
+          }
         );
         // max retry time: 10s
         state.provider.connect();