diff --git a/packages/backend/server/src/core/comment/index.ts b/packages/backend/server/src/core/comment/index.ts index 75957eb060..51bcd8ebc2 100644 --- a/packages/backend/server/src/core/comment/index.ts +++ b/packages/backend/server/src/core/comment/index.ts @@ -2,6 +2,7 @@ import { Module } from '@nestjs/common'; import { ServerConfigModule } from '../config'; import { PermissionModule } from '../permission'; +import { QuotaServiceModule } from '../quota'; import { StorageModule } from '../storage'; import { CommentRealtimeModule } from './realtime.module'; import { CommentResolver } from './resolver'; @@ -9,6 +10,7 @@ import { CommentResolver } from './resolver'; @Module({ imports: [ PermissionModule, + QuotaServiceModule, StorageModule, ServerConfigModule, CommentRealtimeModule, diff --git a/packages/backend/server/src/core/comment/resolver.ts b/packages/backend/server/src/core/comment/resolver.ts index 7f33c0a8cd..5dcf950b69 100644 --- a/packages/backend/server/src/core/comment/resolver.ts +++ b/packages/backend/server/src/core/comment/resolver.ts @@ -26,6 +26,7 @@ import { Comment, DocMode, Models, Reply } from '../../models'; import { CurrentUser } from '../auth/session'; import { ServerFeature, ServerService } from '../config'; import { DocAction, PermissionAccess } from '../permission'; +import { QuotaService } from '../quota'; import { RealtimePublisher } from '../realtime'; import { CommentAttachmentStorage } from '../storage'; import { UserType } from '../user'; @@ -56,6 +57,7 @@ export class CommentResolver { private readonly service: CommentService, private readonly ac: PermissionAccess, private readonly commentAttachmentStorage: CommentAttachmentStorage, + private readonly quota: QuotaService, private readonly queue: JobQueue, private readonly models: Models, private readonly server: ServerService, @@ -354,13 +356,19 @@ export class CommentResolver { 'Doc.Comments.Create' ); - // TODO(@fengmk2): should check total attachment quota in the future version const buffer = await readableToBuffer(attachment.createReadStream()); // max attachment size is 10MB if (buffer.length > 10 * 1024 * 1024) { throw new CommentAttachmentQuotaExceeded(); } + const checkExceeded = + await this.quota.getWorkspaceQuotaCalculator(workspaceId); + const result = checkExceeded(buffer.length); + if (result?.blobQuotaExceeded || result?.storageQuotaExceeded) { + throw new CommentAttachmentQuotaExceeded(); + } + const key = randomUUID(); await this.commentAttachmentStorage.put( workspaceId, diff --git a/packages/backend/server/src/core/quota/state.ts b/packages/backend/server/src/core/quota/state.ts index 71ff12284a..e02e550cd9 100644 --- a/packages/backend/server/src/core/quota/state.ts +++ b/packages/backend/server/src/core/quota/state.ts @@ -287,17 +287,30 @@ export class QuotaStateService { } private async getWorkspaceStorageUsage(workspaceId: string) { - const sum = await this.db.blob.aggregate({ - where: { - workspaceId, - deletedAt: null, - }, - _sum: { - size: true, - }, - }); + const [blobSum, commentAttachmentSum] = await Promise.all([ + this.db.blob.aggregate({ + where: { + workspaceId, + deletedAt: null, + }, + _sum: { + size: true, + }, + }), + this.db.commentAttachment.aggregate({ + where: { + workspaceId, + }, + _sum: { + size: true, + }, + }), + ]); - return BigInt(sum._sum.size ?? 0); + return ( + BigInt(blobSum._sum.size ?? 0) + + BigInt(commentAttachmentSum._sum.size ?? 0) + ); } private hasStandaloneWorkspaceQuota(plan: string) {