refactor(server): use new storage providers (#5433)

packages/backend/server/src/modules/index.ts

@@ -8,6 +8,7 @@ import { DocModule } from './doc';
 import { PaymentModule } from './payment';
 import { QuotaModule } from './quota';
 import { SelfHostedModule } from './self-hosted';
+import { StorageModule } from './storage';
 import { SyncModule } from './sync';
 import { UsersModule } from './users';
 import { WorkspaceModule } from './workspaces';
@@ -27,7 +28,8 @@ switch (SERVER_FLAVOR) {
       WorkspaceModule,
       UsersModule,
       SyncModule,
-      DocModule
+      DocModule,
+      StorageModule
     );
     break;
   case 'graphql':
@@ -39,7 +41,8 @@ switch (SERVER_FLAVOR) {
       UsersModule,
       DocModule,
       PaymentModule,
-      QuotaModule
+      QuotaModule,
+      StorageModule
     );
     break;
   case 'allinone':
@@ -53,7 +56,8 @@ switch (SERVER_FLAVOR) {
       QuotaModule,
       SyncModule,
       DocModule,
-      PaymentModule
+      PaymentModule,
+      StorageModule
     );
     break;
 }

packages/backend/server/src/modules/quota/index.ts

@@ -1,5 +1,6 @@
 import { Module } from '@nestjs/common';
 
+import { StorageModule } from '../storage';
 import { PermissionService } from '../workspaces/permission';
 import { QuotaService } from './service';
 import { QuotaManagementService } from './storage';
@@ -11,6 +12,7 @@ import { QuotaManagementService } from './storage';
  * - quota statistics
  */
 @Module({
+  imports: [StorageModule],
   providers: [PermissionService, QuotaService, QuotaManagementService],
   exports: [QuotaService, QuotaManagementService],
 })

packages/backend/server/src/modules/quota/storage.ts

@@ -1,7 +1,6 @@
-import type { Storage } from '@affine/storage';
-import { Inject, Injectable, NotFoundException } from '@nestjs/common';
+import { Injectable, NotFoundException } from '@nestjs/common';
 
-import { StorageProvide } from '../../storage';
+import { WorkspaceBlobStorage } from '../storage';
 import { PermissionService } from '../workspaces/permission';
 import { QuotaService } from './service';
 
