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

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

apps/server/migrations/20230906100042_user_feature_gates/migration.sql

@@ -0,0 +1,13 @@
+-- CreateTable
+CREATE TABLE "user_feature_gates" (
+    "id" VARCHAR NOT NULL,
+    "user_id" VARCHAR NOT NULL,
+    "feature" VARCHAR NOT NULL,
+    "reason" VARCHAR NOT NULL,
+    "created_at" TIMESTAMPTZ(6) NOT NULL DEFAULT CURRENT_TIMESTAMP,
+
+    CONSTRAINT "user_feature_gates_pkey" PRIMARY KEY ("id")
+);
+
+-- AddForeignKey
+ALTER TABLE "user_feature_gates" ADD CONSTRAINT "user_feature_gates_user_id_fkey" FOREIGN KEY ("user_id") REFERENCES "users"("id") ON DELETE CASCADE ON UPDATE CASCADE;

apps/server/schema.prisma

@@ -1,6 +1,6 @@
 generator client {
-  provider      = "prisma-client-js"
+  provider        = "prisma-client-js"
-  binaryTargets = ["native", "debian-openssl-3.0.x"]
+  binaryTargets   = ["native", "debian-openssl-3.0.x"]
   previewFeatures = ["metrics", "tracing"]
 }
 
@@ -48,10 +48,22 @@ model User {
   createdAt     DateTime                  @default(now()) @map("created_at") @db.Timestamptz(6)
   /// Not available if user signed up through OAuth providers
   password      String?                   @db.VarChar
+  features      UserFeatureGates[]
 
   @@map("users")
 }
 
+model UserFeatureGates {
+  id        String   @id @default(uuid()) @db.VarChar
+  userId    String   @map("user_id") @db.VarChar
+  feature   String   @db.VarChar
+  reason    String   @db.VarChar
+  createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(6)
+  user      User     @relation(fields: [userId], references: [id], onDelete: Cascade)
+
+  @@map("user_feature_gates")
+}
+
 model Account {
   id                String  @id @default(cuid())
   userId            String  @map("user_id")

apps/server/src/config/def.ts

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

apps/server/src/config/default.ts

@@ -175,7 +175,8 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => {
       fs: {
         path: join(homedir(), '.affine-storage'),
       },
-      quota: 10 * 1024 * 1024,
+      // 10GB
+      quota: 10 * 1024 * 1024 * 1024,
     },
     rateLimiter: {
       ttl: 60,

apps/server/src/modules/auth/next-auth-options.ts

@@ -1,16 +1,10 @@
-import { randomUUID } from 'node:crypto';
-
 import { PrismaAdapter } from '@auth/prisma-adapter';
-import { BadRequestException, FactoryProvider, Logger } from '@nestjs/common';
+import { FactoryProvider, Logger } from '@nestjs/common';
 import { verify } from '@node-rs/argon2';
-import { Algorithm, sign, verify as jwtVerify } from '@node-rs/jsonwebtoken';
 import { assign, omit } from 'lodash-es';
-import { nanoid } from 'nanoid';
 import { NextAuthOptions } from 'next-auth';
 import Credentials from 'next-auth/providers/credentials';
-import Email, {
+import Email from 'next-auth/providers/email';
-  type SendVerificationRequestParams,
-} from 'next-auth/providers/email';
 import Github from 'next-auth/providers/github';
 import Google from 'next-auth/providers/google';
 
@@ -20,7 +14,12 @@ import { SessionService } from '../../session';
 import { NewFeaturesKind } from '../users/types';
 import { isStaff } from '../users/utils';
 import { MailService } from './mailer';
-import { getUtcTimestamp, UserClaim } from './service';
+import {
+  decode,
+  encode,
+  sendVerificationRequest,
+  SendVerificationRequestParams,
+} from './utils';
 
 export const NextAuthOptionsProvide = Symbol('NextAuthOptions');
 
@@ -78,44 +77,8 @@ export const NextAuthOptionsProvider: FactoryProvider<NextAuthOptions> = {
             },
           },
           from: config.auth.email.sender,
-          async sendVerificationRequest(params: SendVerificationRequestParams) {
+          sendVerificationRequest: (params: SendVerificationRequestParams) =>
-            const { identifier, url, provider } = params;
+            sendVerificationRequest(config, logger, mailer, session, params),
-            const urlWithToken = new URL(url);
-            const callbackUrl =
-              urlWithToken.searchParams.get('callbackUrl') || '';
-            if (!callbackUrl) {
-              throw new Error('callbackUrl is not set');
-            } else {
-              const newCallbackUrl = new URL(callbackUrl, config.origin);
-
-              const token = nanoid();
-              await session.set(token, identifier);
-              newCallbackUrl.searchParams.set('token', token);
-
-              urlWithToken.searchParams.set(
-                'callbackUrl',
-                newCallbackUrl.toString()
-              );
-            }
-
-            const result = await mailer.sendSignInEmail(
-              urlWithToken.toString(),
-              {
-                to: identifier,
-                from: provider.from,
-              }
-            );
-            logger.log(
-              `send verification email success: ${result.accepted.join(', ')}`
-            );
-
-            const failed = result.rejected
-              .concat(result.pending)
-              .filter(Boolean);
-            if (failed.length) {
-              throw new Error(`Email (${failed.join(', ')}) could not be sent`);
-            }
-          },
         }),
       ],
       // @ts-expect-error Third part library type mismatch
