From b081dd33e99cee630e00e0cdda04841d6da5ed05 Mon Sep 17 00:00:00 2001
From: atsu1125 <atsuchan@atsuchan.page>
Date: Sun, 4 Feb 2024 19:22:37 +0900
Subject: [PATCH] server: fix drive quota for remote users

---
 src/services/drive/add-file.ts | 42 ++++++++++++++++++++++++++++++++--
 1 file changed, 40 insertions(+), 2 deletions(-)

diff --git a/src/services/drive/add-file.ts b/src/services/drive/add-file.ts
index 90c189fac..20fbf1b89 100644
--- a/src/services/drive/add-file.ts
+++ b/src/services/drive/add-file.ts
@@ -1,6 +1,7 @@
 import * as fs from 'fs';
 
 import { v4 as uuid } from 'uuid';
+import { In, IsNull, getConnection } from 'typeorm';
 
 import { publishMainStream, publishDriveStream } from '../stream';
 import { deleteFile } from './delete-file';
@@ -289,6 +290,40 @@ async function deleteOldFile(user: IRemoteUser) {
 	}
 }
 
+async function expireOldFiles(user: IRemoteUser, driveCapacity: number): Promise<void> {
+	// Delete as many files as necessary so the total usage is below driveCapacity,
+	// oldest files first, and exclude avatar and banner.
+	//
+	// Using a window function, i.e. `OVER (ORDER BY "createdAt" DESC)` means that
+	// the `SUM` will be a running total.
+	const conn = await getConnection();
+	const exceededFileIds = await conn.query('SELECT "id" FROM ('
+		+ 'SELECT "id", SUM("size") OVER (ORDER BY "createdAt" DESC) AS "total" FROM "drive_file" WHERE "userId" = $1 AND NOT "isLink"'
+		+ (user.avatarId ? ' AND "id" != $2' : '')
+		+ (user.bannerId ? ' AND "id" != $3' : '')
+		+ ') AS "totals" WHERE "total" > $4',
+		[
+			user.id,
+			user.avatarId ?? '',
+			user.bannerId ?? '',
+			driveCapacity,
+		]
+	);
+
+	if (exceededFileIds.length === 0) {
+		// no files to expire, avatar and banner if present are already the only files
+		throw new Error('remote user drive quota met by avatar and banner');
+	}
+
+	const files = await DriveFiles.find({
+		id: In(exceededFileIds.map(x => x.id)),
+	});
+
+	for (const file of files) {
+		deleteFile(file, true);
+	}
+}
+
 /**
  * Add file to drive
  *
@@ -350,6 +385,7 @@ export default async function(
 	//#region Check drive usage
 	if (user && !isLink) {
 		const usage = await DriveFiles.calcDriveUsageOf(user);
+		const isLocalUser = Users.isLocalUser(user);
 
 		const instance = await fetchMeta();
 		const driveCapacity = 1024 * 1024 * getDriveCapacity(user, instance);
@@ -358,11 +394,13 @@ export default async function(
 
 		// If usage limit exceeded
 		if (usage + info.size > driveCapacity) {
-			if (Users.isLocalUser(user)) {
+			if (isLocalUser) {
 				// throw new IdentifiableError('c6244ed2-a39a-4e1c-bf93-f0fbd7764fa6', 'No free space.');
 			} else {
 				// (アバターまたはバナーを含まず)最も古いファイルを削除する
-				deleteOldFile(user as IRemoteUser);
+				//deleteOldFile(user as IRemoteUser);
+				// delete older files to make space for new file
+				expireOldFiles(await Users.findOneOrFail({ id: user.id }) as IRemoteUser, driveCapacity - info.size);
 			}
 		}
 	}