@@ -10,7 +9,7 @@ export class QuotaManagementService {
   constructor(
     private readonly quota: QuotaService,
     private readonly permissions: PermissionService,
-    @Inject(StorageProvide) private readonly storage: Storage
+    private readonly storage: WorkspaceBlobStorage
   ) {}
 
   async getUserQuota(userId: string) {
@@ -29,7 +28,12 @@ export class QuotaManagementService {
   // TODO: lazy calc, need to be optimized with cache
   async getUserUsage(userId: string) {
     const workspaces = await this.permissions.getOwnedWorkspaces(userId);
-    return this.storage.blobsSize(workspaces);
+
+    const sizes = await Promise.all(
+      workspaces.map(workspace => this.storage.totalSize(workspace))
+    );
+
+    return sizes.reduce((total, size) => total + size, 0);
   }
 
   // get workspace's owner quota and total size of used

packages/backend/server/src/modules/storage/fs.ts

@@ -1,30 +0,0 @@
-import { randomUUID } from 'node:crypto';
-import { createWriteStream } from 'node:fs';
-import { mkdir } from 'node:fs/promises';
-import { join } from 'node:path';
-import { pipeline } from 'node:stream/promises';
-
-import { Injectable } from '@nestjs/common';
-
-import { Config } from '../../config';
-import { FileUpload } from '../../types';
-
-@Injectable()
-export class FSService {
-  constructor(private readonly config: Config) {}
-
-  async writeFile(key: string, file: FileUpload) {
-    const dest = this.config.objectStorage.fs.path;
-    const fileName = `${key}-${randomUUID()}`;
-    const prefix = this.config.node.dev
-      ? `${this.config.https ? 'https' : 'http'}://${this.config.host}:${
-          this.config.port
-        }`
-      : '';
-    await mkdir(dest, { recursive: true });
-    const destFile = join(dest, fileName);
-    await pipeline(file.createReadStream(), createWriteStream(destFile));
-
-    return `${prefix}/assets/${fileName}`;
-  }
-}

packages/backend/server/src/modules/storage/index.ts

@@ -1,11 +1,11 @@
 import { Module } from '@nestjs/common';
 
-import { FSService } from './fs';
-import { S3 } from './s3';
-import { StorageService } from './storage.service';
+import { AvatarStorage, WorkspaceBlobStorage } from './wrappers';
 
 @Module({
-  providers: [S3, StorageService, FSService],
-  exports: [StorageService],
+  providers: [WorkspaceBlobStorage, AvatarStorage],
+  exports: [WorkspaceBlobStorage, AvatarStorage],
 })
 export class StorageModule {}
+
+export { AvatarStorage, WorkspaceBlobStorage };

packages/backend/server/src/modules/storage/providers/fs.ts

@@ -34,6 +34,8 @@ export class FsStorageProvider implements StorageProvider {
   private readonly path: string;
   private readonly logger: Logger;
 
+  readonly type = 'fs';
+
   constructor(
     config: FsStorageConfig,
     public readonly bucket: string

packages/backend/server/src/modules/storage/providers/provider.ts

@@ -1,5 +1,7 @@
 import type { Readable } from 'node:stream';
 
+import { StorageProviderType } from '../../../config';
+
 export interface GetObjectMetadata {
   /**
    * @default 'application/octet-stream'
@@ -26,6 +28,7 @@ export type BlobInputType = Buffer | Readable | string;
 export type BlobOutputType = Readable;
 
 export interface StorageProvider {
+  readonly type: StorageProviderType;
   put(
     key: string,
     body: BlobInputType,

packages/backend/server/src/modules/storage/providers/r2.ts

@@ -2,6 +2,8 @@ import { R2StorageConfig } from '../../../config/storage';
 import { S3StorageProvider } from './s3';
 
 export class R2StorageProvider extends S3StorageProvider {
+  override readonly type = 'r2' as any /* cast 'r2' to 's3' */;
+
   constructor(config: R2StorageConfig, bucket: string) {
     super(
       {

packages/backend/server/src/modules/storage/providers/s3.ts

@@ -21,8 +21,11 @@ import {
 import { autoMetadata, toBuffer } from './utils';
 
 export class S3StorageProvider implements StorageProvider {
-  logger: Logger;
-  client: S3Client;
+  private readonly logger: Logger;
+  protected client: S3Client;
+
+  readonly type = 's3';
+
   constructor(
     config: S3StorageConfig,
     public readonly bucket: string

packages/backend/server/src/modules/storage/s3.ts

@@ -1,22 +0,0 @@
-import { S3Client } from '@aws-sdk/client-s3';
-import { FactoryProvider } from '@nestjs/common';
-
-import { Config } from '../../config';
-
-export const S3_SERVICE = Symbol('S3_SERVICE');
-
-export const S3: FactoryProvider<S3Client> = {
-  provide: S3_SERVICE,
-  useFactory: (config: Config) => {
-    const s3 = new S3Client({
-      region: 'auto',
-      endpoint: `https://${config.objectStorage.r2.accountId}.r2.cloudflarestorage.com`,
-      credentials: {
-        accessKeyId: config.objectStorage.r2.accessKeyId,
-        secretAccessKey: config.objectStorage.r2.secretAccessKey,
-      },
-    });
-    return s3;
-  },
-  inject: [Config],
-};

packages/backend/server/src/modules/storage/storage.service.ts

@@ -1,43 +0,0 @@
-import { PutObjectCommand, S3Client } from '@aws-sdk/client-s3';
-import { Inject, Injectable } from '@nestjs/common';
-import { crc32 } from '@node-rs/crc32';
-import { fileTypeFromBuffer } from 'file-type';
-// eslint-disable-next-line @typescript-eslint/ban-ts-comment
-// @ts-ignore - no types
-import { getStreamAsBuffer } from 'get-stream';
-
-import { Config } from '../../config';
-import { FileUpload } from '../../types';
-import { FSService } from './fs';
-import { S3_SERVICE } from './s3';
-
-@Injectable()
-export class StorageService {
-  constructor(
-    @Inject(S3_SERVICE) private readonly s3: S3Client,
-    private readonly fs: FSService,
-    private readonly config: Config
-  ) {}
-
-  async uploadFile(key: string, file: FileUpload) {
-    if (this.config.objectStorage.r2.enabled) {
-      const readableFile = file.createReadStream();
-      const fileBuffer = await getStreamAsBuffer(readableFile);
-      const mime = (await fileTypeFromBuffer(fileBuffer))?.mime;
-      const crc32Value = crc32(fileBuffer);
-      const keyWithCrc32 = `${crc32Value}-${key}`;
-      await this.s3.send(
-        new PutObjectCommand({
-          Body: fileBuffer,
-          Bucket: this.config.objectStorage.r2.bucket,
-          Key: keyWithCrc32,
-          ContentLength: fileBuffer.length,
-          ContentType: mime,
-        })
-      );
-      return `https://avatar.affineassets.com/${keyWithCrc32}`;
-    } else {
-      return this.fs.writeFile(key, file);
-    }
-  }
-}

packages/backend/server/src/modules/storage/wrappers/avatar.ts

@@ -1,6 +1,7 @@
 import { Injectable } from '@nestjs/common';
 
-import { Config } from '../../../config';
+import { AFFiNEStorageConfig, Config } from '../../../config';
+import { type EventPayload, OnEvent } from '../../../event';
 import {
   BlobInputType,
   createStorageProvider,
@@ -11,20 +12,36 @@ import {
 @Injectable()
 export class AvatarStorage {
   public readonly provider: StorageProvider;
+  private readonly storageConfig: AFFiNEStorageConfig['storages']['avatar'];
 
-  constructor({ storage }: Config) {
-    this.provider = createStorageProvider(storage, 'avatar');
+  constructor(private readonly config: Config) {
+    this.provider = createStorageProvider(this.config.storage, 'avatar');
+    this.storageConfig = this.config.storage.storages.avatar;
   }
 
-  put(key: string, blob: BlobInputType, metadata?: PutObjectMetadata) {
-    return this.provider.put(key, blob, metadata);
+  async put(key: string, blob: BlobInputType, metadata?: PutObjectMetadata) {
+    await this.provider.put(key, blob, metadata);
+    let link = this.storageConfig.publicLinkFactory(key);
+
+    if (link.startsWith('/')) {
+      link = this.config.baseUrl + link;
+    }
+
+    return link;
   }
 
   get(key: string) {
     return this.provider.get(key);
   }
 
-  async delete(key: string) {
+  delete(key: string) {
     return this.provider.delete(key);
   }
+
+  @OnEvent('user.deleted')
+  async onUserDeleted(user: EventPayload<'user.deleted'>) {
+    if (user.avatarUrl) {
+      await this.delete(user.avatarUrl);
+    }
+  }
 }

packages/backend/server/src/modules/storage/wrappers/blob.ts

@@ -1,6 +1,7 @@
 import { Injectable } from '@nestjs/common';
 
 import { Config } from '../../../config';
+import { EventEmitter, type EventPayload, OnEvent } from '../../../event';
 import {
   BlobInputType,
   createStorageProvider,
@@ -10,7 +11,10 @@ import {
 @Injectable()
 export class WorkspaceBlobStorage {
   public readonly provider: StorageProvider;
-  constructor({ storage }: Config) {
+  constructor(
+    private readonly event: EventEmitter,
+    { storage }: Config
+  ) {
     this.provider = createStorageProvider(storage, 'blob');
   }
 
@@ -42,4 +46,25 @@ export class WorkspaceBlobStorage {
     // how could we ignore the ones get soft-deleted?
     return blobs.reduce((acc, item) => acc + item.size, 0);
   }
+
+  @OnEvent('workspace.deleted')
+  async onWorkspaceDeleted(workspaceId: EventPayload<'workspace.deleted'>) {
+    const blobs = await this.list(workspaceId);
+
+    // to reduce cpu time holding
+    blobs.forEach(blob => {
+      this.event.emit('workspace.blob.deleted', {
+        workspaceId: workspaceId,
+        name: blob.key,
+      });
+    });
+  }
+
+  @OnEvent('workspace.blob.deleted')
+  async onDeleteWorkspaceBlob({
+    workspaceId,
+    name,
+  }: EventPayload<'workspace.blob.deleted'>) {
+    await this.delete(workspaceId, name);
+  }
 }

packages/backend/server/src/modules/users/controller.ts

@@ -0,0 +1,40 @@
+import {
+  Controller,
+  ForbiddenException,
+  Get,
+  NotFoundException,
+  Param,
+  Res,
+} from '@nestjs/common';
+import type { Response } from 'express';
+
+import { AvatarStorage } from '../storage';
+
+@Controller('/api/avatars')
+export class UserAvatarController {
+  constructor(private readonly storage: AvatarStorage) {}
+
+  @Get('/:id')
+  async getAvatar(@Res() res: Response, @Param('id') id: string) {
+    if (this.storage.provider.type !== 'fs') {
+      throw new ForbiddenException(
+        'Only available when avatar storage provider set to fs.'
+      );
+    }
+
+    const { body, metadata } = await this.storage.get(id);
+
+    if (!body) {
+      throw new NotFoundException(`Avatar ${id} not found.`);
+    }
+
+    // metadata should always exists if body is not null
+    if (metadata) {
+      res.setHeader('content-type', metadata.contentType);
+      res.setHeader('last-modified', metadata.lastModified.toISOString());
+      res.setHeader('content-length', metadata.contentLength);
+    }
+
+    body.pipe(res);
+  }
+}

packages/backend/server/src/modules/users/index.ts

@@ -3,12 +3,14 @@ import { Module } from '@nestjs/common';
 import { FeatureModule } from '../features';
 import { QuotaModule } from '../quota';
 import { StorageModule } from '../storage';
+import { UserAvatarController } from './controller';
 import { UserResolver } from './resolver';
 import { UsersService } from './users';
 
 @Module({
   imports: [StorageModule, FeatureModule, QuotaModule],
   providers: [UserResolver, UsersService],
+  controllers: [UserAvatarController],
   exports: [UsersService],
 })
 export class UsersModule {}

packages/backend/server/src/modules/users/resolver.ts

@@ -17,6 +17,7 @@ import type { User } from '@prisma/client';
 import { GraphQLError } from 'graphql';
 import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
 
+import { EventEmitter } from '../../event';
 import { PrismaService } from '../../prisma/service';
 import { CloudThrottlerGuard, Throttle } from '../../throttler';
 import type { FileUpload } from '../../types';
@@ -24,7 +25,7 @@ import { Auth, CurrentUser, Public, Publicable } from '../auth/guard';
 import { AuthService } from '../auth/service';
 import { FeatureManagementService } from '../features';
 import { QuotaService } from '../quota';
-import { StorageService } from '../storage/storage.service';
+import { AvatarStorage } from '../storage';
 import { DeleteAccount, RemoveAvatar, UserQuotaType, UserType } from './types';
 import { UsersService } from './users';
 
@@ -39,10 +40,11 @@ export class UserResolver {
   constructor(
     private readonly auth: AuthService,
     private readonly prisma: PrismaService,
-    private readonly storage: StorageService,
+    private readonly storage: AvatarStorage,
     private readonly users: UsersService,
     private readonly feature: FeatureManagementService,
-    private readonly quota: QuotaService
+    private readonly quota: QuotaService,
+    private readonly event: EventEmitter
   ) {}
 
   @Throttle({
@@ -147,10 +149,20 @@ export class UserResolver {
     if (!user) {
       throw new BadRequestException(`User not found`);
     }
-    const url = await this.storage.uploadFile(`${user.id}-avatar`, avatar);
+
+    const link = await this.storage.put(
+      `${user.id}-avatar`,
+      avatar.createReadStream(),
+      {
+        contentType: avatar.mimetype,
+      }
+    );
+
     return this.prisma.user.update({
       where: { id: user.id },
-      data: { avatarUrl: url },
+      data: {
+        avatarUrl: link,
+      },
     });
   }
 
@@ -183,7 +195,8 @@ export class UserResolver {
   })
   @Mutation(() => DeleteAccount)
   async deleteAccount(@CurrentUser() user: UserType): Promise<DeleteAccount> {
-    await this.users.deleteUser(user.id);
+    const deletedUser = await this.users.deleteUser(user.id);
+    this.event.emit('user.deleted', deletedUser);
     return { success: true };
   }
 

packages/backend/server/src/modules/workspaces/controller.ts

@@ -1,9 +1,8 @@
-import type { Storage } from '@affine/storage';
 import {
   Controller,
   ForbiddenException,
   Get,
-  Inject,
+  Logger,
   NotFoundException,
   Param,
   Res,
@@ -12,18 +11,19 @@ import type { Response } from 'express';
 
 import { CallTimer } from '../../metrics';
 import { PrismaService } from '../../prisma';
-import { StorageProvide } from '../../storage';
 import { DocID } from '../../utils/doc';
 import { Auth, CurrentUser, Publicable } from '../auth';
 import { DocHistoryManager, DocManager } from '../doc';
+import { WorkspaceBlobStorage } from '../storage';
 import { UserType } from '../users';
 import { PermissionService, PublicPageMode } from './permission';
 import { Permission } from './types';
 
 @Controller('/api/workspaces')
 export class WorkspacesController {
+  logger = new Logger(WorkspacesController.name);
   constructor(
-    @Inject(StorageProvide) private readonly storage: Storage,
+    private readonly storage: WorkspaceBlobStorage,
     private readonly permission: PermissionService,
     private readonly docManager: DocManager,
     private readonly historyManager: DocHistoryManager,
@@ -40,19 +40,26 @@ export class WorkspacesController {
     @Param('name') name: string,
     @Res() res: Response
   ) {
-    const blob = await this.storage.getBlob(workspaceId, name);
+    const { body, metadata } = await this.storage.get(workspaceId, name);
 
-    if (!blob) {
+    if (!body) {
       throw new NotFoundException(
         `Blob not found in workspace ${workspaceId}: ${name}`
       );
     }
 
-    res.setHeader('content-type', blob.contentType);
-    res.setHeader('last-modified', blob.lastModified);
-    res.setHeader('content-length', blob.size);
+    // metadata should always exists if body is not null
+    if (metadata) {
+      res.setHeader('content-type', metadata.contentType);
+      res.setHeader('last-modified', metadata.lastModified.toISOString());
+      res.setHeader('content-length', metadata.contentLength);
+      res.setHeader('x-checksum-crc32', metadata.checksumCRC32);
+    } else {
+      this.logger.warn(`Blob ${workspaceId}/${name} has no metadata`);
+    }
 
-    res.send(blob.data);
+    res.setHeader('cache-control', 'public, max-age=31536000, immutable');
+    body.pipe(res);
   }
 
   // get doc binary

packages/backend/server/src/modules/workspaces/index.ts

@@ -2,14 +2,19 @@ import { Module } from '@nestjs/common';
 
 import { DocModule } from '../doc';
 import { QuotaModule } from '../quota';
+import { StorageModule } from '../storage';
 import { UsersService } from '../users';
 import { WorkspacesController } from './controller';
-import { DocHistoryResolver } from './history.resolver';
 import { PermissionService } from './permission';
-import { PagePermissionResolver, WorkspaceResolver } from './resolver';
+import {
+  DocHistoryResolver,
+  PagePermissionResolver,
+  WorkspaceBlobResolver,
+  WorkspaceResolver,
+} from './resolvers';
 
 @Module({
-  imports: [DocModule, QuotaModule],
+  imports: [DocModule, QuotaModule, StorageModule],
   controllers: [WorkspacesController],
   providers: [
     WorkspaceResolver,
@@ -17,8 +22,9 @@ import { PagePermissionResolver, WorkspaceResolver } from './resolver';
     UsersService,
     PagePermissionResolver,
     DocHistoryResolver,
+    WorkspaceBlobResolver,
   ],
   exports: [PermissionService],
 })
 export class WorkspaceModule {}
-export { InvitationType, WorkspaceType } from './resolver';
+export { InvitationType, WorkspaceType } from './resolvers';

packages/backend/server/src/modules/workspaces/resolvers/blob.ts

@@ -0,0 +1,175 @@
+import { ForbiddenException, Logger, UseGuards } from '@nestjs/common';
+import {
+  Args,
+  Float,
+  Int,
+  Mutation,
+  Parent,
+  Query,
+  ResolveField,
+  Resolver,
+} from '@nestjs/graphql';
+import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
+
+import { MakeCache, PreventCache } from '../../../cache';
+import { CloudThrottlerGuard } from '../../../throttler';
+import type { FileUpload } from '../../../types';
+import { Auth, CurrentUser } from '../../auth';
+import { QuotaManagementService } from '../../quota';
+import { WorkspaceBlobStorage } from '../../storage';
+import { UserType } from '../../users';
+import { PermissionService } from '../permission';
+import { Permission } from '../types';
+import { WorkspaceBlobSizes, WorkspaceType } from './workspace';
+
+@UseGuards(CloudThrottlerGuard)
+@Auth()
+@Resolver(() => WorkspaceType)
+export class WorkspaceBlobResolver {
+  logger = new Logger(WorkspaceBlobResolver.name);
+  constructor(
+    private readonly permissions: PermissionService,
+    private readonly quota: QuotaManagementService,
+    private readonly storage: WorkspaceBlobStorage
+  ) {}
+
+  @ResolveField(() => Int, {
+    description: 'Blobs size of workspace',
+    complexity: 2,
+  })
+  async blobsSize(@Parent() workspace: WorkspaceType) {
+    return this.storage.totalSize(workspace.id);
+  }
+
+  /**
+   * @deprecated use `workspace.blobs` instead
+   */
+  @Query(() => [String], {
+    description: 'List blobs of workspace',
+    deprecationReason: 'use `workspace.blobs` instead',
+  })
+  @MakeCache(['blobs'], ['workspaceId'])
+  async listBlobs(
+    @CurrentUser() user: UserType,
+    @Args('workspaceId') workspaceId: string
+  ) {
+    await this.permissions.checkWorkspace(workspaceId, user.id);
+
+    return this.storage
+      .list(workspaceId)
+      .then(list => list.map(item => item.key));
+  }
+
+  /**
+   * @deprecated use `user.storageUsage` instead
+   */
+  @Query(() => WorkspaceBlobSizes, {
+    deprecationReason: 'use `user.storageUsage` instead',
+  })
+  async collectAllBlobSizes(@CurrentUser() user: UserType) {
+    const size = await this.quota.getUserUsage(user.id);
+    return { size };
+  }
+
+  /**
+   * @deprecated mutation `setBlob` will check blob limit & quota usage
+   */
+  @Query(() => WorkspaceBlobSizes, {
+    deprecationReason: 'no more needed',
+  })
+  async checkBlobSize(
+    @CurrentUser() user: UserType,
+    @Args('workspaceId') workspaceId: string,
+    @Args('size', { type: () => Float }) blobSize: number
+  ) {
+    const canWrite = await this.permissions.tryCheckWorkspace(
+      workspaceId,
+      user.id,
+      Permission.Write
+    );
+    if (canWrite) {
+      const size = await this.quota.checkBlobQuota(workspaceId, blobSize);
+      return { size };
+    }
+    return false;
+  }
+
+  @Mutation(() => String)
+  @PreventCache(['blobs'], ['workspaceId'])
+  async setBlob(
+    @CurrentUser() user: UserType,
+    @Args('workspaceId') workspaceId: string,
+    @Args({ name: 'blob', type: () => GraphQLUpload })
+    blob: FileUpload
+  ) {
+    await this.permissions.checkWorkspace(
+      workspaceId,
+      user.id,
+      Permission.Write
+    );
+
+    const { quota, size } = await this.quota.getWorkspaceUsage(workspaceId);
+
+    const checkExceeded = (recvSize: number) => {
+      if (!quota) {
+        throw new ForbiddenException('cannot find user quota');
+      }
+      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', () => {
+        const buffer = Buffer.concat(chunks);
+
+        if (checkExceeded(buffer.length)) {
+          reject(new ForbiddenException('storage size limit exceeded'));
+        } else {
+          resolve(buffer);
+        }
+      });
+    });
+
+    if (!(await this.quota.checkBlobQuota(workspaceId, buffer.length))) {
+      throw new ForbiddenException('blob size limit exceeded');
+    }
+
+    await this.storage.put(workspaceId, blob.filename, buffer);
+    return blob.filename;
+  }
+
+  @Mutation(() => Boolean)
+  @PreventCache(['blobs'], ['workspaceId'])
+  async deleteBlob(
+    @CurrentUser() user: UserType,
+    @Args('workspaceId') workspaceId: string,
+    @Args('hash') name: string
+  ) {
+    await this.permissions.checkWorkspace(workspaceId, user.id);
+
+    await this.storage.delete(workspaceId, name);
+
+    return true;
+  }
+}

packages/backend/server/src/modules/workspaces/resolvers/history.ts

@@ -1,3 +1,4 @@
+import { UseGuards } from '@nestjs/common';
 import {
   Args,
   Field,
@@ -11,13 +12,14 @@ import {
 } from '@nestjs/graphql';
 import type { SnapshotHistory } from '@prisma/client';
 
-import { DocID } from '../../utils/doc';
-import { Auth, CurrentUser } from '../auth';
-import { DocHistoryManager } from '../doc/history';
-import { UserType } from '../users';
-import { PermissionService } from './permission';
-import { WorkspaceType } from './resolver';
-import { Permission } from './types';
+import { CloudThrottlerGuard } from '../../../throttler';
+import { DocID } from '../../../utils/doc';
+import { Auth, CurrentUser } from '../../auth';
+import { DocHistoryManager } from '../../doc/history';
+import { UserType } from '../../users';
+import { PermissionService } from '../permission';
+import { Permission } from '../types';
+import { WorkspaceType } from './workspace';
 
 @ObjectType()
 class DocHistoryType implements Partial<SnapshotHistory> {
@@ -31,6 +33,7 @@ class DocHistoryType implements Partial<SnapshotHistory> {
   timestamp!: Date;
 }
 
+@UseGuards(CloudThrottlerGuard)
 @Resolver(() => WorkspaceType)
 export class DocHistoryResolver {
   constructor(

packages/backend/server/src/modules/workspaces/resolvers/index.ts

@@ -0,0 +1,4 @@
+export * from './blob';
+export * from './history';
+export * from './page';
+export * from './workspace';

packages/backend/server/src/modules/workspaces/resolvers/page.ts

@@ -0,0 +1,164 @@
+import { ForbiddenException, UseGuards } from '@nestjs/common';
+import {
+  Args,
+  Field,
+  Mutation,
+  ObjectType,
+  Parent,
+  registerEnumType,
+  ResolveField,
+  Resolver,
+} from '@nestjs/graphql';
+import type { WorkspacePage as PrismaWorkspacePage } from '@prisma/client';
+
+import { PrismaService } from '../../../prisma';
+import { CloudThrottlerGuard } from '../../../throttler';
+import { DocID } from '../../../utils/doc';
+import { Auth, CurrentUser } from '../../auth';
+import { UserType } from '../../users';
+import { PermissionService, PublicPageMode } from '../permission';
+import { Permission } from '../types';
+import { WorkspaceType } from './workspace';
+
+registerEnumType(PublicPageMode, {
+  name: 'PublicPageMode',
+  description: 'The mode which the public page default in',
+});
+
+@ObjectType()
+class WorkspacePage implements Partial<PrismaWorkspacePage> {
+  @Field(() => String, { name: 'id' })
+  pageId!: string;
+
+  @Field()
+  workspaceId!: string;
+
+  @Field(() => PublicPageMode)
+  mode!: PublicPageMode;
+
+  @Field()
+  public!: boolean;
+}
+
+@UseGuards(CloudThrottlerGuard)
+@Auth()
+@Resolver(() => WorkspaceType)
+export class PagePermissionResolver {
+  constructor(
+    private readonly prisma: PrismaService,
+    private readonly permission: PermissionService
+  ) {}
+
+  /**
+   * @deprecated
+   */
+  @ResolveField(() => [String], {
+    description: 'Shared pages of workspace',
+    complexity: 2,
+    deprecationReason: 'use WorkspaceType.publicPages',
+  })
+  async sharedPages(@Parent() workspace: WorkspaceType) {
+    const data = await this.prisma.workspacePage.findMany({
+      where: {
+        workspaceId: workspace.id,
+        public: true,
+      },
+    });
+
+    return data.map(row => row.pageId);
+  }
+
+  @ResolveField(() => [WorkspacePage], {
+    description: 'Public pages of a workspace',
+    complexity: 2,
+  })
+  async publicPages(@Parent() workspace: WorkspaceType) {
+    return this.prisma.workspacePage.findMany({
+      where: {
+        workspaceId: workspace.id,
+        public: true,
+      },
+    });
+  }
+
+  /**
+   * @deprecated
+   */
+  @Mutation(() => Boolean, {
+    name: 'sharePage',
+    deprecationReason: 'renamed to publicPage',
+  })
+  async deprecatedSharePage(
+    @CurrentUser() user: UserType,
+    @Args('workspaceId') workspaceId: string,
+    @Args('pageId') pageId: string
+  ) {
+    await this.publishPage(user, workspaceId, pageId, PublicPageMode.Page);
+    return true;
+  }
+
+  @Mutation(() => WorkspacePage)
+  async publishPage(
+    @CurrentUser() user: UserType,
+    @Args('workspaceId') workspaceId: string,
+    @Args('pageId') pageId: string,
+    @Args({
+      name: 'mode',
+      type: () => PublicPageMode,
+      nullable: true,
+      defaultValue: PublicPageMode.Page,
+    })
+    mode: PublicPageMode
+  ) {
+    const docId = new DocID(pageId, workspaceId);
+
+    if (docId.isWorkspace) {
+      throw new ForbiddenException('Expect page not to be workspace');
+    }
+
+    await this.permission.checkWorkspace(
+      docId.workspace,
+      user.id,
+      Permission.Read
+    );
+
+    return this.permission.publishPage(docId.workspace, docId.guid, mode);
+  }
+
+  /**
+   * @deprecated
+   */
+  @Mutation(() => Boolean, {
+    name: 'revokePage',
+    deprecationReason: 'use revokePublicPage',
+  })
+  async deprecatedRevokePage(
+    @CurrentUser() user: UserType,
+    @Args('workspaceId') workspaceId: string,
+    @Args('pageId') pageId: string
+  ) {
+    await this.revokePublicPage(user, workspaceId, pageId);
+    return true;
+  }
+
+  @Mutation(() => WorkspacePage)
+  async revokePublicPage(
+    @CurrentUser() user: UserType,
+    @Args('workspaceId') workspaceId: string,
+    @Args('pageId') pageId: string
+  ) {
+    const docId = new DocID(pageId, workspaceId);
+
+    if (docId.isWorkspace) {
+      throw new ForbiddenException('Expect page not to be workspace');
+    }
+
+    await this.permission.checkWorkspace(
+      docId.workspace,
+      user.id,
+      Permission.Read
+    );
+
+    return this.permission.revokePublicPage(docId.workspace, docId.guid);
+  }
+}

packages/backend/server/src/modules/workspaces/resolvers/workspace.ts

@@ -1,7 +1,5 @@
-import type { Storage } from '@affine/storage';
 import {
   ForbiddenException,
-  Inject,
   InternalServerErrorException,
   Logger,
   NotFoundException,
@@ -25,29 +23,23 @@ import {
   ResolveField,
   Resolver,
 } from '@nestjs/graphql';
-import type {
-  User,
-  Workspace,
-  WorkspacePage as PrismaWorkspacePage,
-} from '@prisma/client';
+import type { User, Workspace } from '@prisma/client';
+import { getStreamAsBuffer } from 'get-stream';
 import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs';
 import { applyUpdate, Doc } from 'yjs';
 
-import { MakeCache, PreventCache } from '../../cache';
-import { EventEmitter } from '../../event';
-import { PrismaService } from '../../prisma';
-import { StorageProvide } from '../../storage';
-import { CloudThrottlerGuard, Throttle } from '../../throttler';
-import type { FileUpload } from '../../types';
-import { DocID } from '../../utils/doc';
-import { Auth, CurrentUser, Public } from '../auth';
-import { MailService } from '../auth/mailer';
-import { AuthService } from '../auth/service';
-import { QuotaManagementService } from '../quota';
-import { UsersService, UserType } from '../users';
-import { PermissionService, PublicPageMode } from './permission';
-import { Permission } from './types';
-import { defaultWorkspaceAvatar } from './utils';
+import { EventEmitter } from '../../../event';
+import { PrismaService } from '../../../prisma';
+import { CloudThrottlerGuard, Throttle } from '../../../throttler';
+import type { FileUpload } from '../../../types';
+import { Auth, CurrentUser, Public } from '../../auth';
+import { MailService } from '../../auth/mailer';
+import { AuthService } from '../../auth/service';
+import { WorkspaceBlobStorage } from '../../storage';
+import { UsersService, UserType } from '../../users';
+import { PermissionService } from '../permission';
+import { Permission } from '../types';
+import { defaultWorkspaceAvatar } from '../utils';
 
 registerEnumType(Permission, {
   name: 'Permission',
@@ -149,8 +141,7 @@ export class WorkspaceResolver {
     private readonly permissions: PermissionService,
     private readonly users: UsersService,
     private readonly event: EventEmitter,
-    private readonly quota: QuotaManagementService,
-    @Inject(StorageProvide) private readonly storage: Storage
+    private readonly blobStorage: WorkspaceBlobStorage
   ) {}
 
   @ResolveField(() => Permission, {
@@ -235,14 +226,6 @@ export class WorkspaceResolver {
       }));
   }
 
-  @ResolveField(() => Int, {
-    description: 'Blobs size of workspace',
-    complexity: 2,
-  })
-  async blobsSize(@Parent() workspace: WorkspaceType) {
-    return this.storage.blobsSize([workspace.id]);
-  }
-
   @Query(() => Boolean, {
     description: 'Get is owner of workspace',
     complexity: 2,
@@ -565,11 +548,14 @@ export class WorkspaceResolver {
     let avatar = '';
 
     if (metaJSON.avatar) {
-      const avatarBlob = await this.storage.getBlob(
+      const avatarBlob = await this.blobStorage.get(
         workspaceId,
         metaJSON.avatar
       );
-      avatar = avatarBlob?.data.toString('base64') || '';
+
+      if (avatarBlob.body) {
+        avatar = (await getStreamAsBuffer(avatarBlob.body)).toString('base64');
+      }
     }
 
     return {
@@ -653,256 +639,4 @@ export class WorkspaceResolver {
 
     return this.permissions.revokeWorkspace(workspaceId, user.id);
   }
-
-  @Query(() => [String], {
-    description: 'List blobs of workspace',
-  })
-  @MakeCache(['blobs'], ['workspaceId'])
-  async listBlobs(
-    @CurrentUser() user: UserType,
-    @Args('workspaceId') workspaceId: string
-  ) {
-    await this.permissions.checkWorkspace(workspaceId, user.id);
-
-    return this.storage.listBlobs(workspaceId);
-  }
-
-  @Query(() => WorkspaceBlobSizes)
-  async collectAllBlobSizes(@CurrentUser() user: UserType) {
-    const size = await this.quota.getUserUsage(user.id);
-    return { size };
-  }
-
-  @Query(() => WorkspaceBlobSizes)
-  async checkBlobSize(
-    @CurrentUser() user: UserType,
-    @Args('workspaceId') workspaceId: string,
-    @Args('size', { type: () => Float }) blobSize: number
-  ) {
-    const canWrite = await this.permissions.tryCheckWorkspace(
-      workspaceId,
-      user.id,
-      Permission.Write
-    );
-    if (canWrite) {
-      const size = await this.quota.checkBlobQuota(workspaceId, blobSize);
-      return { size };
-    }
-    return false;
-  }
-
-  @Mutation(() => String)
-  @PreventCache(['blobs'], ['workspaceId'])
-  async setBlob(
-    @CurrentUser() user: UserType,
-    @Args('workspaceId') workspaceId: string,
-    @Args({ name: 'blob', type: () => GraphQLUpload })
-    blob: FileUpload
-  ) {
-    await this.permissions.checkWorkspace(
-      workspaceId,
-      user.id,
-      Permission.Write
-    );
-
-    const { quota, size } = await this.quota.getWorkspaceUsage(workspaceId);
-
-    const checkExceeded = (recvSize: number) => {
-      if (!quota) {
-        throw new ForbiddenException('cannot find user quota');
-      }
-      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', () => {
-        const buffer = Buffer.concat(chunks);
-
-        if (checkExceeded(buffer.length)) {
-          reject(new ForbiddenException('storage size limit exceeded'));
-        } else {
-          resolve(buffer);
-        }
-      });
-    });
-
-    return this.storage.uploadBlob(workspaceId, buffer);
-  }
-
-  @Mutation(() => Boolean)
-  @PreventCache(['blobs'], ['workspaceId'])
-  async deleteBlob(
-    @CurrentUser() user: UserType,
-    @Args('workspaceId') workspaceId: string,
-    @Args('hash') hash: string
-  ) {
-    await this.permissions.checkWorkspace(workspaceId, user.id);
-
-    return this.storage.deleteBlob(workspaceId, hash);
-  }
-}
-
-registerEnumType(PublicPageMode, {
-  name: 'PublicPageMode',
-  description: 'The mode which the public page default in',
-});
-
-@ObjectType()
-class WorkspacePage implements Partial<PrismaWorkspacePage> {
-  @Field(() => String, { name: 'id' })
-  pageId!: string;
-
-  @Field()
-  workspaceId!: string;
-
-  @Field(() => PublicPageMode)
-  mode!: PublicPageMode;
-
-  @Field()
-  public!: boolean;
-}
-
-@UseGuards(CloudThrottlerGuard)
-@Auth()
-@Resolver(() => WorkspaceType)
-export class PagePermissionResolver {
-  constructor(
-    private readonly prisma: PrismaService,
-    private readonly permission: PermissionService
-  ) {}
-
-  /**
-   * @deprecated
-   */
-  @ResolveField(() => [String], {
-    description: 'Shared pages of workspace',
-    complexity: 2,
-    deprecationReason: 'use WorkspaceType.publicPages',
-  })
-  async sharedPages(@Parent() workspace: WorkspaceType) {
-    const data = await this.prisma.workspacePage.findMany({
-      where: {
-        workspaceId: workspace.id,
-        public: true,
-      },
-    });
-
-    return data.map(row => row.pageId);
-  }
-
-  @ResolveField(() => [WorkspacePage], {
-    description: 'Public pages of a workspace',
-    complexity: 2,
-  })
-  async publicPages(@Parent() workspace: WorkspaceType) {
-    return this.prisma.workspacePage.findMany({
-      where: {
-        workspaceId: workspace.id,
-        public: true,
-      },
-    });
-  }
-
-  /**
-   * @deprecated
-   */
-  @Mutation(() => Boolean, {
-    name: 'sharePage',
-    deprecationReason: 'renamed to publicPage',
-  })
-  async deprecatedSharePage(
-    @CurrentUser() user: UserType,
-    @Args('workspaceId') workspaceId: string,
-    @Args('pageId') pageId: string
-  ) {
-    await this.publishPage(user, workspaceId, pageId, PublicPageMode.Page);
-    return true;
-  }
-
-  @Mutation(() => WorkspacePage)
-  async publishPage(
-    @CurrentUser() user: UserType,
-    @Args('workspaceId') workspaceId: string,
-    @Args('pageId') pageId: string,
-    @Args({
-      name: 'mode',
-      type: () => PublicPageMode,
-      nullable: true,
-      defaultValue: PublicPageMode.Page,
-    })
-    mode: PublicPageMode
-  ) {
-    const docId = new DocID(pageId, workspaceId);
-
-    if (docId.isWorkspace) {
-      throw new ForbiddenException('Expect page not to be workspace');
-    }
-
-    await this.permission.checkWorkspace(
-      docId.workspace,
-      user.id,
-      Permission.Read
-    );
-
-    return this.permission.publishPage(docId.workspace, docId.guid, mode);
-  }
-
-  /**
-   * @deprecated
-   */
-  @Mutation(() => Boolean, {
-    name: 'revokePage',
-    deprecationReason: 'use revokePublicPage',
-  })
-  async deprecatedRevokePage(
-    @CurrentUser() user: UserType,
-    @Args('workspaceId') workspaceId: string,
-    @Args('pageId') pageId: string
-  ) {
-    await this.revokePublicPage(user, workspaceId, pageId);
-    return true;
-  }
-
-  @Mutation(() => WorkspacePage)
-  async revokePublicPage(
-    @CurrentUser() user: UserType,
-    @Args('workspaceId') workspaceId: string,
-    @Args('pageId') pageId: string
-  ) {
-    const docId = new DocID(pageId, workspaceId);
-
-    if (docId.isWorkspace) {
-      throw new ForbiddenException('Expect page not to be workspace');
-    }
-
-    await this.permission.checkWorkspace(
-      docId.workspace,
-      user.id,
-      Permission.Read
-    );
-
-    return this.permission.revokePublicPage(docId.workspace, docId.guid);
-  }
 }