@@ -200,66 +163,9 @@ export const NextAuthOptionsProvider: FactoryProvider<NextAuthOptions> = {
     }
 
     nextAuthOptions.jwt = {
-      encode: async ({ token, maxAge }) => {
+      encode: async ({ token, maxAge }) =>
-        if (!token?.email) {
+        encode(config, prisma, token, maxAge),
-          throw new BadRequestException('Missing email in jwt token');
+      decode: async ({ token }) => decode(config, token),
-        }
-        const user = await prisma.user.findFirstOrThrow({
-          where: {
-            email: token.email,
-          },
-        });
-        const now = getUtcTimestamp();
-        return sign(
-          {
-            data: {
-              id: user.id,
-              name: user.name,
-              email: user.email,
-              emailVerified: user.emailVerified?.toISOString(),
-              picture: user.avatarUrl,
-              createdAt: user.createdAt.toISOString(),
-              hasPassword: Boolean(user.password),
-            },
-            iat: now,
-            exp: now + (maxAge ?? config.auth.accessTokenExpiresIn),
-            iss: config.serverId,
-            sub: user.id,
-            aud: user.name,
-            jti: randomUUID({
-              disableEntropyCache: true,
-            }),
-          },
-          config.auth.privateKey,
-          {
-            algorithm: Algorithm.ES256,
-          }
-        );
-      },
-      decode: async ({ token }) => {
-        if (!token) {
-          return null;
-        }
-        const { name, email, emailVerified, id, picture, hasPassword } = (
-          await jwtVerify(token, config.auth.publicKey, {
-            algorithms: [Algorithm.ES256],
-            iss: [config.serverId],
-            leeway: config.auth.leeway,
-            requiredSpecClaims: ['exp', 'iat', 'iss', 'sub'],
-          })
-        ).data as Omit<UserClaim, 'avatarUrl'> & {
-          picture: string | undefined;
-        };
-        return {
-          name,
-          email,
-          emailVerified,
-          picture,
-          sub: id,
-          id,
-          hasPassword,
-        };
-      },
     };
     nextAuthOptions.secret ??= config.auth.nextAuthSecret;
 

apps/server/src/modules/auth/next-auth.controller.ts

@@ -23,7 +23,7 @@ import { AuthHandler } from 'next-auth/core';
 import { Config } from '../../config';
 import { Metrics } from '../../metrics/metrics';
 import { PrismaService } from '../../prisma/service';
-import { CloudThrottlerGuard, Throttle } from '../../throttler';
+import { AuthThrottlerGuard, Throttle } from '../../throttler';
 import { NextAuthOptionsProvide } from './next-auth-options';
 import { AuthService } from './service';
 
@@ -47,7 +47,7 @@ export class NextAuthController {
     this.callbackSession = nextAuthOptions.callbacks!.session;
   }
 
