Merge pull request #913 from syuilo/write-file-use-stream

ファイル書き込みをバッファを使用しない形にする

3 files changed, 237 insertions, 149 deletions

diff --git a/src/api/common/add-file-to-drive.ts b/src/api/common/add-file-to-drive.ts
index a96906d291..7defbc631a 100644
--- a/src/api/common/add-file-to-drive.ts
+++ b/src/api/common/add-file-to-drive.ts
@@ -9,73 +9,83 @@ import DriveFolder from '../models/drive-folder';
 import serialize from '../serializers/drive-file';
 import event from '../event';
 import config from '../../conf';
-import { Duplex } from 'stream';
+import { Buffer } from 'buffer';
+import * as fs from 'fs';
+import * as tmp from 'tmp';
+import * as stream from 'stream';
 
 const log = debug('misskey:register-drive-file');
 
-const addToGridFS = (name, binary, type, metadata): Promise<any> => new Promise(async (resolve, reject) => {
-	const dataStream = new Duplex();
-	dataStream.push(binary);
-	dataStream.push(null);
-
-	const bucket = await getGridFSBucket();
-	const writeStream = bucket.openUploadStream(name, { contentType: type, metadata });
-	writeStream.once('finish', (doc) => { resolve(doc); });
-	writeStream.on('error', reject);
-	dataStream.pipe(writeStream);
+const tmpFile = (): Promise<string> => new Promise((resolve, reject) => {
+	tmp.file((e, path) => {
+		if (e) return reject(e);
+		resolve(path);
+	});
 });
 
-/**
- * Add file to drive
- *
- * @param user User who wish to add file
- * @param fileName File name
- * @param data Contents
- * @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 (
+const addToGridFS = (name: string, readable: stream.Readable, type: string, metadata: any): Promise<any> =>
+	getGridFSBucket()
+		.then(bucket => new Promise((resolve, reject) => {
+			const writeStream = bucket.openUploadStream(name, { contentType: type, metadata });
+			writeStream.once('finish', (doc) => { resolve(doc); });
+			writeStream.on('error', reject);
+			readable.pipe(writeStream);
+		}));
+
+const addFile = async (
 	user: any,
-	data: Buffer,
+	path: string,
 	name: string = null,
 	comment: string = null,
 	folderId: mongodb.ObjectID = null,
 	force: boolean = false
-) => new Promise<any>(async (resolve, reject) => {
+) => {
 	log(`registering ${name} (user: ${user.username})`);
 
-	// File size
-	const size = data.byteLength;
-
-	log(`size is ${size}`);
-
-	// File type
-	let mime = 'application/octet-stream';
-	const type = fileType(data);
-	if (type !== null) {
-		mime = type.mime;
-
-		if (name === null) {
-			name = `untitled.${type.ext}`;
-		}
-	} else {
-		if (name === null) {
-			name = 'untitled';
-		}
-	}
-
-	log(`type is ${mime}`);
+	// 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(['application/octet-stream', null]);
+					}
+					return res([type.mime, type.ext]);
+				});
+		}))(),
+		// size
+		((): Promise<number> => new Promise((res, rej) => {
+			fs.stat(path, (err, stats) => {
+				if (err) return rej(err);
+				res(stats.size);
+			});
+		}))()
+	]);
 
-	// Generate hash
-	const hash = crypto
-		.createHash('md5')
-		.update(data)
-		.digest('hex') as string;
+	log(`hash: ${hash}, mime: ${mime}, ext: ${ext}, size: ${size}`);
 
-	log(`hash is ${hash}`);
+	// detect name
+	const detectedName: string = name || (ext ? `untitled.${ext}` : 'untitled');
 
 	if (!force) {
 		// Check if there is a file with the same hash
@@ -86,92 +96,155 @@ export default (
 
 		if (much !== null) {
 			log('file with same hash is found');
-			return resolve(much);
+			return much;
 		} else {
 			log('file with same hash is not found');
 		}
 	}
 
-	// Calculate drive usage
-	const usage = ((await DriveFile
-		.aggregate([
-			{ $match: { 'metadata.user_id': user._id } },
-			{ $project: {
-				length: true
-			}},
-			{ $group: {
-				_id: null,
-				usage: { $sum: '$length' }
-			}}
-		]))[0] || {
-			usage: 0
-		}).usage;
-
-	log(`drive usage is ${usage}`);
-
-	// If usage limit exceeded
-	if (usage + size > user.drive_capacity) {
-		return reject('no-free-space');
-	}
-
-	// If the folder is specified
-	let folder: any = null;
-	if (folderId !== null) {
-		folder = await DriveFolder
-			.findOne({
+	const [properties, folder] = await Promise.all([
+		// properties
+		(async () => {
+			if (!/^image\/.*$/.test(mime)) {
+				return null;
+			}
+			// If the file is an image, calculate width and height to save in property
+			const g = gm(fs.createReadStream(path), name);
+			const size = await prominence(g).size();
+			const properties = {
+				width: size.width,
+				height: size.height
+			};
+			log('image width and height is calculated');
+			return properties;
+		})(),
+		// folder
+		(async () => {
+			if (!folderId) {
+				return null;
+			}
+			const driveFolder = await DriveFolder.findOne({
 				_id: folderId,
 				user_id: user._id
 			});
+			if (!driveFolder) {
+				throw 'folder-not-found';
+			}
+			return driveFolder;
+		})(),
+		// usage checker
+		(async () => {
+			// Calculate drive usage
+			const usage = await DriveFile
+				.aggregate([
+					{ $match: { 'metadata.user_id': user._id } },
+					{
+						$project: {
+							length: true
+						}
+					},
+					{
+						$group: {
+							_id: null,
+							usage: { $sum: '$length' }
+						}
+					}
+				])
+				.then((aggregates: any[]) => {
+					if (aggregates.length > 0) {
+						return aggregates[0].usage;
+					}
+					return 0;
+				});
 
-		if (folder === null) {
-			return reject('folder-not-found');
-		}
-	}
+			log(`drive usage is ${usage}`);
 
-	let properties: any = null;
+			// If usage limit exceeded
+			if (usage + size > user.drive_capacity) {
+				throw 'no-free-space';
+			}
+		})()
+	]);
 
-	// If the file is an image
-	if (/^image\/.*$/.test(mime)) {
-		// Calculate width and height to save in property
-		const g = gm(data, name);
-		const size = await prominence(g).size();
-		properties = {
-			width: size.width,
-			height: size.height
-		};
+	const readable = fs.createReadStream(path);
 
-		log('image width and height is calculated');
-	}
-
-	// Create DriveFile document
-	const file = await addToGridFS(name, data, mime, {
+	return addToGridFS(detectedName, readable, mime, {
 		user_id: user._id,
 		folder_id: folder !== null ? folder._id : null,
 		comment: comment,
 		properties: properties
 	});
+};
 
-	log(`drive file has been created ${file._id}`);
-
-	resolve(file);
-
-	// Serialize
-	const fileObj = await serialize(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) => {
+	// Get file path
+	new Promise((res: (v: [string, boolean]) => void, rej) => {
+		if (typeof file === 'string') {
+			res([file, false]);
+			return;
+		}
+		if (typeof file === 'object' && typeof file.read === 'function') {
+			tmpFile()
+				.then(path => {
+					const readable: stream.Readable = file;
+					const writable = fs.createWriteStream(path);
+					readable
+						.on('error', rej)
+						.on('end', () => {
+							res([path, true]);
+						})
+						.pipe(writable)
+						.on('error', rej);
+				})
+				.catch(rej);
+		}
+		rej(new Error('un-compatible file.'));
+	}).then(([path, remove]): Promise<any> => new Promise((res, rej) => {
+		addFile(user, path, ...args)
+			.then(file => {
+				res(file);
+				if (remove) {
+					fs.unlink(path, (e) => {
+						if (e) log(e.stack);
+					});
+				}
+			})
+			.catch(rej);
+	}))
+		.then(file => {
+			log(`drive file has been created ${file._id}`);
+			resolve(file);
 
-	// Publish drive_file_created event
-	event(user._id, 'drive_file_created', fileObj);
+			serialize(file)
+				.then(serializedFile => {
+					// Publish drive_file_created event
+					event(user._id, 'drive_file_created', serializedFile);
 
-	// Register to search database
-	if (config.elasticsearch.enable) {
-		const es = require('../../db/elasticsearch');
-		es.index({
-			index: 'misskey',
-			type: 'drive_file',
-			id: file._id.toString(),
-			body: {
-				name: file.name,
-				user_id: user._id.toString()
-			}
-		});
-	}
+					// Register to search database
+					if (config.elasticsearch.enable) {
+						const es = require('../../db/elasticsearch');
+						es.index({
+							index: 'misskey',
+							type: 'drive_file',
+							id: file._id.toString(),
+							body: {
+								name: file.name,
+								user_id: user._id.toString()
+							}
+						});
+					}
+				});
+		})
+		.catch(reject);
 });
diff --git a/src/api/endpoints/drive/files/create.ts b/src/api/endpoints/drive/files/create.ts
index 7967c31187..7546eca309 100644
--- a/src/api/endpoints/drive/files/create.ts
+++ b/src/api/endpoints/drive/files/create.ts
@@ -1,7 +1,6 @@
 /**
  * Module dependencies
  */
-import * as fs from 'fs';
 import $ from 'cafy';
 import { validateFileName } from '../../../models/drive-file';
 import serialize from '../../../serializers/drive-file';
@@ -15,15 +14,11 @@ import create from '../../../common/add-file-to-drive';
  * @param {any} user
  * @return {Promise<any>}
  */
-module.exports = (file, params, user) => new Promise(async (res, rej) => {
+module.exports = async (file, params, user): Promise<any> => {
 	if (file == null) {
-		return rej('file is required');
+		throw 'file is required';
 	}
 
-	// TODO: 非同期にしたい。Promise対応してないんだろうか...
-	const buffer = fs.readFileSync(file.path);
-	fs.unlink(file.path, (err) => { if (err) console.log(err); });
-
 	// Get 'name' parameter
 	let name = file.originalname;
 	if (name !== undefined && name !== null) {
@@ -33,7 +28,7 @@ module.exports = (file, params, user) => new Promise(async (res, rej) => {
 		} else if (name === 'blob') {
 			name = null;
 		} else if (!validateFileName(name)) {
-			return rej('invalid name');
+			throw 'invalid name';
 		}
 	} else {
 		name = null;
@@ -41,14 +36,11 @@ module.exports = (file, params, user) => new Promise(async (res, rej) => {
 
 	// Get 'folder_id' parameter
 	const [folderId = null, folderIdErr] = $(params.folder_id).optional.nullable.id().$;
-	if (folderIdErr) return rej('invalid folder_id param');
+	if (folderIdErr) throw 'invalid folder_id param';
 
 	// Create file
-	const driveFile = await create(user, buffer, name, null, folderId);
+	const driveFile = await create(user, file.path, name, null, folderId);
 
 	// Serialize
-	const fileObj = await serialize(driveFile);
-
-	// Response
-	res(fileObj);
-});
+	return serialize(driveFile);
+};
diff --git a/src/api/endpoints/drive/files/upload_from_url.ts b/src/api/endpoints/drive/files/upload_from_url.ts
index 46cfffb69c..519e0bdf65 100644
--- a/src/api/endpoints/drive/files/upload_from_url.ts
+++ b/src/api/endpoints/drive/files/upload_from_url.ts
@@ -2,11 +2,16 @@
  * Module dependencies
  */
 import * as URL from 'url';
-const download = require('download');
 import $ from 'cafy';
 import { validateFileName } from '../../../models/drive-file';
 import serialize from '../../../serializers/drive-file';
 import create from '../../../common/add-file-to-drive';
+import * as debug from 'debug';
+import * as tmp from 'tmp';
+import * as fs from 'fs';
+import * as request from 'request';
+
+const log = debug('misskey:endpoint:upload_from_url');
 
 /**
  * Create a file from a URL
@@ -15,11 +20,11 @@ import create from '../../../common/add-file-to-drive';
  * @param {any} user
  * @return {Promise<any>}
  */
-module.exports = (params, user) => new Promise(async (res, rej) => {
+module.exports = async (params, user): Promise<any> => {
 	// Get 'url' parameter
 	// TODO: Validate this url
 	const [url, urlErr] = $(params.url).string().$;
-	if (urlErr) return rej('invalid url param');
+	if (urlErr) throw 'invalid url param';
 
 	let name = URL.parse(url).pathname.split('/').pop();
 	if (!validateFileName(name)) {
@@ -28,17 +33,35 @@ module.exports = (params, user) => new Promise(async (res, rej) => {
 
 	// Get 'folder_id' parameter
 	const [folderId = null, folderIdErr] = $(params.folder_id).optional.nullable.id().$;
-	if (folderIdErr) return rej('invalid folder_id param');
+	if (folderIdErr) throw 'invalid folder_id param';
+
+	// Create temp file
+	const path = await new Promise((res: (string) => void, rej) => {
+		tmp.file((e, path) => {
+			if (e) return rej(e);
+			res(path);
+		});
+	});
 
-	// Download file
-	const data = await download(url);
+	// write content at URL to temp file
+	await new Promise((res, rej) => {
+		const writable = fs.createWriteStream(path);
+		request(url)
+			.on('error', rej)
+			.on('end', () => {
+				writable.close();
+				res(path);
+			})
+			.pipe(writable)
+			.on('error', rej);
+	});
 
-	// Create file
-	const driveFile = await create(user, data, name, null, folderId);
+	const driveFile = await create(user, path, name, null, folderId);
 
-	// Serialize
-	const fileObj = await serialize(driveFile);
+	// clean-up
+	fs.unlink(path, (e) => {
+		if (e) log(e.stack);
+	});
 
-	// Response
-	res(fileObj);
-});
+	return serialize(driveFile);
+};