Refactor and some fixes

src/server/api/endpoints/drive/files/create.ts

@@ -1,6 +1,7 @@
 /**
  * Module dependencies
  */
+import * as fs from 'fs';
 import $ from 'cafy'; import ID from '../../../../../cafy-id';
 import { validateFileName, pack } from '../../../../../models/drive-file';
 import create from '../../../../../services/drive/add-file';
@@ -32,15 +33,23 @@ module.exports = async (file, params, user): Promise<any> => {
 	const [folderId = null, folderIdErr] = $.type(ID).optional().nullable().get(params.folderId);
 	if (folderIdErr) throw 'invalid folderId param';
 
+	function cleanup() {
+		fs.unlink(file.path, () => {});
+	}
+
 	try {
 		// Create file
 		const driveFile = await create(user, file.path, name, null, folderId);
 
+		cleanup();
+
 		// Serialize
 		return pack(driveFile);
 	} catch (e) {
 		console.error(e);
 
+		cleanup();
+
 		throw e;
 	}
 };

src/services/drive/add-file.ts

@@ -1,6 +1,5 @@
 import { Buffer } from 'buffer';
 import * as fs from 'fs';
-import * as tmp from 'tmp';
 import * as stream from 'stream';
 
 import * as mongodb from 'mongodb';
@@ -14,8 +13,7 @@ import DriveFile, { IMetadata, getDriveFileBucket, IDriveFile, DriveFileChunk }
 import DriveFolder from '../../models/drive-folder';
 import { pack } from '../../models/drive-file';
 import event, { publishDriveStream } from '../../publishers/stream';
-import getAcct from '../../acct/render';
-import { IUser, isLocalUser, isRemoteUser } from '../../models/user';
+import { isLocalUser, IRemoteUser } from '../../models/user';
 import DriveFileThumbnail, { getDriveFileThumbnailBucket, DriveFileThumbnailChunk } from '../../models/drive-file-thumbnail';
 import genThumbnail from '../../drive/gen-thumbnail';
 
