From f4340da4786058eedd242277833331047e1dc86a Mon Sep 17 00:00:00 2001 From: DarkSky <25152247+darkskygit@users.noreply.github.com> Date: Fri, 8 Sep 2023 05:32:41 +0800 Subject: [PATCH] refactor(server): rate limit and permission (#4198) Co-authored-by: LongYinan --- .../migration.sql | 13 ++ apps/server/schema.prisma | 16 +- apps/server/src/config/def.ts | 4 +- apps/server/src/config/default.ts | 3 +- .../src/modules/auth/next-auth-options.ts | 120 ++--------- .../src/modules/auth/next-auth.controller.ts | 10 +- apps/server/src/modules/auth/utils/index.ts | 3 + apps/server/src/modules/auth/utils/jwt.ts | 76 +++++++ .../src/modules/auth/utils/send-mail.ts | 41 ++++ apps/server/src/modules/users/gates.ts | 42 ++++ apps/server/src/modules/users/index.ts | 4 +- apps/server/src/modules/users/resolver.ts | 52 ++--- apps/server/src/modules/users/users.ts | 68 ++++++ apps/server/src/modules/workspaces/index.ts | 3 +- .../src/modules/workspaces/permission.ts | 36 ++++ .../server/src/modules/workspaces/resolver.ts | 194 +++++++++--------- apps/server/src/schema.gql | 1 + apps/server/src/tests/utils.ts | 22 ++ apps/server/src/tests/workspace-blobs.spec.ts | 17 ++ apps/server/src/throttler.ts | 18 ++ packages/graphql/src/fetcher.ts | 33 ++- .../graphql/src/graphql/blob-check-size.gql | 5 + packages/graphql/src/graphql/index.ts | 13 ++ packages/graphql/src/schema.ts | 15 ++ .../workspace/src/blob/cloud-blob-storage.ts | 18 +- 25 files changed, 555 insertions(+), 272 deletions(-) create mode 100644 apps/server/migrations/20230906100042_user_feature_gates/migration.sql create mode 100644 apps/server/src/modules/auth/utils/index.ts create mode 100644 apps/server/src/modules/auth/utils/jwt.ts create mode 100644 apps/server/src/modules/auth/utils/send-mail.ts create mode 100644 apps/server/src/modules/users/gates.ts create mode 100644 apps/server/src/modules/users/users.ts create mode 100644 packages/graphql/src/graphql/blob-check-size.gql diff --git a/apps/server/migrations/20230906100042_user_feature_gates/migration.sql b/apps/server/migrations/20230906100042_user_feature_gates/migration.sql new file mode 100644 index 0000000000..abb4da2b7f --- /dev/null +++ b/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; diff --git a/apps/server/schema.prisma b/apps/server/schema.prisma index e5ad0e5073..93b481b081 100644 --- a/apps/server/schema.prisma +++ b/apps/server/schema.prisma @@ -1,6 +1,6 @@ generator client { - provider = "prisma-client-js" - binaryTargets = ["native", "debian-openssl-3.0.x"] + provider = "prisma-client-js" + 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") diff --git a/apps/server/src/config/def.ts b/apps/server/src/config/def.ts index ef434575ec..efe6b0b297 100644 --- a/apps/server/src/config/def.ts +++ b/apps/server/src/config/def.ts @@ -187,8 +187,8 @@ export interface AFFiNEConfig { path: string; }; /** - * Free user storage quota - * @default 10 * 1024 * 1024 (10GB) + * default storage quota + * @default 10 * 1024 * 1024 * 1024 (10GB) */ quota: number; }; diff --git a/apps/server/src/config/default.ts b/apps/server/src/config/default.ts index 8f9b9eb52d..eeea425d32 100644 --- a/apps/server/src/config/default.ts +++ b/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, diff --git a/apps/server/src/modules/auth/next-auth-options.ts b/apps/server/src/modules/auth/next-auth-options.ts index 0e060d66d7..e0d40b4e76 100644 --- a/apps/server/src/modules/auth/next-auth-options.ts +++ b/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, { - type SendVerificationRequestParams, -} from 'next-auth/providers/email'; +import Email 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 = { }, }, from: config.auth.email.sender, - async sendVerificationRequest(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`); - } - }, + sendVerificationRequest: (params: SendVerificationRequestParams) => + sendVerificationRequest(config, logger, mailer, session, params), }), ], // @ts-expect-error Third part library type mismatch @@ -200,66 +163,9 @@ export const NextAuthOptionsProvider: FactoryProvider = { } nextAuthOptions.jwt = { - encode: async ({ token, maxAge }) => { - 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, - } - ); - }, - 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 & { - picture: string | undefined; - }; - return { - name, - email, - emailVerified, - picture, - sub: id, - id, - hasPassword, - }; - }, + encode: async ({ token, maxAge }) => + encode(config, prisma, token, maxAge), + decode: async ({ token }) => decode(config, token), }; nextAuthOptions.secret ??= config.auth.nextAuthSecret; diff --git a/apps/server/src/modules/auth/next-auth.controller.ts b/apps/server/src/modules/auth/next-auth.controller.ts index 4092a1d7d7..a0eae5c7b6 100644 --- a/apps/server/src/modules/auth/next-auth.controller.ts +++ b/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 { + private async verifyUserFromRequest(req: Request): Promise { const token = req.headers.authorization; if (!token) { const session = await AuthHandler({ diff --git a/apps/server/src/modules/auth/utils/index.ts b/apps/server/src/modules/auth/utils/index.ts new file mode 100644 index 0000000000..c00310ed8f --- /dev/null +++ b/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'; diff --git a/apps/server/src/modules/auth/utils/jwt.ts b/apps/server/src/modules/auth/utils/jwt.ts new file mode 100644 index 0000000000..b666636011 --- /dev/null +++ b/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 & { + picture: string | undefined; + }; + return { + name, + email, + emailVerified, + picture, + sub: id, + id, + hasPassword, + }; +}; diff --git a/apps/server/src/modules/auth/utils/send-mail.ts b/apps/server/src/modules/auth/utils/send-mail.ts new file mode 100644 index 0000000000..c2cf89b578 --- /dev/null +++ b/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`); + } +} diff --git a/apps/server/src/modules/users/gates.ts b/apps/server/src/modules/users/gates.ts new file mode 100644 index 0000000000..191e00752e --- /dev/null +++ b/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; + +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; + +export const FeatureGates = { + ...UserType, + ...UserLevel, +} satisfies UserFeatureGate; diff --git a/apps/server/src/modules/users/index.ts b/apps/server/src/modules/users/index.ts index c885a06534..40dc84ae46 100644 --- a/apps/server/src/modules/users/index.ts +++ b/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'; diff --git a/apps/server/src/modules/users/resolver.ts b/apps/server/src/modules/users/resolver.ts index f5bcdf74e3..dd673a12d8 100644 --- a/apps/server/src/modules/users/resolver.ts +++ b/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({ - where: { id: user.id }, - }); + const storedUser = await this.users.findUserById(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)) { - const hasEarlyAccess = await this.prisma.newFeaturesWaitingList - .findUnique({ - where: { email, type: NewFeaturesKind.EarlyAccess }, - }) - .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 - ); - } + if (!(await this.users.canEarlyAccess(email))) { + 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 - .findUnique({ - where: { email }, - }) - .catch(() => { - return null; - }); + const user = await this.users.findUserByEmail(email); 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 { - await this.prisma.user.delete({ - where: { - id: user.id, - }, - }); - await this.prisma.session.deleteMany({ - where: { - userId: user.id, - }, - }); - return { - success: true, - }; + await this.users.deleteUser(user.id); + return { success: true }; } @Throttle(10, 60) @@ -194,7 +168,7 @@ export class UserResolver { type: NewFeaturesKind, @Args('email') email: string ): Promise { - 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({ diff --git a/apps/server/src/modules/users/users.ts b/apps/server/src/modules/users/users.ts new file mode 100644 index 0000000000..477adc1fd3 --- /dev/null +++ b/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 } }); + } +} diff --git a/apps/server/src/modules/workspaces/index.ts b/apps/server/src/modules/workspaces/index.ts index 6ea516e1dc..c1ee3acc00 100644 --- a/apps/server/src/modules/workspaces/index.ts +++ b/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 {} diff --git a/apps/server/src/modules/workspaces/permission.ts b/apps/server/src/modules/workspaces/permission.ts index c35e65815a..2376d71c0b 100644 --- a/apps/server/src/modules/workspaces/permission.ts +++ b/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 { 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: { diff --git a/apps/server/src/modules/workspaces/resolver.ts b/apps/server/src/modules/workspaces/resolver.ts index 7fc5115ad9..df63cf185b 100644 --- a/apps/server/src/modules/workspaces/resolver.ts +++ b/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((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); } diff --git a/apps/server/src/schema.gql b/apps/server/src/schema.gql index e511c7c8c4..3deac9fa04 100644 --- a/apps/server/src/schema.gql +++ b/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! diff --git a/apps/server/src/tests/utils.ts b/apps/server/src/tests/utils.ts index bd43f4da2c..debd99b45c 100644 --- a/apps/server/src/tests/utils.ts +++ b/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 { + 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, diff --git a/apps/server/src/tests/workspace-blobs.spec.ts b/apps/server/src/tests/workspace-blobs.spec.ts index dbcb436922..41701d6974 100644 --- a/apps/server/src/tests/workspace-blobs.spec.ts +++ b/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'); }); diff --git a/apps/server/src/throttler.ts b/apps/server/src/throttler.ts index 60b125342a..3ebcfc97e6 100644 --- a/apps/server/src/throttler.ts +++ b/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 { + 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 }; diff --git a/packages/graphql/src/fetcher.ts b/packages/graphql/src/fetcher.ts index d40d497cd0..36f61b97e7 100644 --- a/packages/graphql/src/fetcher.ts +++ b/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) - .then(response => { - // eslint-disable-next-line @typescript-eslint/no-non-null-assertion - traceReporter!.cacheTrace(traceId, spanId, startTime, { - requestId, - ...(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); + try { + const response = await fetch(input, init); + // eslint-disable-next-line @typescript-eslint/no-non-null-assertion + traceReporter!.cacheTrace(traceId, spanId, startTime, { + requestId, + ...(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 await Promise.reject(err); + } }; diff --git a/packages/graphql/src/graphql/blob-check-size.gql b/packages/graphql/src/graphql/blob-check-size.gql new file mode 100644 index 0000000000..be76a0a1ab --- /dev/null +++ b/packages/graphql/src/graphql/blob-check-size.gql @@ -0,0 +1,5 @@ +query checkBlobSizes($workspaceId: String!, $size: Float!) { + checkBlobSize(workspaceId: $workspaceId, size: $size) { + size + } +} diff --git a/packages/graphql/src/graphql/index.ts b/packages/graphql/src/graphql/index.ts index 8afc4bc527..fce5cea8dc 100644 --- a/packages/graphql/src/graphql/index.ts +++ b/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', diff --git a/packages/graphql/src/schema.ts b/packages/graphql/src/schema.ts index ba066d815e..0dbb1bc67a 100644 --- a/packages/graphql/src/schema.ts +++ b/packages/graphql/src/schema.ts @@ -50,6 +50,16 @@ export interface UpdateWorkspaceInput { public: InputMaybe; } +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; diff --git a/packages/workspace/src/blob/cloud-blob-storage.ts b/packages/workspace/src/blob/cloud-blob-storage.ts index 7e0fb35db9..7761229ea0 100644 --- a/packages/workspace/src/blob/cloud-blob-storage.ts +++ b/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: {