From d6f65ea4142219679a5e8a2fff7c0e35ba941282 Mon Sep 17 00:00:00 2001 From: DarkSky Date: Thu, 11 Jan 2024 10:21:40 +0000 Subject: [PATCH] feat: blob size limit with quota (#5524) fix AFF-506 TOV-342 --- packages/backend/server/package.json | 1 + .../server/src/modules/quota/schema.ts | 17 +++++ .../server/src/modules/quota/storage.ts | 4 +- .../backend/server/src/modules/quota/types.ts | 8 +- .../src/modules/workspaces/resolvers/blob.ts | 75 +++++++++++++++---- .../modules/workspaces/resolvers/workspace.ts | 12 ++- packages/backend/server/src/schema.gql | 3 + .../server/tests/workspace-blobs.spec.ts | 56 ++++++++++++-- 8 files changed, 150 insertions(+), 26 deletions(-) diff --git a/packages/backend/server/package.json b/packages/backend/server/package.json index bc479fa867..7b85432795 100644 --- a/packages/backend/server/package.json +++ b/packages/backend/server/package.json @@ -114,6 +114,7 @@ "typescript": "^5.3.2" }, "ava": { + "timeout": "1m", "extensions": { "ts": "module" }, diff --git a/packages/backend/server/src/modules/quota/schema.ts b/packages/backend/server/src/modules/quota/schema.ts index a24f0db52a..10e4dbad79 100644 --- a/packages/backend/server/src/modules/quota/schema.ts +++ b/packages/backend/server/src/modules/quota/schema.ts @@ -37,6 +37,23 @@ export const Quotas: Quota[] = [ memberLimit: 10, }, }, + { + feature: QuotaType.RestrictedPlanV1, + type: FeatureKind.Quota, + version: 1, + configs: { + // quota name + name: 'Restricted', + // single blob limit 10MB + blobLimit: OneMB, + // total blob limit 1GB + storageQuota: 10 * OneMB, + // history period of validity 30 days + historyPeriod: 30 * OneDay, + // member limit 10 + memberLimit: 10, + }, + }, ]; export const Quota_FreePlanV1 = { diff --git a/packages/backend/server/src/modules/quota/storage.ts b/packages/backend/server/src/modules/quota/storage.ts index dde3620626..fa4fdcd28b 100644 --- a/packages/backend/server/src/modules/quota/storage.ts +++ b/packages/backend/server/src/modules/quota/storage.ts @@ -44,11 +44,11 @@ export class QuotaManagementService { const { user: owner } = await this.permissions.getWorkspaceOwner(workspaceId); if (!owner) throw new NotFoundException('Workspace owner not found'); - const { storageQuota } = await this.getUserQuota(owner.id); + const { storageQuota, blobLimit } = await this.getUserQuota(owner.id); // get all workspaces size of owner used const usageSize = await this.getUserUsage(owner.id); - return { quota: storageQuota, size: usageSize }; + return { quota: storageQuota, size: usageSize, limit: blobLimit }; } async checkBlobQuota(workspaceId: string, size: number) { diff --git a/packages/backend/server/src/modules/quota/types.ts b/packages/backend/server/src/modules/quota/types.ts index 6097d6bc57..5f94952713 100644 --- a/packages/backend/server/src/modules/quota/types.ts +++ b/packages/backend/server/src/modules/quota/types.ts @@ -8,10 +8,16 @@ import { ByteUnit, OneDay, OneKB } from './constant'; export enum QuotaType { FreePlanV1 = 'free_plan_v1', ProPlanV1 = 'pro_plan_v1', + // only for test, smaller quota + RestrictedPlanV1 = 'restricted_plan_v1', } const quotaPlan = z.object({ - feature: z.enum([QuotaType.FreePlanV1, QuotaType.ProPlanV1]), + feature: z.enum([ + QuotaType.FreePlanV1, + QuotaType.ProPlanV1, + QuotaType.RestrictedPlanV1, + ]), configs: z.object({ name: z.string(), blobLimit: z.number().positive().int(), diff --git a/packages/backend/server/src/modules/workspaces/resolvers/blob.ts b/packages/backend/server/src/modules/workspaces/resolvers/blob.ts index d2a03994ac..ae19169dcb 100644 --- a/packages/backend/server/src/modules/workspaces/resolvers/blob.ts +++ b/packages/backend/server/src/modules/workspaces/resolvers/blob.ts @@ -1,4 +1,4 @@ -import { ForbiddenException, Logger, UseGuards } from '@nestjs/common'; +import { HttpStatus, Logger, UseGuards } from '@nestjs/common'; import { Args, Float, @@ -9,12 +9,14 @@ import { ResolveField, Resolver, } from '@nestjs/graphql'; +import { GraphQLError } from 'graphql'; import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs'; import { MakeCache, PreventCache } from '../../../cache'; import { CloudThrottlerGuard } from '../../../throttler'; import type { FileUpload } from '../../../types'; import { Auth, CurrentUser } from '../../auth'; +import { FeatureManagementService, FeatureType } from '../../features'; import { QuotaManagementService } from '../../quota'; import { WorkspaceBlobStorage } from '../../storage'; import { UserType } from '../../users'; @@ -28,10 +30,26 @@ export class WorkspaceBlobResolver { logger = new Logger(WorkspaceBlobResolver.name); constructor( private readonly permissions: PermissionService, + private readonly feature: FeatureManagementService, private readonly quota: QuotaManagementService, private readonly storage: WorkspaceBlobStorage ) {} + @ResolveField(() => [String], { + description: 'List blobs of workspace', + complexity: 2, + }) + async blobs( + @CurrentUser() user: UserType, + @Parent() workspace: WorkspaceType + ) { + await this.permissions.checkWorkspace(workspace.id, user.id); + + return this.storage + .list(workspace.id) + .then(list => list.map(item => item.key)); + } + @ResolveField(() => Int, { description: 'Blobs size of workspace', complexity: 2, @@ -107,16 +125,30 @@ export class WorkspaceBlobResolver { Permission.Write ); - const { quota, size } = await this.quota.getWorkspaceUsage(workspaceId); + const { quota, size, limit } = + await this.quota.getWorkspaceUsage(workspaceId); + + const unlimited = await this.feature.hasWorkspaceFeature( + workspaceId, + FeatureType.UnlimitedWorkspace + ); const checkExceeded = (recvSize: number) => { if (!quota) { - throw new ForbiddenException('cannot find user quota'); + throw new GraphQLError('cannot find user quota', { + extensions: { + status: HttpStatus[HttpStatus.FORBIDDEN], + code: HttpStatus.FORBIDDEN, + }, + }); } - if (size + recvSize > quota) { - this.logger.log( - `storage size limit exceeded: ${size + recvSize} > ${quota}` - ); + const total = size + recvSize; + // only skip total storage check if workspace has unlimited feature + if (total > quota && !unlimited) { + this.logger.log(`storage size limit exceeded: ${total} > ${quota}`); + return true; + } else if (recvSize > limit) { + this.logger.log(`blob size limit exceeded: ${recvSize} > ${limit}`); return true; } else { return false; @@ -124,7 +156,12 @@ export class WorkspaceBlobResolver { }; if (checkExceeded(0)) { - throw new ForbiddenException('storage size limit exceeded'); + throw new GraphQLError('storage or blob size limit exceeded', { + extensions: { + status: HttpStatus[HttpStatus.PAYLOAD_TOO_LARGE], + code: HttpStatus.PAYLOAD_TOO_LARGE, + }, + }); } const buffer = await new Promise((resolve, reject) => { const stream = blob.createReadStream(); @@ -135,7 +172,14 @@ export class WorkspaceBlobResolver { // check size after receive each chunk to avoid unnecessary memory usage const bufferSize = chunks.reduce((acc, cur) => acc + cur.length, 0); if (checkExceeded(bufferSize)) { - reject(new ForbiddenException('storage size limit exceeded')); + reject( + new GraphQLError('storage or blob size limit exceeded', { + extensions: { + status: HttpStatus[HttpStatus.PAYLOAD_TOO_LARGE], + code: HttpStatus.PAYLOAD_TOO_LARGE, + }, + }) + ); } }); stream.on('error', reject); @@ -143,17 +187,20 @@ export class WorkspaceBlobResolver { const buffer = Buffer.concat(chunks); if (checkExceeded(buffer.length)) { - reject(new ForbiddenException('storage size limit exceeded')); + reject( + new GraphQLError('storage limit exceeded', { + extensions: { + status: HttpStatus[HttpStatus.PAYLOAD_TOO_LARGE], + code: HttpStatus.PAYLOAD_TOO_LARGE, + }, + }) + ); } else { resolve(buffer); } }); }); - if (!(await this.quota.checkBlobQuota(workspaceId, buffer.length))) { - throw new ForbiddenException('blob size limit exceeded'); - } - await this.storage.put(workspaceId, blob.filename, buffer); return blob.filename; } diff --git a/packages/backend/server/src/modules/workspaces/resolvers/workspace.ts b/packages/backend/server/src/modules/workspaces/resolvers/workspace.ts index 806952155b..328744e628 100644 --- a/packages/backend/server/src/modules/workspaces/resolvers/workspace.ts +++ b/packages/backend/server/src/modules/workspaces/resolvers/workspace.ts @@ -1,7 +1,6 @@ import { ForbiddenException, HttpStatus, - InternalServerErrorException, Logger, NotFoundException, UseGuards, @@ -398,8 +397,15 @@ export class WorkspaceResolver { `failed to send ${workspaceId} invite email to ${email}, but successfully revoked permission: ${e}` ); } - - return new InternalServerErrorException(e); + return new GraphQLError( + 'failed to send invite email, please try again', + { + extensions: { + status: HttpStatus[HttpStatus.INTERNAL_SERVER_ERROR], + code: HttpStatus.INTERNAL_SERVER_ERROR, + }, + } + ); } } return inviteId; diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index c16ac448b9..57b78e22d9 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -144,6 +144,9 @@ type WorkspaceType { publicPages: [WorkspacePage!]! histories(guid: String!, before: DateTime, take: Int): [DocHistoryType!]! + """List blobs of workspace""" + blobs: [String!]! + """Blobs size of workspace""" blobsSize: Int! } diff --git a/packages/backend/server/tests/workspace-blobs.spec.ts b/packages/backend/server/tests/workspace-blobs.spec.ts index 37abda1e45..4a1969ed59 100644 --- a/packages/backend/server/tests/workspace-blobs.spec.ts +++ b/packages/backend/server/tests/workspace-blobs.spec.ts @@ -3,6 +3,7 @@ import test from 'ava'; import request from 'supertest'; import { AppModule } from '../src/app'; +import { FeatureManagementService, FeatureType } from '../src/modules/features'; import { QuotaService, QuotaType } from '../src/modules/quota'; import { checkBlobSize, @@ -15,8 +16,11 @@ import { signUp, } from './utils'; +const OneMB = 1024 * 1024; + let app: INestApplication; let quota: QuotaService; +let feature: FeatureManagementService; test.beforeEach(async () => { const { app: testApp } = await createTestingApp({ @@ -25,6 +29,7 @@ test.beforeEach(async () => { app = testApp; quota = app.get(QuotaService); + feature = app.get(FeatureManagementService); }); test.afterEach.always(async () => { @@ -163,12 +168,51 @@ test('should be able calc quota after switch plan', async t => { t.is(size2, 0, 'failed to check pro plan blob size'); }); -test('should reject blob exceeded limit', t => { - // TODO - t.true(true); +test('should reject blob exceeded limit', async t => { + const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1'); + + const workspace1 = await createWorkspace(app, u1.token.token); + await quota.switchUserQuota(u1.id, QuotaType.RestrictedPlanV1); + + const buffer1 = Buffer.from(Array.from({ length: OneMB + 1 }, () => 0)); + await t.throwsAsync(setBlob(app, u1.token.token, workspace1.id, buffer1)); + + await quota.switchUserQuota(u1.id, QuotaType.FreePlanV1); + + const buffer2 = Buffer.from(Array.from({ length: OneMB + 1 }, () => 0)); + await t.notThrowsAsync(setBlob(app, u1.token.token, workspace1.id, buffer2)); + + const buffer3 = Buffer.from(Array.from({ length: 10 * OneMB + 1 }, () => 0)); + await t.throwsAsync(setBlob(app, u1.token.token, workspace1.id, buffer3)); }); -test('should reject blob exceeded quota', t => { - // TODO - t.true(true); +test('should reject blob exceeded quota', async t => { + const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1'); + + const workspace = await createWorkspace(app, u1.token.token); + await quota.switchUserQuota(u1.id, QuotaType.RestrictedPlanV1); + + const buffer = Buffer.from(Array.from({ length: OneMB }, () => 0)); + + for (let i = 0; i < 10; i++) { + await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); + } + + await t.throwsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); +}); + +test('should accept blob even storage out of quota if workspace has unlimited feature', async t => { + const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1'); + + const workspace = await createWorkspace(app, u1.token.token); + await quota.switchUserQuota(u1.id, QuotaType.RestrictedPlanV1); + feature.addWorkspaceFeatures(workspace.id, FeatureType.UnlimitedWorkspace); + + const buffer = Buffer.from(Array.from({ length: OneMB }, () => 0)); + + for (let i = 0; i < 10; i++) { + await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); + } + + await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); });