ファイルと画像認識処理の改善

.imgbotconfig

@@ -0,0 +1,5 @@
+{
+	"ignoredFiles": [
+		"test/resources/*"
+	]
+}

package.json

@@ -180,6 +180,7 @@
 		"portscanner": "2.2.0",
 		"postcss-loader": "3.0.0",
 		"prismjs": "1.18.0",
+		"probe-image-size": "5.0.0",
 		"progress-bar-webpack-plugin": "1.12.1",
 		"promise-limit": "2.7.0",
 		"promise-sequential": "1.1.1",

src/@types/probe-image-size.d.ts

@@ -0,0 +1,27 @@
+declare module 'probe-image-size' {
+	import { ReadStream } from 'fs';
+
+	type ProbeOptions = {
+		retries: 1;
+		timeout: 30000;
+	};
+
+	type ProbeResult = {
+		width: number;
+		height: number;
+		length?: number;
+		type: string;
+		mime: string;
+		wUnits: 'in' | 'mm' | 'cm' | 'pt' | 'pc' | 'px' | 'em' | 'ex';
+		hUnits: 'in' | 'mm' | 'cm' | 'pt' | 'pc' | 'px' | 'em' | 'ex';
+		url?: string;
+	};
+
+	function probeImageSize(src: string | ReadStream, options?: ProbeOptions): Promise<ProbeResult>;
+	function probeImageSize(src: string | ReadStream, callback: (err: Error | null, result?: ProbeResult) => void): void;
+	function probeImageSize(src: string | ReadStream, options: ProbeOptions, callback: (err: Error | null, result?: ProbeResult) => void): void;
+
+	namespace probeImageSize {} // Hack
+
+	export = probeImageSize;
+}

src/misc/detect-url-mine.ts → src/misc/detect-url-mime.ts

@@ -1,14 +1,14 @@
 import { createTemp } from './create-temp';
 import { downloadUrl } from './donwload-url';
-import { detectMine } from './detect-mine';
+import { detectType } from './get-file-info';
 