-  @UseGuards(CloudThrottlerGuard)
+  @UseGuards(AuthThrottlerGuard)
   @Throttle(60, 60)
   @All('*')
   async auth(
@@ -197,7 +197,7 @@ export class NextAuthController {
       throw new BadRequestException(`Invalid new session data`);
     }
     if (password) {
-      const user = await this.getUserFromRequest(req);
+      const user = await this.verifyUserFromRequest(req);
       const { password: userPassword } = user;
       if (!oldPassword) {
         if (userPassword) {
@@ -223,7 +223,7 @@ export class NextAuthController {
       }
       return user;
     } else {
-      const user = await this.getUserFromRequest(req);
+      const user = await this.verifyUserFromRequest(req);
       return this.prisma.user.update({
         where: {
           id: user.id,
@@ -233,7 +233,7 @@ export class NextAuthController {
     }
   }
 
-  private async getUserFromRequest(req: Request): Promise<User> {
+  private async verifyUserFromRequest(req: Request): Promise<User> {
     const token = req.headers.authorization;
     if (!token) {
       const session = await AuthHandler({

apps/server/src/modules/auth/utils/index.ts

@@ -0,0 +1,3 @@
+export { jwtDecode as decode, jwtEncode as encode } from './jwt';
+export { sendVerificationRequest } from './send-mail';
+export type { SendVerificationRequestParams } from 'next-auth/providers/email';

apps/server/src/modules/auth/utils/jwt.ts

@@ -0,0 +1,76 @@
+import { randomUUID } from 'node:crypto';
+
+import { BadRequestException } from '@nestjs/common';
+import { Algorithm, sign, verify as jwtVerify } from '@node-rs/jsonwebtoken';
+import { JWT } from 'next-auth/jwt';
+
+import { Config } from '../../../config';
+import { PrismaService } from '../../../prisma';
+import { getUtcTimestamp, UserClaim } from '../service';
+
+export const jwtEncode = async (
+  config: Config,
+  prisma: PrismaService,
+  token: JWT | undefined,
+  maxAge: number | undefined
+) => {
+  if (!token?.email) {
+    throw new BadRequestException('Missing email in jwt token');
+  }
+  const user = await prisma.user.findFirstOrThrow({
+    where: {
+      email: token.email,
+    },
+  });
+  const now = getUtcTimestamp();
+  return sign(
+    {
+      data: {
+        id: user.id,
+        name: user.name,
+        email: user.email,
+        emailVerified: user.emailVerified?.toISOString(),
+        picture: user.avatarUrl,
+        createdAt: user.createdAt.toISOString(),
+        hasPassword: Boolean(user.password),
+      },
+      iat: now,
+      exp: now + (maxAge ?? config.auth.accessTokenExpiresIn),
+      iss: config.serverId,
+      sub: user.id,
+      aud: user.name,
+      jti: randomUUID({
+        disableEntropyCache: true,
+      }),
+    },
+    config.auth.privateKey,
+    {
+      algorithm: Algorithm.ES256,
+    }
+  );
+};
+
+export const jwtDecode = async (config: Config, token: string | undefined) => {
+  if (!token) {
+    return null;
+  }
+  const { name, email, emailVerified, id, picture, hasPassword } = (
+    await jwtVerify(token, config.auth.publicKey, {
+      algorithms: [Algorithm.ES256],
+      iss: [config.serverId],
+      leeway: config.auth.leeway,
+      requiredSpecClaims: ['exp', 'iat', 'iss', 'sub'],
+    })
+  ).data as Omit<UserClaim, 'avatarUrl'> & {
+    picture: string | undefined;
+  };
+  return {
+    name,
+    email,
+    emailVerified,
+    picture,
+    sub: id,
+    id,
+    hasPassword,
+  };
+};

apps/server/src/modules/auth/utils/send-mail.ts

@@ -0,0 +1,41 @@
+import { Logger } from '@nestjs/common';
+import { nanoid } from 'nanoid';
+import type { SendVerificationRequestParams } from 'next-auth/providers/email';
+
+import { Config } from '../../../config';
+import { SessionService } from '../../../session';
+import { MailService } from '../mailer';
+
+export async function sendVerificationRequest(
+  config: Config,
+  logger: Logger,
+  mailer: MailService,
+  session: SessionService,
+  params: SendVerificationRequestParams
+) {
+  const { identifier, url, provider } = params;
+  const urlWithToken = new URL(url);
+  const callbackUrl = urlWithToken.searchParams.get('callbackUrl') || '';
+  if (!callbackUrl) {
+    throw new Error('callbackUrl is not set');
+  } else {
+    const newCallbackUrl = new URL(callbackUrl, config.origin);
+
+    const token = nanoid();
+    await session.set(token, identifier);
+    newCallbackUrl.searchParams.set('token', token);
+
+    urlWithToken.searchParams.set('callbackUrl', newCallbackUrl.toString());
+  }
+
+  const result = await mailer.sendSignInEmail(urlWithToken.toString(), {
+    to: identifier,
+    from: provider.from,
+  });
+  logger.log(`send verification email success: ${result.accepted.join(', ')}`);
+
+  const failed = result.rejected.concat(result.pending).filter(Boolean);
+  if (failed.length) {
+    throw new Error(`Email (${failed.join(', ')}) could not be sent`);
+  }
+}

apps/server/src/modules/users/gates.ts

@@ -0,0 +1,42 @@
+type FeatureEarlyAccessPreview = {
+  whitelist: RegExp[];
+};
+
+type FeatureStorageLimit = {
+  storageQuota: number;
+};
+
+type UserFeatureGate = {
+  earlyAccessPreview: FeatureEarlyAccessPreview;
+  freeUser: FeatureStorageLimit;
+  proUser: FeatureStorageLimit;
+};
+
+const UserLevel = {
+  freeUser: {
+    storageQuota: 10 * 1024 * 1024 * 1024,
+  },
+  proUser: {
+    storageQuota: 100 * 1024 * 1024 * 1024,
+  },
+} satisfies Pick<UserFeatureGate, 'freeUser' | 'proUser'>;
+
+export function getStorageQuota(features: string[]) {
+  for (const feature of features) {
+    if (feature in UserLevel) {
+      return UserLevel[feature as keyof typeof UserLevel].storageQuota;
+    }
+  }
+  return null;
+}
+
+const UserType = {
+  earlyAccessPreview: {
+    whitelist: [/@toeverything\.info$/],
+  },
+} satisfies Pick<UserFeatureGate, 'earlyAccessPreview'>;
+
+export const FeatureGates = {
+  ...UserType,
+  ...UserLevel,
+} satisfies UserFeatureGate;

apps/server/src/modules/users/index.ts

@@ -2,11 +2,13 @@ import { Module } from '@nestjs/common';
 
 import { StorageModule } from '../storage';
 import { UserResolver } from './resolver';
+import { UsersService } from './users';
 
 @Module({
   imports: [StorageModule],
-  providers: [UserResolver],
+  providers: [UserResolver, UsersService],
 })
 export class UsersModule {}
 
 export { UserType } from './resolver';
+export { UsersService } from './users';

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

@@ -18,13 +18,13 @@ import type { User } from '@prisma/client';
 // @ts-expect-error graphql-upload is not typed
 import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
 
-import { Config } from '../../config';
 import { PrismaService } from '../../prisma/service';
 import { CloudThrottlerGuard, Throttle } from '../../throttler';
 import type { FileUpload } from '../../types';
 import { Auth, CurrentUser, Public } from '../auth/guard';
 import { StorageService } from '../storage/storage.service';
 import { NewFeaturesKind } from './types';
+import { UsersService } from './users';
 import { isStaff } from './utils';
 
 registerEnumType(NewFeaturesKind, {
@@ -83,7 +83,7 @@ export class UserResolver {
   constructor(
     private readonly prisma: PrismaService,
     private readonly storage: StorageService,
-    private readonly config: Config
+    private readonly users: UsersService
   ) {}
 
   @Throttle(10, 60)
@@ -92,9 +92,7 @@ export class UserResolver {
     description: 'Get current user',
   })
   async currentUser(@CurrentUser() user: UserType) {
-    const storedUser = await this.prisma.user.findUnique({
+    const storedUser = await this.users.findUserById(user.id);
-      where: { id: user.id },
-    });
     if (!storedUser) {
       throw new BadRequestException(`User ${user.id} not found in db`);
     }
@@ -117,27 +115,14 @@ export class UserResolver {
   })
   @Public()
   async user(@Args('email') email: string) {
-    if (this.config.featureFlags.earlyAccessPreview && !isStaff(email)) {
+    if (!(await this.users.canEarlyAccess(email))) {
-      const hasEarlyAccess = await this.prisma.newFeaturesWaitingList
+      return new HttpException(
-        .findUnique({
+        `You don't have early access permission\nVisit https://community.affine.pro/c/insider-general/ for more information`,
-          where: { email, type: NewFeaturesKind.EarlyAccess },
+        401
-        })
+      );
-        .catch(() => false);
-      if (!hasEarlyAccess) {
-        return new HttpException(
-          `You don't have early access permission\nVisit https://community.affine.pro/c/insider-general/ for more information`,
-          401
-        );
-      }
     }
     // TODO: need to limit a user can only get another user witch is in the same workspace
-    const user = await this.prisma.user
+    const user = await this.users.findUserByEmail(email);
-      .findUnique({
-        where: { email },
-      })
-      .catch(() => {
-        return null;
-      });
     if (user?.password) {
       const userResponse: UserType = user;
       userResponse.hasPassword = true;
@@ -155,7 +140,7 @@ export class UserResolver {
     @Args({ name: 'avatar', type: () => GraphQLUpload })
     avatar: FileUpload
   ) {
-    const user = await this.prisma.user.findUnique({ where: { id } });
+    const user = await this.users.findUserById(id);
     if (!user) {
       throw new BadRequestException(`User ${id} not found`);
     }
@@ -169,19 +154,8 @@ export class UserResolver {
   @Throttle(10, 60)
   @Mutation(() => DeleteAccount)
   async deleteAccount(@CurrentUser() user: UserType): Promise<DeleteAccount> {
-    await this.prisma.user.delete({
+    await this.users.deleteUser(user.id);
-      where: {
+    return { success: true };
-        id: user.id,
-      },
-    });
-    await this.prisma.session.deleteMany({
-      where: {
-        userId: user.id,
-      },
-    });
-    return {
-      success: true,
-    };
   }
 
   @Throttle(10, 60)
@@ -194,7 +168,7 @@ export class UserResolver {
     type: NewFeaturesKind,
     @Args('email') email: string
   ): Promise<AddToNewFeaturesWaitingList> {
-    if (!user.email.endsWith('@toeverything.info')) {
+    if (!isStaff(user.email)) {
       throw new ForbiddenException('You are not allowed to do this');
     }
     await this.prisma.newFeaturesWaitingList.create({

apps/server/src/modules/users/users.ts

@@ -0,0 +1,68 @@
+import { Injectable } from '@nestjs/common';
+
+import { Config } from '../../config';
+import { PrismaService } from '../../prisma';
+import { getStorageQuota } from './gates';
+import { NewFeaturesKind } from './types';
+import { isStaff } from './utils';
+
+@Injectable()
+export class UsersService {
+  constructor(
+    private readonly prisma: PrismaService,
+    private readonly config: Config
+  ) {}
+
+  async canEarlyAccess(email: string) {
+    if (this.config.featureFlags.earlyAccessPreview && !isStaff(email)) {
+      return this.prisma.newFeaturesWaitingList
+        .findUnique({
+          where: { email, type: NewFeaturesKind.EarlyAccess },
+        })
+        .catch(() => false);
+    } else {
+      return true;
+    }
+  }
+
+  async getStorageQuotaById(id: string) {
+    const features = await this.prisma.user
+      .findUnique({
+        where: { id },
+        select: {
+          features: {
+            select: {
+              feature: true,
+            },
+          },
+        },
+      })
+      .then(user => user?.features.map(f => f.feature) ?? []);
+
+    return getStorageQuota(features) || this.config.objectStorage.quota;
+  }
+
+  async findUserByEmail(email: string) {
+    return this.prisma.user
+      .findUnique({
+        where: { email },
+      })
+      .catch(() => {
+        return null;
+      });
+  }
+
+  async findUserById(id: string) {
+    return this.prisma.user
+      .findUnique({
+        where: { id },
+      })
+      .catch(() => {
+        return null;
+      });
+  }
+
+  async deleteUser(id: string) {
+    return this.prisma.user.delete({ where: { id } });
+  }
+}

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

@@ -1,6 +1,7 @@
 import { Module } from '@nestjs/common';
 
 import { DocModule } from '../doc';
+import { UsersService } from '../users';
 import { WorkspacesController } from './controller';
 import { PermissionService } from './permission';
 import { WorkspaceResolver } from './resolver';
@@ -8,7 +9,7 @@ import { WorkspaceResolver } from './resolver';
 @Module({
   imports: [DocModule.forFeature()],
   controllers: [WorkspacesController],
-  providers: [WorkspaceResolver, PermissionService],
+  providers: [WorkspaceResolver, PermissionService, UsersService],
   exports: [PermissionService],
 })
 export class WorkspaceModule {}

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

@@ -21,6 +21,30 @@ export class PermissionService {
     return data?.type as Permission;
   }
 
+  async getWorkspaceOwner(workspaceId: string) {
+    return this.prisma.userWorkspacePermission.findFirstOrThrow({
+      where: {
+        workspaceId,
+        type: Permission.Owner,
+      },
+      include: {
+        user: true,
+      },
+    });
+  }
+
+  async tryGetWorkspaceOwner(workspaceId: string) {
+    return this.prisma.userWorkspacePermission.findFirst({
+      where: {
+        workspaceId,
+        type: Permission.Owner,
+      },
+      include: {
+        user: true,
+      },
+    });
+  }
+
   async isAccessible(ws: string, id: string, user?: string): Promise<boolean> {
     if (user) {
       return await this.tryCheck(ws, user);
@@ -157,6 +181,18 @@ export class PermissionService {
       .then(p => p.id);
   }
 
+  async getInvitationById(inviteId: string, workspaceId: string) {
+    return this.prisma.userWorkspacePermission.findUniqueOrThrow({
+      where: {
+        id: inviteId,
+        workspaceId,
+      },
+      include: {
+        user: true,
+      },
+    });
+  }
+
   async acceptById(ws: string, id: string) {
     const result = await this.prisma.userWorkspacePermission.updateMany({
       where: {

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({
+    const data = await this.permissions.getWorkspaceOwner(workspace.id);
-      where: {
-        workspaceId: workspace.id,
-        type: Permission.Owner,
-      },
-      include: {
-        user: true,
-      },
-    });
 
     return data.user;
   }
@@ -241,15 +236,7 @@ export class WorkspaceResolver {
     @CurrentUser() user: UserType,
     @Args('workspaceId') workspaceId: string
   ) {
-    const data = await this.prisma.userWorkspacePermission.findFirst({
+    const data = await this.permissions.tryGetWorkspaceOwner(workspaceId);
-      where: {
-        workspaceId,
-        type: Permission.Owner,
-      },
-      include: {
-        user: true,
-      },
-    });
 
     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({
+    const target = await this.users.findUserByEmail(email);
-      where: {
-        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 =
+    const workspaceId = await this.prisma.userWorkspacePermission
-      await this.prisma.userWorkspacePermission.findUniqueOrThrow({
+      .findUniqueOrThrow({
         where: {
           id: inviteId,
         },
-      });
+        select: {
+          workspaceId: true,
+        },
+      })
+      .then(({ workspaceId }) => workspaceId);
 
     const snapshot = await this.prisma.snapshot.findFirstOrThrow({
       where: {
-        id: permission.workspaceId,
+        id: workspaceId,
-        workspaceId: permission.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({
+    const owner = await this.permissions.getWorkspaceOwner(workspaceId);
-      where: {
+    const invitee = await this.permissions.getInvitationById(
-        workspaceId: permission.workspaceId,
+      inviteId,
-        type: Permission.Owner,
+      workspaceId
-      },
-      include: {
-        user: true,
-      },
-    });
-    const invitee = await this.prisma.userWorkspacePermission.findUniqueOrThrow(
-      {
-        where: {
-          id: inviteId,
-          workspaceId: permission.workspaceId,
-        },
-        include: {
-          user: true,
-        },
-      }
     );
 
     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({
+    const owner = await this.permissions.getWorkspaceOwner(workspaceId);
-      where: {
-        workspaceId,
-        type: Permission.Owner,
-      },
-      include: {
-        user: true,
-      },
-    });
 
     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);
+    await this.permissions.check(workspaceId, user.id, Permission.Write);
-    const quota = this.config.objectStorage.quota;
-    const { size } = await this.collectAllBlobSizes(user);
 
-    if (size > quota) {
+    // quota was apply to owner's account
-      this.logger.log(`storage size limit exceeded: ${size} > ${quota}`);
+    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);
   }

apps/server/src/schema.gql

@@ -158,6 +158,7 @@ type Query {
   listBlobs(workspaceId: String!): [String!]!
   collectBlobSizes(workspaceId: String!): WorkspaceBlobSizes!
   collectAllBlobSizes: WorkspaceBlobSizes!
+  checkBlobSize(workspaceId: String!, size: Float!): WorkspaceBlobSizes!
 
   """Get current user"""
   currentUser: UserType!

apps/server/src/tests/utils.ts

@@ -386,6 +386,27 @@ async function collectAllBlobSizes(
   return res.body.data.collectAllBlobSizes.size;
 }
 
+async function checkBlobSize(
+  app: INestApplication,
+  token: string,
+  workspaceId: string,
+  size: number
+): Promise<number> {
+  const res = await request(app.getHttpServer())
+    .post(gql)
+    .auth(token, { type: 'bearer' })
+    .send({
+      query: `query checkBlobSize($workspaceId: String!, $size: Float!) {
+        checkBlobSize(workspaceId: $workspaceId, size: $size) {
+          size
+        }
+      }`,
+      variables: { workspaceId, size },
+    })
+    .expect(200);
+  return res.body.data.checkBlobSize.size;
+}
+
 async function setBlob(
   app: INestApplication,
   token: string,
@@ -480,6 +501,7 @@ async function getInviteInfo(
 export {
   acceptInvite,
   acceptInviteById,
+  checkBlobSize,
   collectAllBlobSizes,
   collectBlobSizes,
   createTestApp,

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

@@ -8,6 +8,7 @@ import request from 'supertest';
 
 import { AppModule } from '../app';
 import {
+  checkBlobSize,
   collectAllBlobSizes,
   collectBlobSizes,
   createWorkspace,
@@ -126,4 +127,20 @@ test('should calc all blobs size', async t => {
 
   const size = await collectAllBlobSizes(app, u1.token.token);
   t.is(size, 8, 'failed to collect all blob sizes');
+
+  const size1 = await checkBlobSize(
+    app,
+    u1.token.token,
+    workspace1.id,
+    10 * 1024 * 1024 * 1024 - 8
+  );
+  t.is(size1, 0, 'failed to check blob size');
+
+  const size2 = await checkBlobSize(
+    app,
+    u1.token.token,
+    workspace1.id,
+    10 * 1024 * 1024 * 1024 - 7
+  );
+  t.is(size2, -1, 'failed to check blob size');
 });

apps/server/src/throttler.ts

@@ -54,4 +54,22 @@ export class CloudThrottlerGuard extends ThrottlerGuard {
   }
 }
 
+@Injectable()
+export class AuthThrottlerGuard extends CloudThrottlerGuard {
+  override async handleRequest(
+    context: ExecutionContext,
+    limit: number,
+    ttl: number
+  ): Promise<boolean> {
+    const { req } = this.getRequestResponse(context);
+
+    if (req?.url === '/api/auth/session') {
+      // relax throttle for session auto renew
+      return super.handleRequest(context, limit * 20, ttl);
+    }
+
+    return super.handleRequest(context, limit, ttl);
+  }
+}
+
 export { Throttle };

packages/graphql/src/fetcher.ts

@@ -215,7 +215,7 @@ export const gqlFetcherFactory = (endpoint: string) => {
   return gqlFetch;
 };
 
-export const fetchWithTraceReport = (
+export const fetchWithTraceReport = async (
   input: RequestInfo | URL,
   init?: RequestInit,
   traceOptions?: { event: string }
@@ -241,21 +241,20 @@ export const fetchWithTraceReport = (
     return fetch(input, init);
   }
 
-  return fetch(input, init)
+  try {
-    .then(response => {
+    const response = await fetch(input, init);
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      traceReporter!.cacheTrace(traceId, spanId, startTime, {
+    traceReporter!.cacheTrace(traceId, spanId, startTime, {
-        requestId,
+      requestId,
-        ...(event ? { event } : {}),
+      ...(event ? { event } : {}),
-      });
-      return response;
-    })
-    .catch(err => {
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      traceReporter!.uploadTrace(traceId, spanId, startTime, {
-        requestId,
-        ...(event ? { event } : {}),
-      });
-      return Promise.reject(err);
     });
+    return response;
+  } catch (err) {
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    traceReporter!.uploadTrace(traceId, spanId, startTime, {
+      requestId,
+      ...(event ? { event } : {}),
+    });
+    return await Promise.reject(err);
+  }
 };

packages/graphql/src/graphql/blob-check-size.gql

@@ -0,0 +1,5 @@
+query checkBlobSizes($workspaceId: String!, $size: Float!) {
+  checkBlobSize(workspaceId: $workspaceId, size: $size) {
+    size
+  }
+}

packages/graphql/src/graphql/index.ts

@@ -7,6 +7,19 @@ export interface GraphQLQuery {
   containsFile?: boolean;
 }
 
+export const checkBlobSizesQuery = {
+  id: 'checkBlobSizesQuery' as const,
+  operationName: 'checkBlobSizes',
+  definitionName: 'checkBlobSize',
+  containsFile: false,
+  query: `
+query checkBlobSizes($workspaceId: String!, $size: Float!) {
+  checkBlobSize(workspaceId: $workspaceId, size: $size) {
+    size
+  }
+}`,
+};
+
 export const deleteBlobMutation = {
   id: 'deleteBlobMutation' as const,
   operationName: 'deleteBlob',

packages/graphql/src/schema.ts

@@ -50,6 +50,16 @@ export interface UpdateWorkspaceInput {
   public: InputMaybe<Scalars['Boolean']['input']>;
 }
 
+export type CheckBlobSizesQueryVariables = Exact<{
+  workspaceId: Scalars['String']['input'];
+  size: Scalars['Float']['input'];
+}>;
+
+export type CheckBlobSizesQuery = {
+  __typename?: 'Query';
+  checkBlobSize: { __typename?: 'WorkspaceBlobSizes'; size: number };
+};
+
 export type DeleteBlobMutationVariables = Exact<{
   workspaceId: Scalars['String']['input'];
   hash: Scalars['String']['input'];
@@ -448,6 +458,11 @@ export type AcceptInviteByWorkspaceIdMutation = {
 };
 
 export type Queries =
+  | {
+      name: 'checkBlobSizesQuery';
+      variables: CheckBlobSizesQueryVariables;
+      response: CheckBlobSizesQuery;
+    }
   | {
       name: 'listBlobsQuery';
       variables: ListBlobsQueryVariables;

packages/workspace/src/blob/cloud-blob-storage.ts

@@ -1,13 +1,13 @@
 import {
+  checkBlobSizesQuery,
   deleteBlobMutation,
   fetchWithTraceReport,
   listBlobsQuery,
   setBlobMutation,
 } from '@affine/graphql';
+import { fetcher } from '@affine/workspace/affine/gql';
 import type { BlobStorage } from '@blocksuite/store';
 
-import { fetcher } from '../affine/gql';
-
 export const createCloudBlobStorage = (workspaceId: string): BlobStorage => {
   return {
     crud: {
@@ -18,6 +18,20 @@ export const createCloudBlobStorage = (workspaceId: string): BlobStorage => {
         ).then(res => res.blob());
       },
       set: async (key, value) => {
+        const {
+          checkBlobSize: { size },
+        } = await fetcher({
+          query: checkBlobSizesQuery,
+          variables: {
+            workspaceId,
+            size: value.size,
+          },
+        });
+
+        if (size <= 0) {
+          throw new Error('Blob size limit exceeded');
+        }
+
         const result = await fetcher({
           query: setBlobMutation,
           variables: {