From 6cdce579fa41b6e8b5dd6824cda73d7c82ee71ed Mon Sep 17 00:00:00 2001
From: Akihiko Odaki <nekomanma@pixiv.co.jp>
Date: Tue, 3 Apr 2018 23:45:13 +0900
Subject: [PATCH] Enforce URI uniquness

---
 src/drive/add-file.ts            | 15 +++++--
 src/drive/upload-from-url.ts     |  4 +-
 src/models/drive-file.ts         | 16 ++++---
 src/models/post.ts               |  3 ++
 src/models/remote-user-object.ts | 15 -------
 src/remote/activitypub/create.ts | 77 +++++++++++++++++++++++---------
 6 files changed, 82 insertions(+), 48 deletions(-)
 delete mode 100644 src/models/remote-user-object.ts

diff --git a/src/drive/add-file.ts b/src/drive/add-file.ts
index f48fada7e..24eb5208d 100644
--- a/src/drive/add-file.ts
+++ b/src/drive/add-file.ts
@@ -10,7 +10,7 @@ import * as debug from 'debug';
 import fileType = require('file-type');
 import prominence = require('prominence');
 
-import DriveFile, { getGridFSBucket } from '../models/drive-file';
+import DriveFile, { IMetadata, getGridFSBucket } from '../models/drive-file';
 import DriveFolder from '../models/drive-folder';
 import { pack } from '../models/drive-file';
 import event, { publishDriveStream } from '../publishers/stream';
