feat: blob size limit with quota (#5524)

fix AFF-506 TOV-342

packages/backend/server/package.json

@@ -114,6 +114,7 @@
     "typescript": "^5.3.2"
   },
   "ava": {
+    "timeout": "1m",
     "extensions": {
       "ts": "module"
     },

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 = {

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) {

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(),

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<Buffer>((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;
   }

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;

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!
 }

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));
 });