@@ -25,13 +23,6 @@ const gm = _gm.subClass({
 
 const log = debug('misskey:drive:add-file');
 
-const tmpFile = (): Promise<[string, any]> => new Promise((resolve, reject) => {
-	tmp.file((e, path, fd, cleanup) => {
-		if (e) return reject(e);
-		resolve([path, cleanup]);
-	});
-});
-
 const writeChunks = (name: string, readable: stream.Readable, type: string, metadata: any) =>
 	getDriveFileBucket()
 		.then(bucket => new Promise((resolve, reject) => {
@@ -55,8 +46,59 @@ const writeThumbnailChunks = (name: string, readable: stream.Readable, originalI
 			readable.pipe(writeStream);
 		}));
 
-const addFile = async (
-	user: IUser,
+async function deleteOldFile(user: IRemoteUser) {
+	const oldFile = await DriveFile.findOne({
+		_id: {
+			$nin: [user.avatarId, user.bannerId]
+		}
+	}, {
+		sort: {
+			_id: 1
+		}
+	});
+
+	if (oldFile) {
+		// チャンクをすべて削除
+		DriveFileChunk.remove({
+			files_id: oldFile._id
+		});
+
+		DriveFile.update({ _id: oldFile._id }, {
+			$set: {
+				'metadata.deletedAt': new Date(),
+				'metadata.isExpired': true
+			}
+		});
+
+		//#region サムネイルもあれば削除
+		const thumbnail = await DriveFileThumbnail.findOne({
+			'metadata.originalId': oldFile._id
+		});
+
+		if (thumbnail) {
+			DriveFileThumbnailChunk.remove({
+				files_id: thumbnail._id
+			});
+
+			DriveFileThumbnail.remove({ _id: thumbnail._id });
+		}
+		//#endregion
+	}
+}
+
+/**
+ * Add file to drive
+ *
+ * @param user User who wish to add file
+ * @param path File path
+ * @param name Name
+ * @param comment Comment
+ * @param folderId Folder ID
+ * @param force If set to true, forcibly upload the file even if there is a file with the same hash.
+ * @return Created drive file
+ */
+export default async function(
+	user: any,
 	path: string,
 	name: string = null,
 	comment: string = null,
@@ -64,55 +106,54 @@ const addFile = async (
 	force: boolean = false,
 	url: string = null,
 	uri: string = null
-): Promise<IDriveFile> => {
-	log(`registering ${name} (user: ${getAcct(user)}, path: ${path})`);
-
-	// Calculate hash, get content type and get file size
-	const [hash, [mime, ext], size] = await Promise.all([
-		// hash
-		((): Promise<string> => new Promise((res, rej) => {
-			const readable = fs.createReadStream(path);
-			const hash = crypto.createHash('md5');
-			const chunks = [];
-			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'));
-				});
-		}))(),
-		// mime
-		((): Promise<[string, string | null]> => new Promise((res, rej) => {
-			const readable = fs.createReadStream(path);
-			readable
-				.on('error', rej)
-				.once('data', (buffer: Buffer) => {
-					readable.destroy();
-					const type = fileType(buffer);
-					if (type) {
-						return res([type.mime, type.ext]);
-					} else {
-						// 種類が同定できなかったら application/octet-stream にする
-						return res(['application/octet-stream', null]);
-					}
-				});
-		}))(),
-		// size
-		((): Promise<number> => new Promise((res, rej) => {
-			fs.stat(path, (err, stats) => {
-				if (err) return rej(err);
-				res(stats.size);
+): Promise<IDriveFile> {
+	// Calc md5 hash
+	const calcHash = new Promise<string>((res, rej) => {
+		const readable = fs.createReadStream(path);
+		const hash = crypto.createHash('md5');
+		const chunks = [];
+		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 content type
+	const detectMime = new Promise<[string, string]>((res, rej) => {
+		const readable = fs.createReadStream(path);
+		readable
+			.on('error', rej)
+			.once('data', (buffer: Buffer) => {
+				readable.destroy();
+				const type = fileType(buffer);
+				if (type) {
+					res([type.mime, type.ext]);
+				} else {
+					// 種類が同定できなかったら application/octet-stream にする
+					res(['application/octet-stream', null]);
+				}
+			});
+	});
+
+	// Get file size
+	const getFileSize = new Promise<number>((res, rej) => {
+		fs.stat(path, (err, stats) => {
+			if (err) return rej(err);
+			res(stats.size);
+		});
+	});
+
+	const [hash, [mime, ext], size] = await Promise.all([calcHash, detectMime, getFileSize]);
 
 	log(`hash: ${hash}, mime: ${mime}, ext: ${ext}, size: ${size}`);
 
 	// detect name
-	const detectedName: string = name || (ext ? `untitled.${ext}` : 'untitled');
+	const detectedName = name || (ext ? `untitled.${ext}` : 'untitled');
 
 	if (!force) {
 		// Check if there is a file with the same hash
@@ -125,26 +166,70 @@ const addFile = async (
 		if (much !== null) {
 			log('file with same hash is found');
 			return much;
-		} else {
-			log('file with same hash is not found');
 		}
 	}
 
-	const [wh, averageColor, folder] = await Promise.all([
-		// Width and height (when image)
-		(async () => {
-			// 画像かどうか
-			if (!/^image\/.*$/.test(mime)) {
-				return null;
+	//#region Check drive usage
+	const usage = await DriveFile
+		.aggregate([{
+			$match: {
+				'metadata.userId': user._id,
+				'metadata.deletedAt': { $exists: false }
 			}
-
-			const imageType = mime.split('/')[1];
-
-			// 画像でもPNGかJPEGかGIFでないならスキップ
-			if (imageType != 'png' && imageType != 'jpeg' && imageType != 'gif') {
-				return null;
+		}, {
+			$project: {
+				length: true
 			}
+		}, {
+			$group: {
+				_id: null,
+				usage: { $sum: '$length' }
+			}
+		}])
+		.then((aggregates: any[]) => {
+			if (aggregates.length > 0) {
+				return aggregates[0].usage;
+			}
+			return 0;
+		});
 
+	log(`drive usage is ${usage}`);
+
+	// If usage limit exceeded
+	if (usage + size > user.driveCapacity) {
+		if (isLocalUser(user)) {
+			throw 'no-free-space';
+		} else {
+			// (アバターまたはバナーを含まず)最も古いファイルを削除する
+			deleteOldFile(user);
+		}
+	}
+	//#endregion
+
+	const fetchFolder = async () => {
+		if (!folderId) {
+			return null;
+		}
+
+		const driveFolder = await DriveFolder.findOne({
+			_id: folderId,
+			userId: user._id
+		});
+
+		if (driveFolder == null) throw 'folder-not-found';
+
+		return driveFolder;
+	};
+
+	const properties = {};
+
+	let propPromises = [];
+
+	const isImage = ['image/jpeg', 'image/gif', 'image/png'].includes(mime);
+
+	if (isImage) {
+		// Calc width and height
+		const calcWh = async () => {
 			log('calculate image width and height...');
 
 			// Calculate width and height
@@ -153,22 +238,12 @@ const addFile = async (
 
 			log(`image width and height is calculated: ${size.width}, ${size.height}`);
 
-			return [size.width, size.height];
-		})(),
-		// average color (when image)
-		(async () => {
-			// 画像かどうか
-			if (!/^image\/.*$/.test(mime)) {
-				return null;
-			}
-
-			const imageType = mime.split('/')[1];
-
-			// 画像でもPNGかJPEGでないならスキップ
-			if (imageType != 'png' && imageType != 'jpeg') {
-				return null;
-			}
+			properties['width'] = size.width;
+			properties['height'] = size.height;
+		};
 
+		// Calc average color
+		const calcAvg = async () => {
 			log('calculate average color...');
 
 			const info = await prominence(gm(fs.createReadStream(path), name)).identify();
@@ -185,112 +260,18 @@ const addFile = async (
 
 			log(`average color is calculated: ${r}, ${g}, ${b}`);
 
-			return isTransparent ? [r, g, b, 255] : [r, g, b];
-		})(),
-		// folder
-		(async () => {
-			if (!folderId) {
-				return null;
-			}
-			const driveFolder = await DriveFolder.findOne({
-				_id: folderId,
-				userId: user._id
-			});
-			if (!driveFolder) {
-				throw 'folder-not-found';
-			}
-			return driveFolder;
-		})(),
-		// usage checker
-		(async () => {
-			// Calculate drive usage
-			const usage = await DriveFile
-				.aggregate([{
-					$match: {
-						'metadata.userId': user._id,
-						'metadata.deletedAt': { $exists: false }
-					}
-				}, {
-					$project: {
-						length: true
-					}
-				}, {
-					$group: {
-						_id: null,
-						usage: { $sum: '$length' }
-					}
-				}])
-				.then((aggregates: any[]) => {
-					if (aggregates.length > 0) {
-						return aggregates[0].usage;
-					}
-					return 0;
-				});
+			const value = isTransparent ? [r, g, b, 255] : [r, g, b];
 
-			log(`drive usage is ${usage}`);
+			properties['avgColor'] = value;
+		};
 
-			// If usage limit exceeded
-			if (usage + size > user.driveCapacity) {
-				if (isLocalUser(user)) {
-					throw 'no-free-space';
-				} else {
-					//#region (アバターまたはバナーを含まず)最も古いファイルを削除する
-					const oldFile = await DriveFile.findOne({
-						_id: {
-							$nin: [user.avatarId, user.bannerId]
-						}
-					}, {
-						sort: {
-							_id: 1
-						}
-					});
+		propPromises = [calcWh(), calcAvg()];
+	}
 
-					if (oldFile) {
-						// チャンクをすべて削除
-						DriveFileChunk.remove({
-							files_id: oldFile._id
-						});
-
-						DriveFile.update({ _id: oldFile._id }, {
-							$set: {
-								'metadata.deletedAt': new Date(),
-								'metadata.isExpired': true
-							}
-						});
-
-						//#region サムネイルもあれば削除
-						const thumbnail = await DriveFileThumbnail.findOne({
-							'metadata.originalId': oldFile._id
-						});
-
-						if (thumbnail) {
-							DriveFileThumbnailChunk.remove({
-								files_id: thumbnail._id
-							});
-
-							DriveFileThumbnail.remove({ _id: thumbnail._id });
-						}
-						//#endregion
-					}
-					//#endregion
-				}
-			}
-		})()
-	]);
+	const [folder] = await Promise.all([fetchFolder(), propPromises]);
 
 	const readable = fs.createReadStream(path);
 
-	const properties = {};
-
-	if (wh) {
-		properties['width'] = wh[0];
-		properties['height'] = wh[1];
-	}
-
-	if (averageColor) {
-		properties['avgColor'] = averageColor;
-	}
-
 	const metadata = {
 		userId: user._id,
 		_user: {
@@ -309,74 +290,24 @@ const addFile = async (
 		metadata.uri = uri;
 	}
 
-	const file = await (writeChunks(detectedName, readable, mime, metadata) as Promise<IDriveFile>);
+	const driveFile = await (writeChunks(detectedName, readable, mime, metadata) as Promise<IDriveFile>);
+
+	log(`drive file has been created ${driveFile._id}`);
+
+	pack(driveFile).then(packedFile => {
+		// Publish drive_file_created event
+		event(user._id, 'drive_file_created', packedFile);
+		publishDriveStream(user._id, 'file_created', packedFile);
+	});
 
 	try {
-		const thumb = await genThumbnail(file);
+		const thumb = await genThumbnail(driveFile);
 		if (thumb) {
-			await writeThumbnailChunks(detectedName, thumb, file._id);
+			await writeThumbnailChunks(detectedName, thumb, driveFile._id);
 		}
 	} catch (e) {
 		// noop
 	}
 
-	return file;
-};
-
-/**
- * Add file to drive
- *
- * @param user User who wish to add file
- * @param file File path or readableStream
- * @param comment Comment
- * @param type File type
- * @param folderId Folder ID
- * @param force If set to true, forcibly upload the file even if there is a file with the same hash.
- * @return Object that represents added file
- */
-export default (user: any, file: string | stream.Readable, ...args) => new Promise<any>((resolve, reject) => {
-	const isStream = typeof file === 'object' && typeof file.read === 'function';
-
-	// Get file path
-	new Promise<[string, any]>((res, rej) => {
-		if (typeof file === 'string') {
-			res([file, null]);
-		} else if (isStream) {
-			tmpFile()
-				.then(([path, cleanup]) => {
-					const readable: stream.Readable = file;
-					const writable = fs.createWriteStream(path);
-					readable
-						.on('error', rej)
-						.on('end', () => {
-							res([path, cleanup]);
-						})
-						.pipe(writable)
-						.on('error', rej);
-				})
-				.catch(rej);
-		} else {
-			rej(new Error('un-compatible file.'));
-		}
-	})
-	.then(([path, cleanup]) => new Promise<IDriveFile>((res, rej) => {
-		addFile(user, path, ...args)
-			.then(file => {
-				res(file);
-				if (cleanup) cleanup();
-			})
-			.catch(rej);
-	}))
-	.then(file => {
-		log(`drive file has been created ${file._id}`);
-
-		resolve(file);
-
-		pack(file).then(packedFile => {
-			// Publish drive_file_created event
-			event(user._id, 'drive_file_created', packedFile);
-			publishDriveStream(user._id, 'file_created', packedFile);
-		});
-	})
-	.catch(reject);
-});
+	return driveFile;
+}