From 7e81e0db6ac1289ae9504f7e3da5db6e56f41a51 Mon Sep 17 00:00:00 2001
From: otofune <otofune@gmail.com>
Date: Mon, 6 Nov 2017 14:37:00 +0900
Subject: [PATCH] support GridFS

---
 src/api/common/add-file-to-drive.ts | 37 ++++++++++++++++++-----------
 src/api/models/drive-file.ts        | 15 ++++++++++--
 src/db/mongodb.ts                   | 35 +++++++++++++++++++++++----
 3 files changed, 67 insertions(+), 20 deletions(-)

diff --git a/src/api/common/add-file-to-drive.ts b/src/api/common/add-file-to-drive.ts
index 714eeb520..f48f0cbcf 100644
--- a/src/api/common/add-file-to-drive.ts
+++ b/src/api/common/add-file-to-drive.ts
@@ -4,14 +4,27 @@ import * as gm from 'gm';
 import * as debug from 'debug';
 import fileType = require('file-type');
 import prominence = require('prominence');
-import DriveFile from '../models/drive-file';
+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
  *
@@ -58,7 +71,7 @@ export default (
 
 	// Generate hash
 	const hash = crypto
-		.createHash('sha256')
+		.createHash('md5')
 		.update(data)
 		.digest('hex') as string;
 
@@ -67,8 +80,10 @@ export default (
 	if (!force) {
 		// Check if there is a file with the same hash
 		const much = await DriveFile.findOne({
-			user_id: user._id,
-			hash: hash
+			md5: hash,
+			metadata: {
+				user_id: user._id
+			}
 		});
 
 		if (much !== null) {
@@ -82,13 +97,13 @@ export default (
 	// Calculate drive usage
 	const usage = ((await DriveFile
 		.aggregate([
-			{ $match: { user_id: user._id } },
+			{ $match: { metadata: { user_id: user._id } } },
 			{ $project: {
-				datasize: true
+				length: true
 			}},
 			{ $group: {
 				_id: null,
-				usage: { $sum: '$datasize' }
+				usage: { $sum: '$length' }
 			}}
 		]))[0] || {
 			usage: 0
@@ -131,21 +146,15 @@ export default (
 	}
 
 	// Create DriveFile document
-	const file = await DriveFile.insert({
-		created_at: new Date(),
+	const file = await addToGridFS(`${user._id}/${name}`, data, {
 		user_id: user._id,
 		folder_id: folder !== null ? folder._id : null,
-		data: data,
-		datasize: size,
 		type: mime,
 		name: name,
 		comment: comment,
-		hash: hash,
 		properties: properties
 	});
 
-	delete file.data;
-
 	log(`drive file has been created ${file._id}`);
 
 	resolve(file);
diff --git a/src/api/models/drive-file.ts b/src/api/models/drive-file.ts
index 8d158cf56..79a87f657 100644
--- a/src/api/models/drive-file.ts
+++ b/src/api/models/drive-file.ts
@@ -1,11 +1,22 @@
-import db from '../../db/mongodb';
+import * as mongodb from 'mongodb';
+import monkDb, { nativeDbConn } from '../../db/mongodb';
 
-const collection = db.get('drive_files');
+const collection = monkDb.get('drive_files.files');
 
 (collection as any).createIndex('hash'); // fuck type definition
 
 export default collection as any; // fuck type definition
 
+const getGridFSBucket = async (): Promise<mongodb.GridFSBucket> => {
+	const db = await nativeDbConn()
+	const bucket = new mongodb.GridFSBucket(db, {
+		bucketName: 'drive_files'
+	})
+	return bucket
+}
+
+export { getGridFSBucket }
+
 export function validateFileName(name: string): boolean {
 	return (
 		(name.trim().length > 0) &&
diff --git a/src/db/mongodb.ts b/src/db/mongodb.ts
index 6ee7f4534..75f1a1d3c 100644
--- a/src/db/mongodb.ts
+++ b/src/db/mongodb.ts
@@ -1,11 +1,38 @@
-import * as mongo from 'monk';
-
 import config from '../conf';
 
 const uri = config.mongodb.user && config.mongodb.pass
-	? `mongodb://${config.mongodb.user}:${config.mongodb.pass}@${config.mongodb.host}:${config.mongodb.port}/${config.mongodb.db}`
-	: `mongodb://${config.mongodb.host}:${config.mongodb.port}/${config.mongodb.db}`;
+? `mongodb://${config.mongodb.user}:${config.mongodb.pass}@${config.mongodb.host}:${config.mongodb.port}/${config.mongodb.db}`
+: `mongodb://${config.mongodb.host}:${config.mongodb.port}/${config.mongodb.db}`;
+
+/**
+ * monk
+ */
+import * as mongo from 'monk';
 
 const db = mongo(uri);
 
 export default db;
+
+/**
+ * MongoDB native module (officialy)
+ */
+import * as mongodb from 'mongodb'
+
+let mdb: mongodb.Db;
+
+const nativeDbConn = async (): Promise<mongodb.Db> => {
+	if (mdb) return mdb;
+
+	const db = await ((): Promise<mongodb.Db> => new Promise((resolve, reject) => {
+		mongodb.MongoClient.connect(uri, (e, db) => {
+			if (e) return reject(e)
+			resolve(db)
+		})
+	}))()
+
+	mdb = db
+
+	return db
+}
+
+export { nativeDbConn }