Added a very basic image API authentication system

api/index.js

@@ -27,6 +27,7 @@ const queue = [];
 const { v4: uuidv4 } = require("uuid");
 
 const MAX_JOBS = process.env.JOBS !== "" && process.env.JOBS !== undefined ? parseInt(process.env.JOBS) : os.cpus().length * 4; // Completely arbitrary, should usually be some multiple of your amount of cores
+const PASS = process.env.PASS !== "" && process.env.PASS !== undefined ? process.env.PASS : undefined;
 let jobAmount = 0;
 
 const acceptJob = (uuid, sock) => {
@@ -108,6 +109,10 @@ httpServer.on("request", async (req, res) => {
     res.statusCode = 405;
     return res.end("405 Method Not Allowed");
   }
+  if (PASS && req.headers.authentication !== PASS) {
+    res.statusCode = 401;
+    return res.end("401 Unauthorized");
+  }
   const reqUrl = new URL(req.url, `http://${req.headers.host}`);
   if (reqUrl.pathname === "/status" && req.method === "GET") {
     log(`Sending server status to ${req.socket.remoteAddress}:${req.socket.remotePort} via HTTP`);
@@ -154,6 +159,12 @@ httpServer.on("request", async (req, res) => {
 httpServer.on("upgrade", (req, sock, head) => {
   const reqUrl = new URL(req.url, `http://${req.headers.host}`);
 
+  if (PASS && req.headers.authentication !== PASS) {
+    sock.write("HTTP/1.1 401 Unauthorized\r\n\r\n");
+    sock.destroy();
+    return;
+  }
+
   if (reqUrl.pathname === "/sock") {
     wss.handleUpgrade(req, sock, head, (ws) => {
       wss.emit("connection", ws, req);

servers.json

@@ -3,7 +3,7 @@
     { "id": "1", "host": "localhost", "port": 2333, "password": "youshallnotpass", "local": true }
   ],
   "image": [
-    "localhost"
+    { "server": "localhost", "auth": "verycoolpass100" }
   ],
   "searx": [
     "https://searx.xyz"

utils/services/image.js

@@ -25,7 +25,7 @@ class ImageWorker extends BaseServiceWorker {
     if (process.env.API === "true") {
       for (const server of this.servers) {
         try {
-          await this.connect(server);
+          await this.connect(server.server, server.auth);
         } catch (e) {
           logger.error(e);
         }
@@ -53,7 +53,13 @@ class ImageWorker extends BaseServiceWorker {
         controller.abort();
       }, 2000);
       try {
-        const statusRequest = await fetch(`http://${address}:8080/running`, { signal: controller.signal });
+        const auth = this.servers.filter((val) => val.server === address)[0].auth;
+        const statusRequest = await fetch(`http://${address}:8080/running`, {
+          signal: controller.signal,
+          headers: {
+            "Authentication": auth && auth !== "" ? auth : undefined
+          }
+        });
         clearTimeout(timeout);
         const status = await statusRequest.json();
         serversLeft--;
@@ -89,7 +95,7 @@ class ImageWorker extends BaseServiceWorker {
     if (serversLeft < this.servers.length) {
       for (const server of this.servers) {
         try {
-          if (!this.connections.has(server)) await this.connect(server);
+          if (!this.connections.has(server.server)) await this.connect(server.server, server.auth);
         } catch (e) {
           logger.error(e);
         }
@@ -108,7 +114,13 @@ class ImageWorker extends BaseServiceWorker {
         controller.abort();
       }, 5000);
       try {
-        const statusRequest = await fetch(`http://${address}:8080/status`, { signal: controller.signal });
+        const auth = this.servers.filter((val) => val.server === address)[0].auth;
+        const statusRequest = await fetch(`http://${address}:8080/status`, {
+          signal: controller.signal,
+          headers: {
+            "Authentication": auth && auth !== "" ? auth : undefined
+          }
+        });
         clearTimeout(timeout);
         const status = await statusRequest.text();
         serversLeft--;
@@ -137,8 +149,12 @@ class ImageWorker extends BaseServiceWorker {
     }
   }
 
-  async connect(server) {
+  async connect(server, auth) {
-    const connection = new WebSocket(`ws://${server}:8080/sock`);
+    const connection = new WebSocket(`ws://${server}:8080/sock`, {
+      headers: {
+        "Authentication": auth && auth !== "" ? auth : undefined
+      }
+    });
     connection.on("message", async (msg) => {
       const opcode = msg.readUint8(0);
       const req = msg.slice(37, msg.length);
@@ -150,7 +166,11 @@ class ImageWorker extends BaseServiceWorker {
       } else if (opcode === 0x01) { // Job completed successfully
         // the image API sends all job responses over the same socket; make sure this is ours
         if (this.jobs[uuid]) {
-          const imageReq = await fetch(`http://${server}:8080/image?id=${uuid}`);
+          const imageReq = await fetch(`http://${server}:8080/image?id=${uuid}`, {
+            headers: {
+              "Authentication": auth && auth !== "" ? auth : undefined
+            }
+          });
           const image = await imageReq.buffer();
           // The response data is given as the file extension/ImageMagick type of the image (e.g. "png"), followed
           // by a newline, followed by the image data.
@@ -260,7 +280,7 @@ class ImageWorker extends BaseServiceWorker {
         let amount = 0;
         for (const server of this.servers) {
           try {
-            await this.connect(server);
+            await this.connect(server.server, server.auth);
             amount += 1;
           } catch (e) {
             logger.error(e);