hn-hlw-service/hlw/im-image-archive/index.js

const axios = require("axios");
const crypto = require("crypto");
const fs = require("fs");
const fsp = require("fs/promises");
const path = require("path");
const { Transform } = require("stream");
const { pipeline } = require("stream/promises");
const logger = require("../../utils/logger");

const DEFAULT_MAX_FILE_SIZE = 30 * 1024 * 1024;
const RETRY_DELAYS = [60 * 1000, 5 * 60 * 1000, 15 * 60 * 1000];

function resolveDefaultImageRoot(
  moduleDirectory = __dirname,
  isBuilt = process.env.IS_BUILDED === "YES"
) {
  // 源码运行时 __dirname 在 hlw/im-image-archive，需要先回到服务根目录；
  // esbuild 打包后所有代码共用 bundle 所在目录，可直接将其作为服务根目录。
  const serviceRoot = isBuilt
    ? moduleDirectory
    : path.resolve(moduleDirectory, "../..");
  return path.resolve(serviceRoot, "../private-files/im-chat");
}

function getImageRoot() {
  return path.resolve(
    process.env.CONFIG_IM_IMAGE_ROOT || resolveDefaultImageRoot()
  );
}

async function initializeStorage() {
  const imageRoot = getImageRoot();
  await fsp.mkdir(imageRoot, { recursive: true });
  await fsp.access(imageRoot, fs.constants.R_OK | fs.constants.W_OK);
  logger.info("腾讯 IM 图片归档目录已就绪", imageRoot);
  return imageRoot;
}

function getMaxFileSize() {
  const configured = Number(process.env.CONFIG_IM_IMAGE_MAX_BYTES);
  return Number.isFinite(configured) && configured > 0
    ? configured
    : DEFAULT_MAX_FILE_SIZE;
}

function getAllowedHostSuffixes() {
  return (process.env.CONFIG_IM_IMAGE_ALLOWED_HOSTS ||
    ".my-imcloud.com,.im.qcloud.com,.myqcloud.com")
    .split(",")
    .map((item) => item.trim().toLowerCase())
    .filter(Boolean);
}

function findOriginalImage(message) {
  if (!message || !Array.isArray(message.MsgBody)) return null;

  for (const element of message.MsgBody) {
    if (element && element.MsgType === "TIMImageElem") {
      const content = element.MsgContent || {};
      const imageList = Array.isArray(content.ImageInfoArray)
        ? content.ImageInfoArray
        : [];
      const image =
        imageList.find((item) => Number(item && item.Type) === 1 && item.URL) ||
        imageList.find((item) => item && item.URL);
      if (image) {
        return {
          uuid: String(content.UUID || ""),
          imageFormat: content.ImageFormat,
          type: Number(image.Type) || 1,
          url: String(image.URL),
          expectedSize: Number(image.Size) || 0,
          width: Number(image.Width) || 0,
          height: Number(image.Height) || 0,
        };
      }
    }
  }

  return null;
}

function prepareMessageForArchive(message) {
  const source = findOriginalImage(message);
  if (!source) return message;

  const now = Date.now();
  return {
    ...message,
    imageArchive: {
      status: "pending",
      sourceUUID: source.uuid,
      sourceType: source.type,
      sourceURL: source.url,
      expectedSize: source.expectedSize,
      width: source.width,
      height: source.height,
      relativePath: "",
      size: 0,
      mimeType: "",
      sha256: "",
      retryCount: 0,
      lastError: "",
      createdAt: now,
      updatedAt: now,
      archivedAt: null,
    },
  };
}

function isAllowedSourceUrl(value) {
  let parsed;
  try {
    parsed = new URL(value);
  } catch (_error) {
    return false;
  }

  if (parsed.protocol !== "https:") return false;
  const hostname = parsed.hostname.toLowerCase();
  return getAllowedHostSuffixes().some((suffix) => {
    const normalized = suffix.startsWith(".") ? suffix : `.${suffix}`;
    return hostname === normalized.slice(1) || hostname.endsWith(normalized);
  });
}

function safePathPart(value, fallback) {
  const normalized = String(value || "")
    .replace(/[^a-zA-Z0-9._-]/g, "_")
    .slice(0, 160);
  return normalized || fallback;
}

function getMessageDate(message) {
  const seconds = Number(message && message.MsgTime);
  const date = Number.isFinite(seconds) && seconds > 0
    ? new Date(seconds * 1000)
    : new Date();
  return Number.isNaN(date.getTime()) ? new Date() : date;
}

