feat: add user level blob quota (#4114)

apps/server/src/config/def.ts

@@ -186,6 +186,11 @@ export interface AFFiNEConfig {
     fs: {
       path: string;
     };
+    /**
+     * Free user storage quota
+     * @default 10 * 1024 * 1024 (10GB)
+     */
+    quota: number;
   };
 
   /**

apps/server/src/config/default.ts

@@ -55,6 +55,7 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => {
       AFFINE_SERVER_HOST: 'host',
       AFFINE_SERVER_SUB_PATH: 'path',
       AFFINE_ENV: 'affineEnv',
+      AFFINE_FREE_USER_QUOTA: 'objectStorage.quota',
       DATABASE_URL: 'db.url',
       ENABLE_R2_OBJECT_STORAGE: ['objectStorage.r2.enabled', 'boolean'],
       R2_OBJECT_STORAGE_ACCOUNT_ID: 'objectStorage.r2.accountId',
@@ -170,6 +171,7 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => {
       fs: {
         path: join(homedir(), '.affine-storage'),
       },
+      quota: 10 * 1024 * 1024,
     },
     rateLimiter: {
       ttl: 60,

apps/server/src/modules/workspaces/resolver.ts

@@ -27,6 +27,7 @@ import type { User, Workspace } from '@prisma/client';
 import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
 import { applyUpdate, Doc } from 'yjs';
 
+import { Config } from '../../config';
 import { PrismaService } from '../../prisma';
 import { StorageProvide } from '../../storage';
 import { CloudThrottlerGuard, Throttle } from '../../throttler';
@@ -130,6 +131,7 @@ export class UpdateWorkspaceInput extends PickType(
 export class WorkspaceResolver {
   constructor(
     private readonly auth: AuthService,
+    private readonly config: Config,
     private readonly mailer: MailService,
     private readonly prisma: PrismaService,
     private readonly permissionProvider: PermissionService,
@@ -610,17 +612,28 @@ export class WorkspaceResolver {
   ) {
     await this.permissionProvider.check(workspaceId, user.id);
 
-    return this.storage.blobsSize(workspaceId).then(size => ({ size }));
+    return this.storage.blobsSize([workspaceId]).then(size => ({ size }));
   }
 
   @Query(() => WorkspaceBlobSizes)
-  async collectAllBlobSizes(@CurrentUser() user: User) {
-    const workspaces = await this.workspaces(user);
-
-    const size = (
-      await Promise.all(workspaces.map(({ id }) => this.storage.blobsSize(id)))
-    ).reduce((prev, curr) => prev + curr, 0);
+  async collectAllBlobSizes(@CurrentUser() user: UserType) {
+    const workspaces = await this.prisma.userWorkspacePermission
+      .findMany({
+        where: {
+          userId: user.id,
+          accepted: true,
+        },
+        select: {
+          workspace: {
+            select: {
+              id: true,
+            },
+          },
+        },
+      })
+      .then(data => data.map(({ workspace }) => workspace.id));
 
+    const size = await this.storage.blobsSize(workspaces);
     return { size };
   }
 
@@ -632,6 +645,12 @@ export class WorkspaceResolver {
     blob: FileUpload
   ) {
     await this.permissionProvider.check(workspaceId, user.id, Permission.Write);
+    const quota = this.config.objectStorage.quota;
+    const { size } = await this.collectAllBlobSizes(user);
+
+    if (size > quota) {
+      throw new ForbiddenException('storage size limit exceeded');
+    }
 
     const buffer = await new Promise<Buffer>((resolve, reject) => {
       const stream = blob.createReadStream();
@@ -645,6 +664,10 @@ export class WorkspaceResolver {
       });
     });
 
+    if (size + buffer.length > quota) {
+      throw new ForbiddenException('storage size limit exceeded');
+    }
+
     return this.storage.uploadBlob(workspaceId, buffer);
   }
 

apps/server/src/tests/utils.ts

@@ -339,7 +339,6 @@ async function collectBlobSizes(
   const res = await request(app.getHttpServer())
     .post(gql)
     .auth(token, { type: 'bearer' })
-    .set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
     .send({
       query: `
             query {
@@ -353,6 +352,26 @@ async function collectBlobSizes(
   return res.body.data.collectBlobSizes.size;
 }
 
+async function collectAllBlobSizes(
+  app: INestApplication,
+  token: string
+): Promise<number> {
+  const res = await request(app.getHttpServer())
+    .post(gql)
+    .auth(token, { type: 'bearer' })
+    .send({
+      query: `
+            query {
+              collectAllBlobSizes {
+                size
+              }
+            }
+          `,
+    })
+    .expect(200);
+  return res.body.data.collectAllBlobSizes.size;
+}
+
 async function setBlob(
   app: INestApplication,
   token: string,
@@ -447,6 +466,7 @@ async function getInviteInfo(
 export {
   acceptInvite,
   acceptInviteById,
+  collectAllBlobSizes,
   collectBlobSizes,
   createTestApp,
   createWorkspace,

apps/server/src/tests/workspace-blobs.spec.ts

@@ -10,6 +10,7 @@ import request from 'supertest';
 
 import { AppModule } from '../app';
 import {
+  collectAllBlobSizes,
   collectBlobSizes,
   createWorkspace,
   listBlobs,
@@ -108,4 +109,25 @@ describe('Workspace Module - Blobs', () => {
     const size = await collectBlobSizes(app, u1.token.token, workspace.id);
     ok(size === 4, 'failed to collect blob sizes');
   });
+
+  it('should calc all blobs size', async () => {
+    const u1 = await signUp(app, 'u1', 'u1@affine.pro', '1');
+
+    const workspace1 = await createWorkspace(app, u1.token.token);
+
+    const buffer1 = Buffer.from([0, 0]);
+    await setBlob(app, u1.token.token, workspace1.id, buffer1);
+    const buffer2 = Buffer.from([0, 1]);
+    await setBlob(app, u1.token.token, workspace1.id, buffer2);
+
+    const workspace2 = await createWorkspace(app, u1.token.token);
+
+    const buffer3 = Buffer.from([0, 0]);
+    await setBlob(app, u1.token.token, workspace2.id, buffer3);
+    const buffer4 = Buffer.from([0, 1]);
+    await setBlob(app, u1.token.token, workspace2.id, buffer4);
+
+    const size = await collectAllBlobSizes(app, u1.token.token);
+    ok(size === 8, 'failed to collect all blob sizes');
+  });
 });