Add basic daemon connection authentication

app/src/daemon.ts

@@ -1,24 +1,25 @@
 "use strict";
 
+import http, { IncomingMessage } from "http";
+import WebSocket from "ws";
+
 import { spawn } from "child_process";
 import { Readable } from "stream";
 
-import WebSocket from "ws";
-
 import { DecodeTransform, EncodeTransform } from "./transforms";
 
-export function init(port: number) {
+export function init(port: number, serverPassword?: string) {
+    const server = http.createServer();
+    const wss = new WebSocket.Server({ noServer: true });
+
     process.stdout.write("Starting WebSocket server... ");
 
-    const wss = new WebSocket.Server({ port }, () => {
-        // eslint-disable-next-line no-console
-        console.log("Done!");
+    server.on("listening", () => {
+        process.stdout.write("Done!\n");
     });
-
     wss.on("error", err => {
-        // eslint-disable-next-line no-console
-        console.log("Failed!");
-        console.error(err);
+        console.error("Failed!");
+        console.error(err.message);
     });
 
     wss.on("connection", socket => {
@@ -42,14 +43,64 @@ export function init(port: number) {
 
         // bridge.stdout -> socket
         bridge.stdout.pipe(new DecodeTransform()).on("data", data => {
-            // Socket can be CLOSING here
-            if (socket.readyState === WebSocket.OPEN) {
-                socket.send(JSON.stringify(data));
+            if (socket.readyState !== WebSocket.OPEN) {
+                return;
             }
+
+            socket.send(JSON.stringify(data));
         });
 
         // Handle termination
         socket.on("close", () => bridge.kill());
         bridge.on("exit", () => socket.close());
     });
+
+    /**
+     * Authenticates requests by checking password URL param against
+     * server password specified in launch options.
+     */
+    function authenticate(req: IncomingMessage) {
+        if (!serverPassword) return true;
+
+        const password = new URL(
+            req.url!,
+            `http://${req.headers.host}`
+        ).searchParams.get("password");
+
+        return password === serverPassword;
+    }
+
+    server.on("upgrade", (req, socket, head) => {
+        if (
+            // Only accept WebSocket requests from extension origins
+            !req.headers.origin?.startsWith("moz-extension://") ||
+            !authenticate(req)
+        ) {
+            socket.write("HTTP/1.1 401 Unauthorized\r\n\r\n");
+            socket.destroy();
+            return;
+        }
+
+        wss.handleUpgrade(req, socket, head, ws => {
+            wss.emit("connection", ws, req);
+        });
+    });
+
+    /**
+     * JS WebSocket API does not allow access to connection errors, so
+     * provide an endpoint for feedback on invalid authentication.
+     */
+    server.on("request", (req, res) => {
+        if (!authenticate(req)) {
+            res.writeHead(401);
+            res.end();
+
+            return;
+        }
+
+        res.writeHead(200);
+        res.end();
+    });
+
+    server.listen(port);
 }

app/src/main.ts

@@ -6,7 +6,7 @@ import { __applicationVersion } from "../package.json";
 
 const argv = minimist(process.argv.slice(2), {
     boolean: ["daemon", "help", "version"],
-    string: ["__name", "port"],
+    string: ["__name", "port", "password"],
     alias: {
         d: "daemon",
         h: "help",
@@ -36,6 +36,11 @@ Options:
                    options.
   -p, --port       Set port number for WebSocket server. This must match the
                    port set in the extension options.
+  --password       Set an optional password for the WebSocket server. This must
+                   match the password set in the extension options.
+                   WARNING: This password is intended only as a basic access
+                   control measure and is transmitted in plain text even over
+                   remote connections!
 `
     );
 } else if (argv.daemon) {
@@ -46,7 +51,7 @@ Options:
     }
 
     import("./daemon").then(daemon => {
-        daemon.init(port);
+        daemon.init(port, argv.password);
     });
 } else {
     import("./bridge");