function detectImageType(buffer) {
  if (!buffer || buffer.length < 4) return null;
  if (buffer[0] === 0xff && buffer[1] === 0xd8 && buffer[2] === 0xff) {
    return { extension: ".jpg", mimeType: "image/jpeg" };
  }
  if (buffer.subarray(0, 8).equals(Buffer.from("89504e470d0a1a0a", "hex"))) {
    return { extension: ".png", mimeType: "image/png" };
  }
  const signature = buffer.subarray(0, 6).toString("ascii");
  if (signature === "GIF87a" || signature === "GIF89a") {
    return { extension: ".gif", mimeType: "image/gif" };
  }
  if (buffer.subarray(0, 2).toString("ascii") === "BM") {
    return { extension: ".bmp", mimeType: "image/bmp" };
  }
  if (
    buffer.subarray(0, 4).toString("ascii") === "RIFF" &&
    buffer.subarray(8, 12).toString("ascii") === "WEBP"
  ) {
    return { extension: ".webp", mimeType: "image/webp" };
  }
  if (buffer.length >= 12 && buffer.subarray(4, 8).toString("ascii") === "ftyp") {
    const brand = buffer.subarray(8, 12).toString("ascii").toLowerCase();
    if (["heic", "heix", "hevc", "hevx", "mif1", "msf1"].includes(brand)) {
      return { extension: ".heic", mimeType: "image/heic" };
    }
    if (["avif", "avis"].includes(brand)) {
      return { extension: ".avif", mimeType: "image/avif" };
    }
  }
  return null;
}

async function downloadImage(sourceUrl, directory, fileBase) {
  if (!isAllowedSourceUrl(sourceUrl)) {
    throw new Error("图片下载地址不是允许的腾讯 IM HTTPS 域名");
  }

  await fsp.mkdir(directory, { recursive: true });
  const temporaryPath = path.join(
    directory,
    `${fileBase}.${process.pid}.${Date.now()}.part`
  );
  const maxFileSize = getMaxFileSize();
  let downloadedSize = 0;
  const hash = crypto.createHash("sha256");

  try {
    const response = await axios.get(sourceUrl, {
      responseType: "stream",
      timeout: 20000,
      maxRedirects: 0,
      validateStatus: (status) => status === 200,
    });

    const contentLength = Number(response.headers["content-length"] || 0);
    if (contentLength > maxFileSize) {
      response.data.destroy();
      throw new Error(`图片超过最大限制 ${maxFileSize} 字节`);
    }

    const limiter = new Transform({
      transform(chunk, _encoding, callback) {
        downloadedSize += chunk.length;
        if (downloadedSize > maxFileSize) {
          callback(new Error(`图片超过最大限制 ${maxFileSize} 字节`));
          return;
        }
        hash.update(chunk);
        callback(null, chunk);
      },
    });

    await pipeline(
      response.data,
      limiter,
      fs.createWriteStream(temporaryPath, { flags: "wx" })
    );

    const fileHandle = await fsp.open(temporaryPath, "r");
    const header = Buffer.alloc(16);
    try {
      await fileHandle.read(header, 0, header.length, 0);
    } finally {
      await fileHandle.close();
    }

    const imageType = detectImageType(header);
    if (!imageType) {
      throw new Error("下载内容不是支持的图片格式");
    }

    const finalPath = path.join(directory, `${fileBase}${imageType.extension}`);
    try {
      await fsp.access(finalPath, fs.constants.F_OK);
      await fsp.unlink(temporaryPath);
    } catch (_notFound) {
      await fsp.rename(temporaryPath, finalPath);
    }

    return {
      filePath: finalPath,
      size: downloadedSize,
      mimeType: imageType.mimeType,
      sha256: hash.digest("hex"),
    };
  } catch (error) {
    await fsp.unlink(temporaryPath).catch(() => {});
    throw error;
  }
}

