Sharkey/src/api/common/add-file-to-drive.ts

239 lines
6.1 KiB
TypeScript
Raw Normal View History

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');
2017-11-06 07:37:00 +02:00
import DriveFile, { getGridFSBucket } from '../models/drive-file';
2016-12-29 00:49:51 +02:00
import DriveFolder from '../models/drive-folder';
import serialize from '../serializers/drive-file';
import event from '../event';
2017-01-17 02:17:52 +02:00
import config from '../../conf';
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
2017-02-06 14:11:09 +02:00
const log = debug('misskey:register-drive-file');
const tmpFile = (): Promise<string> => new Promise((resolve, reject) => {
tmp.file((e, path) => {
2017-11-13 20:47:42 +02:00
if (e) return reject(e);
resolve(path);
});
});
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);
2017-11-13 20:47:42 +02:00
}));
2017-11-06 07:37:00 +02:00
2016-12-29 00:49:51 +02:00
/**
* Add file to drive
*
* @param user User who wish to add file
* @param file File path or readableStream
2016-12-29 00:49:51 +02:00
* @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,
2016-12-29 00:49:51 +02:00
name: string = null,
comment: string = null,
folderId: mongodb.ObjectID = null,
force: boolean = false
) => new Promise<any>((resolve, reject) => {
2017-02-06 14:11:09 +02:00
log(`registering ${name} (user: ${user.username})`);
// Get file path
new Promise((res: (v: string) => void, rej) => {
if (typeof file === 'string') {
2017-11-13 20:47:42 +02:00
res(file);
return;
2016-12-29 00:49:51 +02:00
}
if (typeof file === 'object' && typeof file.read === 'function') {
tmpFile()
.then(path => {
2017-11-13 20:47:42 +02:00
const readable: stream.Readable = file;
const writable = fs.createWriteStream(path);
readable
.on('error', rej)
.on('end', () => {
2017-11-13 20:47:42 +02:00
res(path);
})
.pipe(writable)
2017-11-13 20:47:42 +02:00
.on('error', rej);
})
2017-11-13 20:47:42 +02:00
.catch(rej);
2016-12-29 00:49:51 +02:00
}
2017-11-13 20:47:42 +02:00
rej(new Error('un-compatible file.'));
})
// Calculate hash, get content type and get file size
.then(path => Promise.all([
path,
// hash
((): Promise<string> => new Promise((res, rej) => {
2017-11-13 20:47:42 +02:00
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) => {
2017-11-13 20:47:42 +02:00
const readable = fs.createReadStream(path);
readable
.on('error', rej)
.once('data', (buffer: Buffer) => {
2017-11-13 20:47:42 +02:00
readable.destroy();
const type = fileType(buffer);
if (!type) {
2017-11-13 20:47:42 +02:00
return res(['application/octet-stream', null]);
}
2017-11-13 20:47:42 +02:00
return res([type.mime, type.ext]);
});
}))(),
// size
((): Promise<number> => new Promise((res, rej) => {
fs.stat(path, (err, stats) => {
2017-11-13 20:47:42 +02:00
if (err) return rej(err);
res(stats.size);
});
}))()
]))
.then(async ([path, hash, [mime, ext], size]) => {
2017-11-13 20:47:42 +02:00
log(`hash: ${hash}, mime: ${mime}, ext: ${ext}, size: ${size}`);
// 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,
'metadata.user_id': user._id
});
if (much !== null) {
log('file with same hash is found');
return resolve(much);
} else {
log('file with same hash is not found');
}
}
2016-12-29 00:49:51 +02:00
const [properties, folder] = await Promise.all([
// properties
(async () => {
if (!/^image\/.*$/.test(mime)) {
2017-11-13 20:47:42 +02:00
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');
2017-11-13 20:47:42 +02:00
return properties;
})(),
// folder
(async () => {
if (!folderId) {
2017-11-13 20:47:42 +02:00
return null;
}
const driveFolder = await DriveFolder.findOne({
_id: folderId,
user_id: user._id
2017-11-13 20:47:42 +02:00
});
if (!driveFolder) {
2017-11-13 20:47:42 +02:00
throw 'folder-not-found';
}
2017-11-13 20:47:42 +02:00
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) {
2017-11-13 20:47:42 +02:00
return aggregates[0].usage;
}
2017-11-13 20:47:42 +02:00
return 0;
});
log(`drive usage is ${usage}`);
// If usage limit exceeded
if (usage + size > user.drive_capacity) {
throw 'no-free-space';
}
})()
2017-11-13 20:47:42 +02:00
]);
2017-11-13 20:47:42 +02:00
const readable = fs.createReadStream(path);
2017-11-13 21:28:51 +02:00
return addToGridFS(detectedName, readable, mime, {
user_id: user._id,
folder_id: folder !== null ? folder._id : null,
comment: comment,
properties: properties
2017-11-13 20:47:42 +02:00
});
})
.then(file => {
log(`drive file has been created ${file._id}`);
2017-11-13 20:47:42 +02:00
resolve(file);
2017-11-13 21:28:51 +02:00
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()
}
});
}
});
})
2017-11-13 20:47:42 +02:00
.catch(reject);
2016-12-29 00:49:51 +02:00
});