refactor(server): rate limit and permission (#4198)

Co-authored-by: LongYinan <lynweklm@gmail.com>

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

@@ -9,6 +9,7 @@ import {
 import {
   Args,
   Field,
+  Float,
   ID,
   InputType,
   Int,
@@ -36,6 +37,7 @@ import type { FileUpload } from '../../types';
 import { Auth, CurrentUser, Public } from '../auth';
 import { MailService } from '../auth/mailer';
 import { AuthService } from '../auth/service';
+import { UsersService } from '../users';
 import { UserType } from '../users/resolver';
 import { PermissionService } from './permission';
 import { Permission } from './types';
@@ -139,7 +141,8 @@ export class WorkspaceResolver {
     private readonly config: Config,
     private readonly mailer: MailService,
     private readonly prisma: PrismaService,
-    private readonly permissionProvider: PermissionService,
+    private readonly permissions: PermissionService,
+    private readonly users: UsersService,
     @Inject(StorageProvide) private readonly storage: Storage
   ) {}
 
@@ -156,7 +159,7 @@ export class WorkspaceResolver {
       return workspace.permission;
     }
 
-    const permission = await this.permissionProvider.get(workspace.id, user.id);
+    const permission = await this.permissions.get(workspace.id, user.id);
 
     if (!permission) {
       throw new ForbiddenException();
@@ -197,15 +200,7 @@ export class WorkspaceResolver {
     complexity: 2,
   })
   async owner(@Parent() workspace: WorkspaceType) {
-    const data = await this.prisma.userWorkspacePermission.findFirstOrThrow({
-      where: {
-        workspaceId: workspace.id,
-        type: Permission.Owner,
-      },
-      include: {
-        user: true,
-      },
-    });
+    const data = await this.permissions.getWorkspaceOwner(workspace.id);
 
     return data.user;
   }
@@ -241,15 +236,7 @@ export class WorkspaceResolver {
     @CurrentUser() user: UserType,
     @Args('workspaceId') workspaceId: string
   ) {
-    const data = await this.prisma.userWorkspacePermission.findFirst({
-      where: {
-        workspaceId,
-        type: Permission.Owner,
-      },
-      include: {
-        user: true,
-      },
-    });
+    const data = await this.permissions.tryGetWorkspaceOwner(workspaceId);
 
     return data?.user?.id === user.id;
   }
@@ -298,7 +285,7 @@ export class WorkspaceResolver {
     description: 'Get workspace by id',
   })
   async workspace(@CurrentUser() user: UserType, @Args('id') id: string) {
-    await this.permissionProvider.check(id, user.id);
+    await this.permissions.check(id, user.id);
     const workspace = await this.prisma.workspace.findUnique({ where: { id } });
 
     if (!workspace) {
@@ -367,7 +354,7 @@ export class WorkspaceResolver {
     @Args({ name: 'input', type: () => UpdateWorkspaceInput })
     { id, ...updates }: UpdateWorkspaceInput
   ) {
-    await this.permissionProvider.check(id, user.id, Permission.Admin);
+    await this.permissions.check(id, user.id, Permission.Admin);
 
     return this.prisma.workspace.update({
       where: {
@@ -379,7 +366,7 @@ export class WorkspaceResolver {
 
   @Mutation(() => Boolean)
   async deleteWorkspace(@CurrentUser() user: UserType, @Args('id') id: string) {
-    await this.permissionProvider.check(id, user.id, Permission.Owner);
+    await this.permissions.check(id, user.id, Permission.Owner);
 
     await this.prisma.workspace.delete({
       where: {
@@ -411,17 +398,13 @@ export class WorkspaceResolver {
     @Args('permission', { type: () => Permission }) permission: Permission,
     @Args('sendInviteMail', { nullable: true }) sendInviteMail: boolean
   ) {
-    await this.permissionProvider.check(workspaceId, user.id, Permission.Admin);
+    await this.permissions.check(workspaceId, user.id, Permission.Admin);
 
     if (permission === Permission.Owner) {
       throw new ForbiddenException('Cannot change owner');
     }
 
-    const target = await this.prisma.user.findUnique({
-      where: {
-        email,
-      },
-    });
+    const target = await this.users.findUserByEmail(email);
 
     if (target) {
       const originRecord = await this.prisma.userWorkspacePermission.findFirst({
@@ -435,7 +418,7 @@ export class WorkspaceResolver {
         return originRecord.id;
       }
 
-      const inviteId = await this.permissionProvider.grant(
+      const inviteId = await this.permissions.grant(
         workspaceId,
         target.id,
         permission
@@ -458,7 +441,7 @@ export class WorkspaceResolver {
       return inviteId;
     } else {
       const user = await this.auth.createAnonymousUser(email);
-      const inviteId = await this.permissionProvider.grant(
+      const inviteId = await this.permissions.grant(
         workspaceId,
         user.id,
         permission
@@ -488,17 +471,21 @@ export class WorkspaceResolver {
     description: 'Update workspace',
   })
   async getInviteInfo(@Args('inviteId') inviteId: string) {
-    const permission =
-      await this.prisma.userWorkspacePermission.findUniqueOrThrow({
+    const workspaceId = await this.prisma.userWorkspacePermission
+      .findUniqueOrThrow({
         where: {
           id: inviteId,
         },
-      });
+        select: {
+          workspaceId: true,
+        },
+      })
+      .then(({ workspaceId }) => workspaceId);
 
     const snapshot = await this.prisma.snapshot.findFirstOrThrow({
       where: {
-        id: permission.workspaceId,
-        workspaceId: permission.workspaceId,
+        id: workspaceId,
+        workspaceId,
       },
     });
 
@@ -507,32 +494,17 @@ export class WorkspaceResolver {
     applyUpdate(doc, new Uint8Array(snapshot.blob));
     const metaJSON = doc.getMap('meta').toJSON();
 
-    const owner = await this.prisma.userWorkspacePermission.findFirstOrThrow({
-      where: {
-        workspaceId: permission.workspaceId,
-        type: Permission.Owner,
-      },
-      include: {
-        user: true,
-      },
-    });
-    const invitee = await this.prisma.userWorkspacePermission.findUniqueOrThrow(
-      {
-        where: {
-          id: inviteId,
-          workspaceId: permission.workspaceId,
-        },
-        include: {
-          user: true,
-        },
-      }
+    const owner = await this.permissions.getWorkspaceOwner(workspaceId);
+    const invitee = await this.permissions.getInvitationById(
+      inviteId,
+      workspaceId
     );
 
     let avatar = '';
 
     if (metaJSON.avatar) {
       const avatarBlob = await this.storage.getBlob(
-        permission.workspaceId,
+        workspaceId,
         metaJSON.avatar
       );
       avatar = avatarBlob?.data.toString('base64') || '';
@@ -542,7 +514,7 @@ export class WorkspaceResolver {
       workspace: {
         name: metaJSON.name || '',
         avatar: avatar || defaultWorkspaceAvatar,
-        id: permission.workspaceId,
+        id: workspaceId,
       },
       user: owner.user,
       invitee: invitee.user,
@@ -555,9 +527,9 @@ export class WorkspaceResolver {
     @Args('workspaceId') workspaceId: string,
     @Args('userId') userId: string
   ) {
-    await this.permissionProvider.check(workspaceId, user.id, Permission.Admin);
+    await this.permissions.check(workspaceId, user.id, Permission.Admin);
 
-    return this.permissionProvider.revoke(workspaceId, userId);
+    return this.permissions.revoke(workspaceId, userId);
   }
 
   @Mutation(() => Boolean)
@@ -586,7 +558,7 @@ export class WorkspaceResolver {
       });
     }
 
-    return this.permissionProvider.acceptById(workspaceId, inviteId);
+    return this.permissions.acceptById(workspaceId, inviteId);
   }
 
   @Mutation(() => Boolean)
@@ -594,7 +566,7 @@ export class WorkspaceResolver {
     @CurrentUser() user: UserType,
     @Args('workspaceId') workspaceId: string
   ) {
-    return this.permissionProvider.accept(workspaceId, user.id);
+    return this.permissions.accept(workspaceId, user.id);
   }
 
   @Mutation(() => Boolean)
@@ -604,17 +576,9 @@ export class WorkspaceResolver {
     @Args('workspaceName') workspaceName: string,
     @Args('sendLeaveMail', { nullable: true }) sendLeaveMail: boolean
   ) {
-    await this.permissionProvider.check(workspaceId, user.id);
+    await this.permissions.check(workspaceId, user.id);
 
-    const owner = await this.prisma.userWorkspacePermission.findFirstOrThrow({
-      where: {
-        workspaceId,
-        type: Permission.Owner,
-      },
-      include: {
-        user: true,
-      },
-    });
+    const owner = await this.permissions.getWorkspaceOwner(workspaceId);
 
     if (!owner.user) {
       throw new ForbiddenException(
@@ -629,7 +593,7 @@ export class WorkspaceResolver {
       });
     }
 
-    return this.permissionProvider.revoke(workspaceId, user.id);
+    return this.permissions.revoke(workspaceId, user.id);
   }
 
   @Mutation(() => Boolean)
@@ -638,9 +602,9 @@ export class WorkspaceResolver {
     @Args('workspaceId') workspaceId: string,
     @Args('pageId') pageId: string
   ) {
-    await this.permissionProvider.check(workspaceId, user.id, Permission.Admin);
+    await this.permissions.check(workspaceId, user.id, Permission.Admin);
 
-    return this.permissionProvider.grantPage(workspaceId, pageId);
+    return this.permissions.grantPage(workspaceId, pageId);
   }
 
   @Mutation(() => Boolean)
@@ -649,9 +613,9 @@ export class WorkspaceResolver {
     @Args('workspaceId') workspaceId: string,
     @Args('pageId') pageId: string
   ) {
-    await this.permissionProvider.check(workspaceId, user.id, Permission.Admin);
+    await this.permissions.check(workspaceId, user.id, Permission.Admin);
 
-    return this.permissionProvider.revokePage(workspaceId, pageId);
+    return this.permissions.revokePage(workspaceId, pageId);
   }
 
   @Query(() => [String], {
@@ -661,7 +625,7 @@ export class WorkspaceResolver {
     @CurrentUser() user: UserType,
     @Args('workspaceId') workspaceId: string
   ) {
-    await this.permissionProvider.check(workspaceId, user.id);
+    await this.permissions.check(workspaceId, user.id);
 
     return this.storage.listBlobs(workspaceId);
   }
@@ -671,7 +635,7 @@ export class WorkspaceResolver {
     @CurrentUser() user: UserType,
     @Args('workspaceId') workspaceId: string
   ) {
-    await this.permissionProvider.check(workspaceId, user.id);
+    await this.permissions.check(workspaceId, user.id);
 
     return this.storage.blobsSize([workspaceId]).then(size => ({ size }));
   }
@@ -699,6 +663,29 @@ export class WorkspaceResolver {
     return { size };
   }
 
+  @Query(() => WorkspaceBlobSizes)
+  async checkBlobSize(
+    @CurrentUser() user: UserType,
+    @Args('workspaceId') workspaceId: string,
+    @Args('size', { type: () => Float }) size: number
+  ) {
+    const canWrite = await this.permissions.tryCheck(
+      workspaceId,
+      user.id,
+      Permission.Write
+    );
+    if (canWrite) {
+      const { user } = await this.permissions.getWorkspaceOwner(workspaceId);
+      if (user) {
+        const quota = await this.users.getStorageQuotaById(user.id);
+        const { size: currentSize } = await this.collectAllBlobSizes(user);
+
+        return { size: quota - (size + currentSize) };
+      }
+    }
+    return false;
+  }
+
   @Mutation(() => String)
   async setBlob(
     @CurrentUser() user: UserType,
@@ -706,36 +693,53 @@ export class WorkspaceResolver {
     @Args({ name: 'blob', type: () => GraphQLUpload })
     blob: FileUpload
   ) {
-    await this.permissionProvider.check(workspaceId, user.id, Permission.Write);
-    const quota = this.config.objectStorage.quota;
-    const { size } = await this.collectAllBlobSizes(user);
+    await this.permissions.check(workspaceId, user.id, Permission.Write);
 
-    if (size > quota) {
-      this.logger.log(`storage size limit exceeded: ${size} > ${quota}`);
+    // quota was apply to owner's account
+    const { user: owner } =
+      await this.permissions.getWorkspaceOwner(workspaceId);
+    if (!owner) return new NotFoundException('Workspace owner not found');
+    const quota = await this.users.getStorageQuotaById(owner.id);
+    const { size } = await this.collectAllBlobSizes(owner);
+
+    const checkExceeded = (recvSize: number) => {
+      if (size + recvSize > quota) {
+        this.logger.log(
+          `storage size limit exceeded: ${size + recvSize} > ${quota}`
+        );
+        return true;
+      } else {
+        return false;
+      }
+    };
+
+    if (checkExceeded(0)) {
       throw new ForbiddenException('storage size limit exceeded');
     }
-
     const buffer = await new Promise<Buffer>((resolve, reject) => {
       const stream = blob.createReadStream();
       const chunks: Uint8Array[] = [];
       stream.on('data', chunk => {
         chunks.push(chunk);
+
+        // 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'));
+        }
       });
       stream.on('error', reject);
       stream.on('end', () => {
-        resolve(Buffer.concat(chunks));
+        const buffer = Buffer.concat(chunks);
+
+        if (checkExceeded(buffer.length)) {
+          reject(new ForbiddenException('storage size limit exceeded'));
+        } else {
+          resolve(buffer);
+        }
       });
     });
 
-    if (size + buffer.length > quota) {
-      this.logger.log(
-        `storage size limit exceeded after blob set: ${size} > ${
-          buffer.length > quota
-        }`
-      );
-      throw new ForbiddenException('storage size limit exceeded');
-    }
-
     return this.storage.uploadBlob(workspaceId, buffer);
   }
 
@@ -745,7 +749,7 @@ export class WorkspaceResolver {
     @Args('workspaceId') workspaceId: string,
     @Args('hash') hash: string
   ) {
-    await this.permissionProvider.check(workspaceId, user.id);
+    await this.permissions.check(workspaceId, user.id);
 
     return this.storage.deleteBlob(workspaceId, hash);
   }