From bd434ed02dd69b4a120de5a727a3e5d321f3cbfe Mon Sep 17 00:00:00 2001 From: syuilo Date: Fri, 25 May 2018 18:41:49 +0900 Subject: [PATCH] Refactor and some fixes --- .../api/endpoints/drive/files/create.ts | 9 + src/services/drive/add-file.ts | 427 ++++++++---------- 2 files changed, 188 insertions(+), 248 deletions(-) diff --git a/src/server/api/endpoints/drive/files/create.ts b/src/server/api/endpoints/drive/files/create.ts index e9348e4e2..dd748d6bb 100644 --- a/src/server/api/endpoints/drive/files/create.ts +++ b/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 => { 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; } }; diff --git a/src/services/drive/add-file.ts b/src/services/drive/add-file.ts index bcd5bee51..8b1d3eef0 100644 --- a/src/services/drive/add-file.ts +++ b/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 => { - 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 => 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 => new Promise((res, rej) => { - fs.stat(path, (err, stats) => { - if (err) return rej(err); - res(stats.size); +): Promise { + // Calc md5 hash + const calcHash = 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')); }); - }))() - ]); + }); + + // 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((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); + const driveFile = await (writeChunks(detectedName, readable, mime, metadata) as Promise); + + 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((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((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; +}