diff --git a/Cargo.lock b/Cargo.lock index 73c6451c30..48166959ff 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1695,7 +1695,7 @@ dependencies = [ [[package]] name = "jwst" version = "0.1.1" -source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819" +source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159" dependencies = [ "async-trait", "base64 0.21.3", @@ -1723,7 +1723,7 @@ dependencies = [ [[package]] name = "jwst-codec" version = "0.1.0" -source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819" +source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159" dependencies = [ "arbitrary", "bitvec", @@ -1743,7 +1743,7 @@ dependencies = [ [[package]] name = "jwst-logger" version = "0.1.0" -source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819" +source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159" dependencies = [ "chrono", "nu-ansi-term", @@ -1756,7 +1756,7 @@ dependencies = [ [[package]] name = "jwst-storage" version = "0.1.0" -source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819" +source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159" dependencies = [ "anyhow", "async-trait", @@ -1786,7 +1786,7 @@ dependencies = [ [[package]] name = "jwst-storage-migration" version = "0.1.0" -source = "git+https://github.com/toeverything/OctoBase.git#b026b44f67e1043e83f31a21360a0baee122e819" +source = "git+https://github.com/toeverything/OctoBase.git#58f3bbdf97f391a535e772d32828a484376c4159" dependencies = [ "sea-orm-migration", "tokio", diff --git a/apps/server/src/config/def.ts b/apps/server/src/config/def.ts index 469e338b12..2a395f722c 100644 --- a/apps/server/src/config/def.ts +++ b/apps/server/src/config/def.ts @@ -186,6 +186,11 @@ export interface AFFiNEConfig { fs: { path: string; }; + /** + * Free user storage quota + * @default 10 * 1024 * 1024 (10GB) + */ + quota: number; }; /** diff --git a/apps/server/src/config/default.ts b/apps/server/src/config/default.ts index b1df181d6b..25f16e2f65 100644 --- a/apps/server/src/config/default.ts +++ b/apps/server/src/config/default.ts @@ -55,6 +55,7 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => { AFFINE_SERVER_HOST: 'host', AFFINE_SERVER_SUB_PATH: 'path', AFFINE_ENV: 'affineEnv', + AFFINE_FREE_USER_QUOTA: 'objectStorage.quota', DATABASE_URL: 'db.url', ENABLE_R2_OBJECT_STORAGE: ['objectStorage.r2.enabled', 'boolean'], R2_OBJECT_STORAGE_ACCOUNT_ID: 'objectStorage.r2.accountId', @@ -170,6 +171,7 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => { fs: { path: join(homedir(), '.affine-storage'), }, + quota: 10 * 1024 * 1024, }, rateLimiter: { ttl: 60, diff --git a/apps/server/src/modules/workspaces/resolver.ts b/apps/server/src/modules/workspaces/resolver.ts index 647f7cbeb8..68919ca137 100644 --- a/apps/server/src/modules/workspaces/resolver.ts +++ b/apps/server/src/modules/workspaces/resolver.ts @@ -27,6 +27,7 @@ import type { User, Workspace } from '@prisma/client'; import GraphQLUpload from 'graphql-upload/GraphQLUpload.mjs'; import { applyUpdate, Doc } from 'yjs'; +import { Config } from '../../config'; import { PrismaService } from '../../prisma'; import { StorageProvide } from '../../storage'; import { CloudThrottlerGuard, Throttle } from '../../throttler'; @@ -130,6 +131,7 @@ export class UpdateWorkspaceInput extends PickType( export class WorkspaceResolver { constructor( private readonly auth: AuthService, + private readonly config: Config, private readonly mailer: MailService, private readonly prisma: PrismaService, private readonly permissionProvider: PermissionService, @@ -610,17 +612,28 @@ export class WorkspaceResolver { ) { await this.permissionProvider.check(workspaceId, user.id); - return this.storage.blobsSize(workspaceId).then(size => ({ size })); + return this.storage.blobsSize([workspaceId]).then(size => ({ size })); } @Query(() => WorkspaceBlobSizes) - async collectAllBlobSizes(@CurrentUser() user: User) { - const workspaces = await this.workspaces(user); - - const size = ( - await Promise.all(workspaces.map(({ id }) => this.storage.blobsSize(id))) - ).reduce((prev, curr) => prev + curr, 0); + async collectAllBlobSizes(@CurrentUser() user: UserType) { + const workspaces = await this.prisma.userWorkspacePermission + .findMany({ + where: { + userId: user.id, + accepted: true, + }, + select: { + workspace: { + select: { + id: true, + }, + }, + }, + }) + .then(data => data.map(({ workspace }) => workspace.id)); + const size = await this.storage.blobsSize(workspaces); return { size }; } @@ -632,6 +645,12 @@ export class WorkspaceResolver { blob: FileUpload ) { await this.permissionProvider.check(workspaceId, user.id, Permission.Write); + const quota = this.config.objectStorage.quota; + const { size } = await this.collectAllBlobSizes(user); + + if (size > quota) { + throw new ForbiddenException('storage size limit exceeded'); + } const buffer = await new Promise((resolve, reject) => { const stream = blob.createReadStream(); @@ -645,6 +664,10 @@ export class WorkspaceResolver { }); }); + if (size + buffer.length > quota) { + throw new ForbiddenException('storage size limit exceeded'); + } + return this.storage.uploadBlob(workspaceId, buffer); } diff --git a/apps/server/src/tests/utils.ts b/apps/server/src/tests/utils.ts index addef78f10..8b517b5806 100644 --- a/apps/server/src/tests/utils.ts +++ b/apps/server/src/tests/utils.ts @@ -339,7 +339,6 @@ async function collectBlobSizes( const res = await request(app.getHttpServer()) .post(gql) .auth(token, { type: 'bearer' }) - .set({ 'x-request-id': 'test', 'x-operation-name': 'test' }) .send({ query: ` query { @@ -353,6 +352,26 @@ async function collectBlobSizes( return res.body.data.collectBlobSizes.size; } +async function collectAllBlobSizes( + app: INestApplication, + token: string +): Promise { + const res = await request(app.getHttpServer()) + .post(gql) + .auth(token, { type: 'bearer' }) + .send({ + query: ` + query { + collectAllBlobSizes { + size + } + } + `, + }) + .expect(200); + return res.body.data.collectAllBlobSizes.size; +} + async function setBlob( app: INestApplication, token: string, @@ -447,6 +466,7 @@ async function getInviteInfo( export { acceptInvite, acceptInviteById, + collectAllBlobSizes, collectBlobSizes, createTestApp, createWorkspace, diff --git a/apps/server/src/tests/workspace-blobs.spec.ts b/apps/server/src/tests/workspace-blobs.spec.ts index c92f03bb24..272d7569ca 100644 --- a/apps/server/src/tests/workspace-blobs.spec.ts +++ b/apps/server/src/tests/workspace-blobs.spec.ts @@ -10,6 +10,7 @@ import request from 'supertest'; import { AppModule } from '../app'; import { + collectAllBlobSizes, collectBlobSizes, createWorkspace, listBlobs, @@ -108,4 +109,25 @@ describe('Workspace Module - Blobs', () => { const size = await collectBlobSizes(app, u1.token.token, workspace.id); ok(size === 4, 'failed to collect blob sizes'); }); + + it('should calc all blobs size', async () => { + const u1 = await signUp(app, 'u1', 'u1@affine.pro', '1'); + + const workspace1 = await createWorkspace(app, u1.token.token); + + const buffer1 = Buffer.from([0, 0]); + await setBlob(app, u1.token.token, workspace1.id, buffer1); + const buffer2 = Buffer.from([0, 1]); + await setBlob(app, u1.token.token, workspace1.id, buffer2); + + const workspace2 = await createWorkspace(app, u1.token.token); + + const buffer3 = Buffer.from([0, 0]); + await setBlob(app, u1.token.token, workspace2.id, buffer3); + const buffer4 = Buffer.from([0, 1]); + await setBlob(app, u1.token.token, workspace2.id, buffer4); + + const size = await collectAllBlobSizes(app, u1.token.token); + ok(size === 8, 'failed to collect all blob sizes'); + }); }); diff --git a/packages/storage/index.d.ts b/packages/storage/index.d.ts index ff780a3863..11051b6732 100644 --- a/packages/storage/index.d.ts +++ b/packages/storage/index.d.ts @@ -16,7 +16,7 @@ export class Storage { /** Delete a blob from workspace storage. */ deleteBlob(workspaceId: string, hash: string): Promise; /** Workspace size taken by blobs. */ - blobsSize(workspaceId: string): Promise; + blobsSize(workspaces: Array): Promise; } export interface Blob { @@ -26,5 +26,8 @@ export interface Blob { data: Buffer; } -/** Merge updates in form like `Y.applyUpdate(doc, update)` way and return the result binary. */ +/** + * Merge updates in form like `Y.applyUpdate(doc, update)` way and return the + * result binary. + */ export function mergeUpdatesInApplyWay(updates: Array): Buffer; diff --git a/packages/storage/src/lib.rs b/packages/storage/src/lib.rs index ea1afb5405..4433a99247 100644 --- a/packages/storage/src/lib.rs +++ b/packages/storage/src/lib.rs @@ -9,7 +9,6 @@ use std::{ use jwst::BlobStorage; use jwst_codec::Doc; use jwst_storage::{BlobStorageType, JwstStorage, JwstStorageError}; - use napi::{bindgen_prelude::*, Error, Result, Status}; #[macro_use] @@ -112,7 +111,11 @@ impl Storage { (id, ext.map(|ext| HashMap::from([("format".into(), ext)]))) }; - let Ok(meta) = self.blobs().get_metadata(Some(workspace_id.clone()), id.clone(), params.clone()).await else { + let Ok(meta) = self + .blobs() + .get_metadata(Some(workspace_id.clone()), id.clone(), params.clone()) + .await + else { return Ok(None); }; @@ -144,12 +147,13 @@ impl Storage { /// Workspace size taken by blobs. #[napi] - pub async fn blobs_size(&self, workspace_id: String) -> Result { - map_err!(self.blobs().get_blobs_size(workspace_id).await) + pub async fn blobs_size(&self, workspaces: Vec) -> Result { + map_err!(self.blobs().get_blobs_size(workspaces).await) } } -/// Merge updates in form like `Y.applyUpdate(doc, update)` way and return the result binary. +/// Merge updates in form like `Y.applyUpdate(doc, update)` way and return the +/// result binary. #[napi(catch_unwind)] pub fn merge_updates_in_apply_way(updates: Vec) -> Result { let mut doc = Doc::default();