@@ -45,7 +45,8 @@ const addFile = async (
 	name: string = null,
 	comment: string = null,
 	folderId: mongodb.ObjectID = null,
-	force: boolean = false
+	force: boolean = false,
+	uri: string = null
 ) => {
 	log(`registering ${name} (user: ${getAcct(user)}, path: ${path})`);
 
@@ -224,12 +225,18 @@ const addFile = async (
 		properties['avgColor'] = averageColor;
 	}
 
-	return addToGridFS(detectedName, readable, mime, {
+	const metadata = {
 		userId: user._id,
 		folderId: folder !== null ? folder._id : null,
 		comment: comment,
 		properties: properties
-	});
+	} as IMetadata;
+
+	if (uri !== null) {
+		metadata.uri = uri;
+	}
+
+	return addToGridFS(detectedName, readable, mime, metadata);
 };
 
 /**
diff --git a/src/drive/upload-from-url.ts b/src/drive/upload-from-url.ts
index 7ff16e9e4..f96af0f26 100644
--- a/src/drive/upload-from-url.ts
+++ b/src/drive/upload-from-url.ts
@@ -8,7 +8,7 @@ import * as request from 'request';
 
 const log = debug('misskey:common:drive:upload_from_url');
 
-export default async (url, user, folderId = null): Promise<IDriveFile> => {
+export default async (url, user, folderId = null, uri = null): Promise<IDriveFile> => {
 	let name = URL.parse(url).pathname.split('/').pop();
 	if (!validateFileName(name)) {
 		name = null;
@@ -35,7 +35,7 @@ export default async (url, user, folderId = null): Promise<IDriveFile> => {
 			.on('error', rej);
 	});
 
-	const driveFile = await create(user, path, name, null, folderId);
+	const driveFile = await create(user, path, name, null, folderId, false, uri);
 
 	// clean-up
 	fs.unlink(path, (e) => {
diff --git a/src/models/drive-file.ts b/src/models/drive-file.ts
index fba1aebda..c86570f0f 100644
--- a/src/models/drive-file.ts
+++ b/src/models/drive-file.ts
@@ -6,6 +6,8 @@ import monkDb, { nativeDbConn } from '../db/mongodb';
 
 const DriveFile = monkDb.get<IDriveFile>('driveFiles.files');
 
+DriveFile.createIndex('metadata.uri', { sparse: true, unique: true });
+
 export default DriveFile;
 
 const getGridFSBucket = async (): Promise<mongodb.GridFSBucket> => {
@@ -18,17 +20,21 @@ const getGridFSBucket = async (): Promise<mongodb.GridFSBucket> => {
 
 export { getGridFSBucket };
 
+export type IMetadata = {
+	properties: any;
+	userId: mongodb.ObjectID;
+	folderId: mongodb.ObjectID;
+	comment: string;
+	uri: string;
+};
+
 export type IDriveFile = {
 	_id: mongodb.ObjectID;
 	uploadDate: Date;
 	md5: string;
 	filename: string;
 	contentType: string;
-	metadata: {
-		properties: any;
-		userId: mongodb.ObjectID;
-		folderId: mongodb.ObjectID;
-	}
+	metadata: IMetadata;
 };
 
 export function validateFileName(name: string): boolean {
diff --git a/src/models/post.ts b/src/models/post.ts
index 798c18e4b..2f2b51b94 100644
--- a/src/models/post.ts
+++ b/src/models/post.ts
@@ -11,6 +11,8 @@ import { pack as packFile } from './drive-file';
 
 const Post = db.get<IPost>('posts');
 
+Post.createIndex('uri', { sparse: true, unique: true });
+
 export default Post;
 
 export function isValidText(text: string): boolean {
@@ -49,6 +51,7 @@ export type IPost = {
 		heading: number;
 		speed: number;
 	};
+	uri: string;
 };
 
 /**
diff --git a/src/models/remote-user-object.ts b/src/models/remote-user-object.ts
deleted file mode 100644
index fb5b337c9..000000000
--- a/src/models/remote-user-object.ts
+++ /dev/null
@@ -1,15 +0,0 @@
-import * as mongodb from 'mongodb';
-import db from '../db/mongodb';
-
-const RemoteUserObject = db.get<IRemoteUserObject>('remoteUserObjects');
-
-export default RemoteUserObject;
-
-export type IRemoteUserObject = {
-	_id: mongodb.ObjectID;
-	uri: string;
-	object: {
-		$ref: string;
-		$id: mongodb.ObjectID;
-	}
-};
diff --git a/src/remote/activitypub/create.ts b/src/remote/activitypub/create.ts
index 3676ba746..46e4c9988 100644
--- a/src/remote/activitypub/create.ts
+++ b/src/remote/activitypub/create.ts
@@ -1,6 +1,8 @@
 import { JSDOM } from 'jsdom';
+import { ObjectID } from 'mongodb';
 import config from '../../config';
-import RemoteUserObject, { IRemoteUserObject } from '../../models/remote-user-object';
+import DriveFile from '../../models/drive-file';
+import Post from '../../models/post';
 import { IRemoteUser } from '../../models/user';
 import uploadFromUrl from '../../drive/upload-from-url';
 import createPost from '../../post/create';
@@ -8,15 +10,13 @@ import distributePost from '../../post/distribute';
 import Resolver from './resolver';
 const createDOMPurify = require('dompurify');
 
-function createRemoteUserObject($ref, $id, { id }) {
-	const object = { $ref, $id };
-
-	if (!id) {
-		return { object };
-	}
-
-	return RemoteUserObject.insert({ uri: id, object });
-}
+type IResult = {
+	resolver: Resolver;
+	object: {
+		$ref: string;
+		$id: ObjectID;
+	};
+};
 
 class Creator {
 	private actor: IRemoteUser;
@@ -27,17 +27,23 @@ class Creator {
 		this.distribute = distribute;
 	}
 
-	private async createImage(image) {
+	private async createImage(resolver: Resolver, image) {
 		if ('attributedTo' in image && this.actor.account.uri !== image.attributedTo) {
 			throw new Error();
 		}
 
-		const { _id } = await uploadFromUrl(image.url, this.actor);
-		return createRemoteUserObject('driveFiles.files', _id, image);
+		const { _id } = await uploadFromUrl(image.url, this.actor, image.id || null);
+		return {
+			resolver,
+			object: { $ref: 'driveFiles.files', $id: _id }
+		};
 	}
 
 	private async createNote(resolver: Resolver, note) {
-		if ('attributedTo' in note && this.actor.account.uri !== note.attributedTo) {
+		if (
+			('attributedTo' in note && this.actor.account.uri !== note.attributedTo) ||
+			typeof note.id !== 'string'
+		) {
 			throw new Error();
 		}
 
@@ -61,10 +67,10 @@ class Creator {
 			userId: this.actor._id,
 			appId: null,
 			viaMobile: false,
-			geo: undefined
+			geo: undefined,
+			uri: note.id
 		}, null, null, []);
 
-		const promisedRemoteUserObject = createRemoteUserObject('posts', inserted._id, note);
 		const promises = [];
 
 		if (this.distribute) {
@@ -89,18 +95,45 @@ class Creator {
 
 		await Promise.all(promises);
 
-		return promisedRemoteUserObject;
+		return {
+			resolver,
+			object: { $ref: 'posts', id: inserted._id }
+		};
 	}
 
-	public async create(parentResolver: Resolver, value): Promise<Array<Promise<IRemoteUserObject>>> {
+	public async create(parentResolver: Resolver, value): Promise<Array<Promise<IResult>>> {
 		const collection = await parentResolver.resolveCollection(value);
 
 		return collection.object.map(async element => {
 			if (typeof element === 'string') {
-				const object = RemoteUserObject.findOne({ uri: element });
+				try {
+					await Promise.all([
+						DriveFile.findOne({ 'metadata.uri': element }).then(file => {
+							if (file === null) {
+								return;
+							}
 
-				if (object !== null) {
-					return object;
+							throw {
+								$ref: 'driveFile.files',
+								$id: file._id
+							};
+						}, () => {}),
+						Post.findOne({ uri: element }).then(post => {
+							if (post === null) {
+								return;
+							}
+
+							throw {
+								$ref: 'posts',
+								$id: post._id
+							};
+						}, () => {})
+					]);
+				} catch (object) {
+					return {
+						resolver: collection.resolver,
+						object
+					};
 				}
 			}
 
@@ -108,7 +141,7 @@ class Creator {
 
 			switch (object.type) {
 			case 'Image':
-				return this.createImage(object);
+				return this.createImage(resolver, object);
 
 			case 'Note':
 				return this.createNote(resolver, object);