johan/FoundKey
1
0
Fork 0
forked from FoundKeyGang/FoundKey
Code Issues Pull requests Projects Releases Packages Wiki Activity

server: fix drive quota for remote users

Browse source 
This deletes as many files as necessary to ensure the drive quota for
remote users is kept. Previously only one file would have been deleted
for each file added.

Changelog: Fixed
Co-authored-by: CGsama <CGsama@outlook.com>
Co-authored-by: tamaina <tamaina@hotmail.co.jp>
This commit is contained in:
Johann150 2023-01-09 21:27:05 +01:00
parent 4a77e93dfd
commit 2fde652b4a
Signed by untrusted user: Johann150
GPG key ID: 9EE6577A2A06F8F1
1 changed files with 33 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

53
packages/backend/src/services/drive/add-file.ts
View file

@ -2,9 +2,10 @@ import * as fs from 'node:fs';
import { v4 as uuid } from 'uuid'; import { v4 as uuid } from 'uuid';
import S3 from 'aws-sdk/clients/s3.js'; import S3 from 'aws-sdk/clients/s3.js';
import { IsNull } from 'typeorm'; import { In, IsNull } from 'typeorm';
import sharp from 'sharp'; import sharp from 'sharp';
import { db } from '@/db/postgre.js';
import { FILE_TYPE_BROWSERSAFE } from '@/const.js'; import { FILE_TYPE_BROWSERSAFE } from '@/const.js';
import { publishMainStream, publishDriveStream } from '@/services/stream.js'; import { publishMainStream, publishDriveStream } from '@/services/stream.js';
import { fetchMeta } from '@/misc/fetch-meta.js'; import { fetchMeta } from '@/misc/fetch-meta.js';
@ -290,25 +291,36 @@ async function upload(key: string, stream: fs.ReadStream | Buffer, _type: string
if (result) logger.debug(`Uploaded: ${result.Bucket}/${result.Key} => ${result.Location}`); if (result) logger.debug(`Uploaded: ${result.Bucket}/${result.Key} => ${result.Location}`);
} }
async function deleteOldFile(user: IRemoteUser): Promise<void> { async function expireOldFiles(user: IRemoteUser, driveCapacity: number): Promise<void> {
const q = DriveFiles.createQueryBuilder('file') // Delete as many files as necessary so the total usage is below driveCapacity,
.where('file.userId = :userId', { userId: user.id }) // oldest files first, and exclude avatar and banner.
.andWhere('NOT file.isLink'); //
// Using a window function, i.e. `OVER (ORDER BY "createdAt" DESC)` means that
// the `SUM` will be a running total.
const exceededFileIds = await db.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 (user.avatarId) { if (exceededFileIds.length === 0) {
q.andWhere('file.id != :avatarId', { avatarId: user.avatarId }); // 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');
} }
if (user.bannerId) { const files = await DriveFiles.findBy({
q.andWhere('file.id != :bannerId', { bannerId: user.bannerId }); id: In(exceededFileIds.map(x => x.id)),
} });
q.orderBy('file.id', 'ASC'); for (const file of files) {
deleteFile(file, true);
const oldFile = await q.getOne();
if (oldFile) {
deleteFile(oldFile, true);
} }
} }
@ -373,19 +385,20 @@ export async function addFile({
//#region Check drive usage //#region Check drive usage
if (user && !isLink) { if (user && !isLink) {
const usage = await DriveFiles.calcDriveUsageOf(user.id); const usage = await DriveFiles.calcDriveUsageOf(user.id);
const isLocalUser = Users.isLocalUser(user);
const instance = await fetchMeta(); const instance = await fetchMeta();
const driveCapacity = 1024 * 1024 * (Users.isLocalUser(user) ? instance.localDriveCapacityMb : instance.remoteDriveCapacityMb); const driveCapacity = 1024 * 1024 * (isLocalUser ? instance.localDriveCapacityMb : instance.remoteDriveCapacityMb);
logger.debug(`drive usage is ${usage} (max: ${driveCapacity})`); logger.debug(`drive usage is ${usage} (max: ${driveCapacity})`);
// If usage limit exceeded // If usage limit exceeded
if (usage + info.size > driveCapacity) { if (usage + info.size > driveCapacity) {
if (Users.isLocalUser(user)) { if (isLocalUser) {
throw new Error('no-free-space'); throw new Error('no-free-space');
} else { } else {
// delete oldest file (excluding banner and avatar) // delete older files to make space for new file
deleteOldFile(await Users.findOneByOrFail({ id: user.id }) as IRemoteUser); expireOldFiles(await Users.findOneByOrFail({ id: user.id }) as IRemoteUser, driveCapacity - info.size);
} }
} }
} }