mrmBot-Matrix/utils/image.js

const magick = require("../build/Release/image.node");
const { Worker } = require("worker_threads");
const fetch = require("node-fetch");
const AbortController = require("abort-controller");
const net = require("net");
const dgram = require("dgram");
const fileType = require("file-type");
const servers = require("../servers.json").image;
const path = require("path");

const formats = ["image/jpeg", "image/png", "image/webp", "image/gif"];

const chooseServer = async (ideal) => {
  if (ideal.length === 0) throw "No available servers";
  const sorted = ideal.sort((a, b) => {
    return b.load - a.load;
  });
  return sorted[0];
};

const getIdeal = () => {
  return new Promise((resolve, reject) => {
    const socket = dgram.createSocket("udp4");
    let serversLeft = servers.length;
    const idealServers = [];
    const timeout = setTimeout(() => {
      socket.close(async () => {
        try {
          const server = await chooseServer(idealServers);
          resolve(server);
        } catch (e) {
          reject(e);
        }
      });
    }, 5000);
    socket.on("message", async (msg, rinfo) => {
      const opcode = msg.readUint8(0);
      const res = parseInt(msg.slice(1, msg.length).toString());
      if (opcode === 0x3) {
        serversLeft--;
        idealServers.push({
          addr: rinfo.address,
          load: res
        });
        if (!serversLeft) {
          clearTimeout(timeout);
          socket.close(async () => {
            try {
              const server = await chooseServer(idealServers);
              resolve(server);
            } catch (e) {
              reject(e);
            }
          });
        }
      }
    });
    for (const server of servers) {
      socket.send(Buffer.from([0x2]), 8080, server, (err) => {
        if (err) reject(err);
      });
    }
  });
};

exports.check = (cmd) => {
  return magick[cmd] ? true : false;
};

exports.getType = async (image) => {
  if (!image.startsWith("http")) {
    const imageType = await fileType.fromFile(image);
    if (imageType && formats.includes(imageType.mime)) {
      return imageType.mime;
    }
    return undefined;
  }
  let type;
  const controller = new AbortController();
  const timeout = setTimeout(() => {
    controller.abort();
  }, 25000);
  try {
    const imageRequest = await fetch(image, { signal: controller.signal, headers: {
      "Range": "bytes=0-1023"
    }});
    clearTimeout(timeout);
    const imageBuffer = await imageRequest.buffer();
    const imageType = await fileType.fromBuffer(imageBuffer);
    if (imageType && formats.includes(imageType.mime)) {
      type = imageType.mime;
    }
  } catch (error) {
    if (error.name === "AbortError") {
      throw Error("Timed out");
    } else {
      throw error;
    }
  } finally {
    clearTimeout(timeout);
  }
  return type;
};

exports.run = object => {
  return new Promise(async (resolve, reject) => {
    if (process.env.API === "true") {
      // Connect to best image server
      const currentServer = await getIdeal();
      const socket = dgram.createSocket("udp4");
      const data = Buffer.concat([Buffer.from([0x1 /* queue job */]), Buffer.from(JSON.stringify(object))]);

      const timeout = setTimeout(() => {
        socket.close();
        reject("UDP timed out");
      }, 25000);

      let jobID;
      socket.on("message", (msg) => {
        const opcode = msg.readUint8(0);
        const req = msg.slice(37, msg.length);
        const uuid = msg.slice(1, 36).toString();
        if (opcode === 0x0) { // Job queued
          clearTimeout(timeout);
          jobID = uuid;
        } else if (opcode === 0x1) { // Job completed successfully
          // the image API sends all job responses over the same socket; make sure this is ours
          if (jobID === uuid) {
            const client = net.createConnection(req.toString(), currentServer.addr);
            const array = [];
            client.on("data", (rawData) => {
              array.push(rawData);
            });
            client.once("end", () => {
              const data = Buffer.concat(array);
              // 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.
              const delimIndex = data.indexOf("\n");
              socket.close();
              if (delimIndex === -1) reject("Could not parse response");
              const payload = {
                // Take just the image data
                buffer: data.slice(delimIndex + 1),
                type: data.slice(0, delimIndex).toString()
              };
              resolve(payload);
            });
            client.on("error", (err) => {
              socket.close();
              reject(err);
            });
          }
        } else if (opcode === 0x2) { // Job errored
          if (jobID === uuid) {
            socket.close();
            reject(req);
          }
        }
      });

      socket.send(data, 8080, currentServer.addr, (err) => {
        if (err) {
          socket.close();
          reject(err);
        }
      });
    } else {
      // Called from command (not using image API)
      const worker = new Worker(path.join(__dirname, "image-runner.js"), {
        workerData: object
      });
      worker.on("message", (data) => {
        resolve({
          buffer: Buffer.from([...data.buffer]),
          type: data.fileExtension
        });
      });
      worker.on("error", reject);
    }
  });
};