import { config } from "dotenv"; config(); import { cpus } from "os"; import { Worker } from "worker_threads"; import { join } from "path"; import { createServer } from "http"; import { WebSocketServer } from "ws"; import { fileURLToPath } from "url"; import { dirname } from "path"; import { createRequire } from "module"; const nodeRequire = createRequire(import.meta.url); const magick = nodeRequire(`../build/${process.env.DEBUG && process.env.DEBUG === "true" ? "Debug" : "Release"}/image.node`); /* Rerror 0x01 Tqueue 0x02 Rqueue 0x03 Tcancel 0x04 Rcancel 0x05 Twait 0x06 Rwait 0x07 Rinit 0x08 */ const Rerror = 0x01; const Tqueue = 0x02; const Rqueue = 0x03; const Tcancel = 0x04; const Rcancel = 0x05; const Twait = 0x06; const Rwait = 0x07; const Rinit = 0x08; const start = process.hrtime(); const log = (msg, jobNum) => { console.log(`[${process.hrtime(start)[1] / 1000000}${jobNum !== undefined ? `:${jobNum}` : ""}]\t ${msg}`); }; class JobCache extends Map { set(key, value) { super.set(key, value); setTimeout(() => { if (super.has(key) && this.get(key) === value && value.data) super.delete(key); }, 300000); // delete jobs if not requested after 5 minutes } } const jobs = new JobCache(); // Should look like ID : { msg: "request", num: } const queue = []; // Array of IDs const MAX_JOBS = process.env.JOBS !== "" && process.env.JOBS !== undefined ? parseInt(process.env.JOBS) : 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 = (id, sock) => { jobAmount++; queue.shift(); const job = jobs.get(id); return runJob({ id: id, msg: job.msg, num: job.num }, sock).then(() => { log(`Job ${id} has finished`); }).catch((err) => { console.error(`Error on job ${id}:`, err); const newJob = jobs.get(id); if (!newJob.tag) { newJob.error = err.message; jobs.set(id, newJob); return; } jobs.delete(id); sock.send(Buffer.concat([Buffer.from([Rerror]), newJob.tag, Buffer.from(err.message)])); }).finally(() => { jobAmount--; if (queue.length > 0) { acceptJob(queue[0], sock); } }); }; const wss = new WebSocketServer({ clientTracking: true, noServer: true }); wss.on("connection", (ws, request) => { log(`WS client ${request.socket.remoteAddress}:${request.socket.remotePort} has connected`); const num = Buffer.alloc(2); num.writeUInt16LE(MAX_JOBS); const formats = {}; for (const cmd of Object.keys(magick)) { formats[cmd] = ["image/png", "image/gif", "image/jpeg", "image/webp"]; } const init = Buffer.concat([Buffer.from([Rinit]), Buffer.from([0x00, 0x00]), num, Buffer.from(JSON.stringify(formats))]); ws.send(init); ws.on("error", (err) => { console.error(err); }); ws.on("message", (msg) => { const opcode = msg.readUint8(0); const tag = msg.slice(1, 3); const req = msg.toString().slice(3); if (opcode == Tqueue) { const id = msg.readUInt32LE(3); const obj = msg.slice(7); const job = { msg: obj, num: jobAmount }; jobs.set(id, job); queue.push(id); const newBuffer = Buffer.concat([Buffer.from([Rqueue]), tag]); ws.send(newBuffer); if (jobAmount < MAX_JOBS) { log(`Got WS request for job ${job.msg} with id ${id}`, job.num); acceptJob(id, ws); } else { log(`Got WS request for job ${job.msg} with id ${id}, queued in position ${queue.indexOf(id)}`, job.num); } } else if (opcode == Tcancel) { delete queue[queue.indexOf(req) - 1]; jobs.delete(req); const cancelResponse = Buffer.concat([Buffer.from([Rcancel]), tag]); ws.send(cancelResponse); } else if (opcode == Twait) { const id = msg.readUInt32LE(3); const job = jobs.get(id); if (!job) { const errorResponse = Buffer.concat([Buffer.from([Rerror]), tag, Buffer.from("Invalid job ID")]); ws.send(errorResponse); return; } if (job.error) { jobs.delete(id); const errorResponse = Buffer.concat([Buffer.from([Rerror]), tag, Buffer.from(job.error)]); ws.send(errorResponse); return; } job.tag = tag; jobs.set(id, job); //const waitResponse = Buffer.concat([Buffer.from([Rwait]), tag]); //ws.send(waitResponse); } else { log("Could not parse WS message"); } }); ws.on("close", () => { log(`WS client ${request.socket.remoteAddress}:${request.socket.remotePort} has disconnected`); }); }); wss.on("error", (err) => { console.error("A WS error occurred: ", err); }); const httpServer = createServer(); httpServer.on("request", async (req, res) => { if (req.method !== "GET") { 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 === "/image" && req.method === "GET") { if (!reqUrl.searchParams.has("id")) { res.statusCode = 400; return res.end("400 Bad Request"); } const id = parseInt(reqUrl.searchParams.get("id")); if (!jobs.has(id)) { res.statusCode = 410; return res.end("410 Gone"); } log(`Sending image data for job ${id} to ${req.socket.remoteAddress}:${req.socket.remotePort} via HTTP`); const ext = jobs.get(id).ext; let contentType; switch (ext) { case "gif": contentType = "image/gif"; break; case "png": contentType = "image/png"; break; case "jpeg": case "jpg": contentType = "image/jpeg"; break; case "webp": contentType = "image/webp"; break; } if (contentType) res.setHeader("Content-Type", contentType); else res.setHeader("Content-Type", ext); const data = jobs.get(id).data; jobs.delete(id); return res.end(data, (err) => { if (err) console.error(err); }); } else { res.statusCode = 404; return res.end("404 Not Found"); } }); 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); }); } else { sock.destroy(); } }); httpServer.on("error", (e) => { console.error("An HTTP error occurred: ", e); }); httpServer.listen(3762, () => { log("HTTP and WS listening on port 3762"); }); const runJob = (job, ws) => { return new Promise((resolve, reject) => { log(`Job ${job.id} starting...`, job.num); const object = JSON.parse(job.msg); // If the image has a path, it must also have a type if (object.path && !object.params.type) { reject(new TypeError("Unknown image type")); } const worker = new Worker(join(dirname(fileURLToPath(import.meta.url)), "../utils/image-runner.js"), { workerData: object }); const timeout = setTimeout(() => { worker.terminate(); reject(new Error("Job timed out")); }, 900000); log(`Job ${job.id} started`, job.num); worker.once("message", (data) => { clearTimeout(timeout); log(`Sending result of job ${job.id} back to the bot`, job.num); const jobObject = jobs.get(job.id); jobObject.data = data.buffer; jobObject.ext = data.fileExtension; jobs.set(job.id, jobObject); const waitResponse = Buffer.concat([Buffer.from([Rwait]), jobObject.tag]); ws.send(waitResponse); resolve(); }); worker.once("error", (e) => { clearTimeout(timeout); reject(e); }); }); };