Skip to content
Snippets Groups Projects
Commit bd434ed0 authored by syuilo's avatar syuilo
Browse files

Refactor and some fixes

parent df89f5c8
No related branches found
Tags 10.87.4
No related merge requests found
/**
* 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<any> => {
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;
}
};
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<IDriveFile> => {
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<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([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);
): Promise<IDriveFile> {
// Calc md5 hash
const calcHash = new Promise<string>((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<number>((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 }
}
}, {
$project: {
length: true
}
}, {
$group: {
_id: null,
usage: { $sum: '$length' }
}
}])
.then((aggregates: any[]) => {
if (aggregates.length > 0) {
return aggregates[0].usage;
}
return 0;
});
const imageType = mime.split('/')[1];
log(`drive usage is ${usage}`);
// 画像でもPNGかJPEGかGIFでないならスキップ
if (imageType != 'png' && imageType != 'jpeg' && imageType != 'gif') {
return null;
}
// 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,111 +260,17 @@ 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;
});
log(`drive usage is ${usage}`);
// 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
}
});
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 readable = fs.createReadStream(path);
const value = isTransparent ? [r, g, b, 255] : [r, g, b];
const properties = {};
properties['avgColor'] = value;
};
if (wh) {
properties['width'] = wh[0];
properties['height'] = wh[1];
propPromises = [calcWh(), calcAvg()];
}
if (averageColor) {
properties['avgColor'] = averageColor;
}
const [folder] = await Promise.all([fetchFolder(), propPromises]);
const readable = fs.createReadStream(path);
const metadata = {
userId: user._id,
......@@ -309,74 +290,24 @@ const addFile = async (
metadata.uri = uri;
}
const file = await (writeChunks(detectedName, readable, mime, metadata) as Promise<IDriveFile>);
const driveFile = await (writeChunks(detectedName, readable, mime, metadata) as Promise<IDriveFile>);
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<any>((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<IDriveFile>((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;
}
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment