2026-08-25 18:43:04 +08:00

422 lines
12 KiB
JavaScript
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

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,
},
};