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

179 lines
3.9 KiB
TypeScript

import * as mongodb from 'mongodb';
import * as crypto from 'crypto';
import * as gm from 'gm';
import * as debug from 'debug';
import fileType = require('file-type');
import prominence = require('prominence');
import DriveFile, { getGridFSBucket } from '../models/drive-file';
import DriveFolder from '../models/drive-folder';
import serialize from '../serializers/drive-file';
import event from '../event';
import config from '../../conf';
import { Duplex } from 'stream';
const log = debug('misskey:register-drive-file');
const addToGridFS = (name, binary, 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, { metadata });
writeStream.once('finish', (doc) => { resolve(doc); });
writeStream.on('error', reject);
dataStream.pipe(writeStream);
});
/**
* 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 (
user: any,
data: Buffer,
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}`);
// Generate hash
const hash = crypto
.createHash('md5')
.update(data)
.digest('hex') as string;
log(`hash is ${hash}`);
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');
}
}
// 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({
_id: folderId,
user_id: user._id
});
if (folder === null) {
return reject('folder-not-found');
}
}
let properties: any = null;
// 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
};
log('image width and height is calculated');
}
// Create DriveFile document
const file = await addToGridFS(name, data, {
user_id: user._id,
folder_id: folder !== null ? folder._id : null,
type: mime,
comment: comment,
properties: properties
});
log(`drive file has been created ${file._id}`);
resolve(file);
// Serialize
const fileObj = await serialize(file);
// Publish drive_file_created event
event(user._id, 'drive_file_created', fileObj);
// 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()
}
});
}
});