async function archiveDocument(db, documentId) {
  const collection = db.collection("im-chat-msg");
  const document = await collection.findOne({ _id: documentId });
  if (!document || !document.imageArchive) return;
  if (document.imageArchive.status === "success") return;

  const claimed = await collection.updateOne(
    {
      _id: documentId,
      "imageArchive.status": { $in: ["pending", "retry", "processing"] },
    },
    {
      $set: {
        "imageArchive.status": "processing",
        "imageArchive.updatedAt": Date.now(),
      },
    }
  );
  if (!claimed.modifiedCount) return;

  try {
    const date = getMessageDate(document);
    const year = String(date.getFullYear());
    const month = String(date.getMonth() + 1).padStart(2, "0");
    const day = String(date.getDate()).padStart(2, "0");
    const msgId = safePathPart(document.MsgId, String(documentId));
    const fileBase = crypto
      .createHash("sha256")
      .update(document.imageArchive.sourceUUID || msgId)
      .digest("hex");
    const directory = path.join(getImageRoot(), year, month, day, msgId);
    const result = await downloadImage(
      document.imageArchive.sourceURL,
      directory,
      fileBase
    );
    const relativePath = path
      .relative(getImageRoot(), result.filePath)
      .replace(/\\/g, "/");

    await collection.updateOne(
      { _id: documentId },
      {
        $set: {
          "imageArchive.status": "success",
          "imageArchive.relativePath": relativePath,
          "imageArchive.size": result.size,
          "imageArchive.mimeType": result.mimeType,
          "imageArchive.sha256": result.sha256,
          "imageArchive.archivedAt": Date.now(),
          "imageArchive.updatedAt": Date.now(),
          "imageArchive.lastError": "",
        },
      }
    );
  } catch (error) {
    await collection.updateOne(
      { _id: documentId },
      {
        $set: {
          "imageArchive.status": "retry",
          "imageArchive.updatedAt": Date.now(),
          "imageArchive.lastError": String(error.message || error).slice(0, 500),
        },
        $inc: { "imageArchive.retryCount": 1 },
      }
    );
    throw error;
  }
}

function enqueueArchive(db, documentId, attempt = 0) {
  const run = () => {
    archiveDocument(db, documentId).catch((error) => {
      logger.error("腾讯 IM 图片归档失败", {
        documentId: String(documentId),
        attempt: attempt + 1,
        error: error.message,
      });
      if (attempt < RETRY_DELAYS.length) {
        setTimeout(
          () => enqueueArchive(db, documentId, attempt + 1),
          RETRY_DELAYS[attempt]
        );
      } else {
        db.collection("im-chat-msg")
          .updateOne(
            { _id: documentId, "imageArchive.status": "retry" },
            {
              $set: {
                "imageArchive.status": "failed",
                "imageArchive.updatedAt": Date.now(),
              },
            }
          )
          .catch((updateError) => {
            logger.error("更新腾讯 IM 图片归档最终失败状态失败", updateError);
          });
      }
    });
  };
  setImmediate(run);
}

async function resumePendingArchives(db) {
  const documents = await db
    .collection("im-chat-msg")
    .find(
      { "imageArchive.status": { $in: ["pending", "retry", "processing"] } },
      { projection: { _id: 1 } }
    )
    .limit(500)
    .toArray();
  documents.forEach((document) => enqueueArchive(db, document._id));
  return documents.length;
}

async function sendArchivedImage(req, res, db) {
  const document = await db.collection("im-chat-msg").findOne(
    { MsgId: req.params.msgId, "imageArchive.status": "success" },
    { projection: { imageArchive: 1 } }
  );
  if (!document || !document.imageArchive || !document.imageArchive.relativePath) {
    res.status(404).json({ success: false, message: "图片不存在或尚未归档" });
    return;
  }

  const root = getImageRoot();
  const filePath = path.resolve(root, document.imageArchive.relativePath);
  if (filePath !== root && !filePath.startsWith(`${root}${path.sep}`)) {
    res.status(400).json({ success: false, message: "图片路径无效" });
    return;
  }

  try {
    await fsp.access(filePath, fs.constants.R_OK);
  } catch (_error) {
    res.status(404).json({ success: false, message: "图片文件不存在" });
    return;
  }

  res.setHeader(
    "Content-Type",
    document.imageArchive.mimeType || "application/octet-stream"
  );
  res.setHeader("Content-Disposition", "inline");
  res.setHeader("Cache-Control", "private, max-age=3600");
  await new Promise((resolve, reject) => {
    res.sendFile(filePath, (error) => (error ? reject(error) : resolve()));
  });
}

module.exports = {
  findOriginalImage,
  prepareMessageForArchive,
  initializeStorage,
  enqueueArchive,
  resumePendingArchives,
  sendArchivedImage,
  // Exported for focused unit tests.
  _test: {
    archiveDocument,
    detectImageType,
    isAllowedSourceUrl,
    getImageRoot,
    resolveDefaultImageRoot,
  },
};