Sharkey/src/services/drive/add-file.ts

374 lines
9.3 KiB
TypeScript
Raw Normal View History

import { Buffer } from 'buffer';
import * as fs from 'fs';
import * as tmp from 'tmp';
import * as stream from 'stream';
2016-12-29 00:49:51 +02:00
import * as mongodb from 'mongodb';
import * as crypto from 'crypto';
import * as _gm from 'gm';
2017-02-06 14:11:09 +02:00
import * as debug from 'debug';
2017-01-02 23:03:19 +02:00
import fileType = require('file-type');
import prominence = require('prominence');
import DriveFile, { IMetadata, getDriveFileBucket, IDriveFile, DriveFileChunk } from '../../models/drive-file';
2018-04-04 21:21:11 +03:00
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 DriveFileThumbnail, { getDriveFileThumbnailBucket, DriveFileThumbnailChunk } from '../../models/drive-file-thumbnail';
import genThumbnail from '../../drive/gen-thumbnail';
2016-12-29 00:49:51 +02:00
const gm = _gm.subClass({
imageMagick: true
});
2018-03-27 10:51:12 +03:00
const log = debug('misskey:drive:add-file');
2017-02-06 14:11:09 +02:00
2018-04-09 22:02:25 +03:00
const tmpFile = (): Promise<[string, any]> => new Promise((resolve, reject) => {
tmp.file((e, path, fd, cleanup) => {
2017-11-13 20:47:42 +02:00
if (e) return reject(e);
2018-04-09 22:02:25 +03:00
resolve([path, cleanup]);
2017-11-13 20:47:42 +02:00
});
});
const writeChunks = (name: string, readable: stream.Readable, type: string, metadata: any) =>
getDriveFileBucket()
.then(bucket => new Promise((resolve, reject) => {
const writeStream = bucket.openUploadStream(name, { contentType: type, metadata });
2018-04-09 20:34:47 +03:00
writeStream.once('finish', resolve);
writeStream.on('error', reject);
readable.pipe(writeStream);
2017-11-13 20:47:42 +02:00
}));
2017-11-06 07:37:00 +02:00
const writeThumbnailChunks = (name: string, readable: stream.Readable, originalId) =>
getDriveFileThumbnailBucket()
.then(bucket => new Promise((resolve, reject) => {
const writeStream = bucket.openUploadStream(name, {
contentType: 'image/jpeg',
metadata: {
originalId
}
});
writeStream.once('finish', resolve);
writeStream.on('error', reject);
readable.pipe(writeStream);
}));
const addFile = async (
2018-04-09 20:34:47 +03:00
user: IUser,
path: string,
2016-12-29 00:49:51 +02:00
name: string = null,
comment: string = null,
folderId: mongodb.ObjectID = null,
2018-04-03 17:45:13 +03:00
force: boolean = false,
uri: string = null
2018-04-09 20:34:47 +03:00
): Promise<IDriveFile> => {
2018-03-27 10:51:12 +03:00
log(`registering ${name} (user: ${getAcct(user)}, path: ${path})`);
2017-02-06 14:11:09 +02:00
// 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'));
2017-11-13 20:47:42 +02:00
});
}))(),
// 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);
});
}))()
]);
log(`hash: ${hash}, mime: ${mime}, ext: ${ext}, size: ${size}`);
2016-12-29 00:49:51 +02:00
// detect name
const detectedName: string = name || (ext ? `untitled.${ext}` : 'untitled');
if (!force) {
// Check if there is a file with the same hash
const much = await DriveFile.findOne({
md5: hash,
2018-03-29 08:48:47 +03:00
'metadata.userId': user._id
});
if (much !== null) {
log('file with same hash is found');
return much;
} else {
log('file with same hash is not found');
}
}
2017-12-11 06:33:33 +02:00
const [wh, averageColor, folder] = await Promise.all([
2017-12-08 12:42:02 +02:00
// Width and height (when image)
(async () => {
// 画像かどうか
if (!/^image\/.*$/.test(mime)) {
return null;
}
const imageType = mime.split('/')[1];
2017-12-08 12:42:02 +02:00
// 画像でもPNGかJPEGかGIFでないならスキップ
if (imageType != 'png' && imageType != 'jpeg' && imageType != 'gif') {
return null;
}
2017-12-11 06:33:33 +02:00
log('calculate image width and height...');
2017-12-08 12:42:02 +02:00
// Calculate width and height
const g = gm(fs.createReadStream(path), name);
const size = await prominence(g).size();
2017-12-11 06:33:33 +02:00
log(`image width and height is calculated: ${size.width}, ${size.height}`);
2017-12-08 12:42:02 +02:00
return [size.width, size.height];
})(),
2017-12-11 06:33:33 +02:00
// 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;
}
log('calculate average color...');
const buffer = await prominence(gm(fs.createReadStream(path), name)
.setFormat('ppm')
.resize(1, 1)) // 1pxのサイズに縮小して平均色を取得するというハック
.toBuffer();
const r = buffer.readUInt8(buffer.length - 3);
const g = buffer.readUInt8(buffer.length - 2);
const b = buffer.readUInt8(buffer.length - 1);
log(`average color is calculated: ${r}, ${g}, ${b}`);
return [r, g, b];
})(),
// folder
(async () => {
if (!folderId) {
return null;
}
const driveFolder = await DriveFolder.findOne({
_id: folderId,
2018-03-29 08:48:47 +03:00
userId: user._id
});
if (!driveFolder) {
throw 'folder-not-found';
}
return driveFolder;
})(),
// usage checker
(async () => {
// Calculate drive usage
const usage = await DriveFile
.aggregate([{
2018-05-04 10:37:15 +03:00
$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;
});
log(`drive usage is ${usage}`);
// If usage limit exceeded
2018-03-29 08:48:47 +03:00
if (usage + size > user.driveCapacity) {
2018-04-17 13:55:58 +03:00
if (isLocalUser(user)) {
throw 'no-free-space';
} else {
//#region (アバターまたはバナーを含まず)最も古いファイルを削除する
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: {
2018-05-04 10:37:15 +03:00
'metadata.deletedAt': new Date(),
'metadata.isExpired': true
2018-04-17 13:55:58 +03:00
}
});
//#region サムネイルもあれば削除
const thumbnail = await DriveFileThumbnail.findOne({
'metadata.originalId': oldFile._id
});
if (thumbnail) {
DriveFileThumbnailChunk.remove({
files_id: thumbnail._id
});
DriveFileThumbnail.remove({ _id: thumbnail._id });
}
//#endregion
2018-04-17 13:55:58 +03:00
}
//#endregion
}
}
})()
]);
const readable = fs.createReadStream(path);
2017-12-08 12:42:02 +02:00
const properties = {};
if (wh) {
properties['width'] = wh[0];
properties['height'] = wh[1];
}
2017-12-11 06:33:33 +02:00
if (averageColor) {
2018-03-29 08:48:47 +03:00
properties['avgColor'] = averageColor;
2017-12-11 06:33:33 +02:00
}
2018-04-03 17:45:13 +03:00
const metadata = {
2018-03-29 08:48:47 +03:00
userId: user._id,
_user: {
host: user.host
},
2018-03-29 08:48:47 +03:00
folderId: folder !== null ? folder._id : null,
comment: comment,
properties: properties
2018-04-03 17:45:13 +03:00
} as IMetadata;
if (uri !== null) {
metadata.uri = uri;
}
const file = await (writeChunks(detectedName, readable, mime, metadata) as Promise<IDriveFile>);
try {
const thumb = await genThumbnail(file);
if (thumb) {
await writeThumbnailChunks(detectedName, thumb, file._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) => {
2018-04-09 20:34:47 +03:00
const isStream = typeof file === 'object' && typeof file.read === 'function';
// Get file path
2018-04-09 22:02:25 +03:00
new Promise<[string, any]>((res, rej) => {
if (typeof file === 'string') {
2018-04-09 22:02:25 +03:00
res([file, null]);
2018-04-09 20:34:47 +03:00
} else if (isStream) {
tmpFile()
2018-04-09 22:02:25 +03:00
.then(([path, cleanup]) => {
const readable: stream.Readable = file;
const writable = fs.createWriteStream(path);
readable
.on('error', rej)
.on('end', () => {
2018-04-09 22:02:25 +03:00
res([path, cleanup]);
})
.pipe(writable)
.on('error', rej);
})
.catch(rej);
2018-04-09 20:34:47 +03:00
} else {
rej(new Error('un-compatible file.'));
}
})
2018-04-09 22:02:25 +03:00
.then(([path, cleanup]) => new Promise<IDriveFile>((res, rej) => {
addFile(user, path, ...args)
.then(file => {
res(file);
2018-04-09 22:02:25 +03:00
if (cleanup) cleanup();
})
.catch(rej);
}))
.then(file => {
log(`drive file has been created ${file._id}`);
2018-04-09 20:34:47 +03:00
resolve(file);
2018-04-09 20:34:47 +03:00
pack(file).then(packedFile => {
// Publish drive_file_created event
2018-04-09 20:34:47 +03:00
event(user._id, 'drive_file_created', packedFile);
publishDriveStream(user._id, 'file_created', packedFile);
});
})
.catch(reject);
2016-12-29 00:49:51 +02:00
});