-export async function detectUrlMine(url: string) {
+export async function detectUrlMime(url: string) {
 	const [path, cleanup] = await createTemp();
 
 	try {
 		await downloadUrl(url, path);
-		const [type] = await detectMine(path);
-		return type;
+		const { mime } = await detectType(path);
+		return mime;
 	} finally {
 		cleanup();
 	}

src/misc/get-file-info.ts

@@ -0,0 +1,201 @@
+import * as fs from 'fs';
+import * as crypto from 'crypto';
+import * as fileType from 'file-type';
+import isSvg from 'is-svg';
+import * as probeImageSize from 'probe-image-size';
+import * as sharp from 'sharp';
+
+export type FileInfo = {
+	size: number;
+	md5: string;
+	type: {
+		mime: string;
+		ext: string | null;
+	};
+	width?: number;
+	height?: number;
+	avgColor?: number[];
+	warnings: string[];
+};
+
+const TYPE_OCTET_STREAM = {
+	mime: 'application/octet-stream',
+	ext: null
+};
+
+const TYPE_SVG = {
+	mime: 'image/svg+xml',
+	ext: 'svg'
+};
+
+/**
+ * Get file information
+ */
+export async function getFileInfo(path: string): Promise<FileInfo> {
+	const warnings = [] as string[];
+
+	const size = await getFileSize(path);
+	const md5 = await calcHash(path);
+
+	let type = await detectType(path);
+
+	// image dimensions
+	let width: number | undefined;
+	let height: number | undefined;
+
+	if (['image/jpeg', 'image/gif', 'image/png', 'image/apng', 'image/webp', 'image/bmp', 'image/tiff', 'image/svg+xml', 'image/vnd.adobe.photoshop'].includes(type.mime)) {
+		const imageSize = await detectImageSize(path).catch(e => {
+			warnings.push(`detectImageSize failed: ${e}`);
+			return undefined;
+		});
+
+		// うまく判定できない画像は octet-stream にする
+		if (!imageSize) {
+			warnings.push(`cannot detect image dimensions`);
+			type = TYPE_OCTET_STREAM;
+		} else if (imageSize.wUnits === 'px') {
+			width = imageSize.width;
+			height = imageSize.height;
+
+			// 制限を超えている画像は octet-stream にする
+			if (imageSize.width > 16383 || imageSize.height > 16383) {
+				warnings.push(`image dimensions exceeds limits`);
+				type = TYPE_OCTET_STREAM;
+			}
+		} else {
+			warnings.push(`unsupported unit type: ${imageSize.wUnits}`);
+		}
+	}
+
+	// average color
+	let avgColor: number[] | undefined;
+
+	if (['image/jpeg', 'image/gif', 'image/png', 'image/apng', 'image/webp', 'image/svg+xml'].includes(type.mime)) {
+		avgColor = await calcAvgColor(path).catch(e => {
+			warnings.push(`calcAvgColor failed: ${e}`);
+			return undefined;
+		});
+	}
+
+	return {
+		size,
+		md5,
+		type,
+		width,
+		height,
+		avgColor,
+		warnings,
+	};
+}
+
+/**
+ * Detect MIME Type and extension
+ */
+export async function detectType(path: string) {
+	// Check 0 byte
+	const fileSize = await getFileSize(path);
+	if (fileSize === 0) {
+		return TYPE_OCTET_STREAM;
+	}
+
+	const type = await fileType.fromFile(path);
+
+	if (type) {
+		// XMLはSVGかもしれない
+		if (type.mime === 'application/xml' && await checkSvg(path)) {
+			return TYPE_SVG;
+		}
+
+		return {
+			mime: type.mime,
+			ext: type.ext
+		};
+	}
+
+	// 種類が不明でもSVGかもしれない
+	if (await checkSvg(path)) {
+		return TYPE_SVG;
+	}
+
+	// それでも種類が不明なら application/octet-stream にする
+	return TYPE_OCTET_STREAM;
+}
+
+/**
+ * Check the file is SVG or not
+ */
+export async function checkSvg(path: string) {
+	try {
+		const size = await getFileSize(path);
+		if (size > 1 * 1024 * 1024) return false;
+		return isSvg(fs.readFileSync(path));
+	} catch {
+		return false;
+	}
+}
+
+/**
+ * Get file size
+ */
+export async function getFileSize(path: string): Promise<number> {
+	return new Promise<number>((res, rej) => {
+		fs.stat(path, (err, stats) => {
+			if (err) return rej(err);
+			res(stats.size);
+		});
+	});
+}
+
+/**
+ * Calculate MD5 hash
+ */
+async function calcHash(path: string): Promise<string> {
+	return new Promise<string>((res, rej) => {
+		const readable = fs.createReadStream(path);
+		const hash = crypto.createHash('md5');
+		const chunks: Buffer[] = [];
+		readable
+			.on('error', rej)
+			.pipe(hash)
+			.on('error', rej)
+			.on('data', chunk => chunks.push(chunk))
+			.on('end', () => {
+				const buffer = Buffer.concat(chunks);
+				res(buffer.toString('hex'));
+			});
+	});
+}
+
+/**
+ * Detect dimensions of image
+ */
+async function detectImageSize(path: string): Promise<{
+	width: number;
+	height: number;
+	wUnits: string;
+	hUnits: string;
+}> {
+	const readable = fs.createReadStream(path);
+	const imageSize = await probeImageSize(readable);
+	readable.destroy();
+	return imageSize;
+}
+
+/**
+ * Calculate average color of image
+ */
+async function calcAvgColor(path: string): Promise<number[]> {
+	const img = sharp(path);
+
+	const info = await (img as any).stats();
+
+	if (info.isOpaque) {
+		const r = Math.round(info.channels[0].mean);
+		const g = Math.round(info.channels[1].mean);
+		const b = Math.round(info.channels[2].mean);
+
+		return [r, g, b];
+	} else {
+		return [255, 255, 255];
+	}
+}

src/server/api/endpoints/admin/emoji/add.ts

@@ -1,6 +1,6 @@
 import $ from 'cafy';
 import define from '../../../define';
-import { detectUrlMine } from '../../../../../misc/detect-url-mine';
+import { detectUrlMime } from '../../../../../misc/detect-url-mime';
 import { Emojis } from '../../../../../models';
 import { genId } from '../../../../../misc/gen-id';
 import { getConnection } from 'typeorm';
@@ -46,7 +46,7 @@ export const meta = {
 };
 
 export default define(meta, async (ps, me) => {
-	const type = await detectUrlMine(ps.url);
+	const type = await detectUrlMime(ps.url);
 
 	const exists = await Emojis.findOne({
 		name: ps.name,

src/server/api/endpoints/admin/emoji/update.ts

@@ -1,6 +1,6 @@
 import $ from 'cafy';
 import define from '../../../define';
-import { detectUrlMine } from '../../../../../misc/detect-url-mine';
+import { detectUrlMime } from '../../../../../misc/detect-url-mime';
 import { ID } from '../../../../../misc/cafy-id';
 import { Emojis } from '../../../../../models';
 import { getConnection } from 'typeorm';
@@ -52,7 +52,7 @@ export default define(meta, async (ps) => {
 
 	if (emoji == null) throw new ApiError(meta.errors.noSuchEmoji);
 
-	const type = await detectUrlMine(ps.url);
+	const type = await detectUrlMime(ps.url);
 
 	await Emojis.update(emoji.id, {
 		updatedAt: new Date(),

src/server/file/send-drive-file.ts

@@ -8,7 +8,7 @@ import { contentDisposition } from '../../misc/content-disposition';
 import { DriveFiles } from '../../models';
 import { InternalStorage } from '../../services/drive/internal-storage';
 import { downloadUrl } from '../../misc/donwload-url';
-import { detectMine } from '../../misc/detect-mine';
+import { detectType } from '../../misc/get-file-info';
 import { convertToJpeg, convertToPng } from '../../services/drive/image-processor';
 import { GenerateVideoThumbnail } from '../../services/drive/generate-video-thumbnail';
 
@@ -52,23 +52,23 @@ export default async function(ctx: Koa.Context) {
 			try {
 				await downloadUrl(file.uri, path);
 
-				const [type, ext] = await detectMine(path);
+				const { mime, ext } = await detectType(path);
 
 				const convertFile = async () => {
 					if (isThumbnail) {
-						if (['image/jpeg', 'image/webp'].includes(type)) {
+						if (['image/jpeg', 'image/webp'].includes(mime)) {
 							return await convertToJpeg(path, 498, 280);
-						} else if (['image/png'].includes(type)) {
+						} else if (['image/png'].includes(mime)) {
 							return await convertToPng(path, 498, 280);
-						} else if (type.startsWith('video/')) {
+						} else if (mime.startsWith('video/')) {
 							return await GenerateVideoThumbnail(path);
 						}
 					}
 
 					return {
 						data: fs.readFileSync(path),
 						ext,
-						type,
+						type: mime,
 					};
 				};