From c53457691d306e606d009496f9b2bcb1f9dfbb18 Mon Sep 17 00:00:00 2001 From: DarkSky <25152247+darkskygit@users.noreply.github.com> Date: Tue, 19 May 2026 22:48:05 +0800 Subject: [PATCH] feat(server): entitlement based model (#14996) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #### PR Dependency Tree * **PR #14996** 👈 This tree was auto-generated by [Charcoal](https://github.com/danerwilliams/charcoal) ## Summary by CodeRabbit * **New Features** * Admin mutations to grant/revoke commercial entitlements. * New Doc comment-update permission. * Realtime user/workspace quota-state endpoints and live-update rooms. * **Bug Fixes** * More accurate readable-doc filtering and permission evaluation. * **Refactor** * Workspace feature management moved to entitlement-based model; permission and quota pipelines redesigned. * Admin workspace UI now edits flags only (feature toggles removed). * **Tests** * Extensive new and updated tests for permissions, entitlements, quota, projection, and backfills. [![Review Change Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/toeverything/AFFiNE/pull/14996?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack) --- .docker/selfhost/schema.json | 16 + .../native/src/permission/candidates.rs | 4 + .../native/src/permission/evaluator.rs | 3 + .../server/src/__tests__/copilot/byok.spec.ts | 140 +- .../src/__tests__/copilot/copilot.e2e.ts | 2 +- .../src/__tests__/copilot/copilot.spec.ts | 10 +- .../__tests__/copilot/tool-call-loop.spec.ts | 12 +- .../server/src/__tests__/doc/history.spec.ts | 15 + .../__tests__/e2e/comment/resolver.spec.ts | 122 +- .../e2e/doc-renderer/controller.spec.ts | 54 +- .../src/__tests__/e2e/doc/resolver.spec.ts | 58 + .../__tests__/e2e/storage/r2-proxy.spec.ts | 29 +- .../e2e/workspace/controller.spec.ts | 73 +- .../__tests__/e2e/workspace/member.spec.ts | 139 +- .../src/__tests__/e2e/workspace/team.spec.ts | 57 +- .../__tests__/mocks/team-workspace.mock.ts | 10 + .../src/__tests__/mocks/workspace.mock.ts | 49 + .../src/__tests__/models/doc-user.spec.ts | 35 + .../src/__tests__/models/feature-user.spec.ts | 26 + .../models/feature-workspace.spec.ts | 14 + .../models/permission-projection.spec.ts | 587 ++++++++ .../__tests__/models/workspace-user.spec.ts | 124 ++ .../src/__tests__/payment/event.spec.ts | 204 +++ .../storage/blob-upload-cleanup.spec.ts | 5 +- .../server/src/__tests__/worker.e2e.ts | 2 +- .../src/__tests__/workspace/blobs.e2e.ts | 41 +- .../__tests__/workspace/controller.spec.ts | 27 + .../server/src/base/metrics/metrics.ts | 1 + .../server/src/core/comment/realtime.ts | 4 +- .../server/src/core/comment/resolver.ts | 10 +- .../server/src/core/config/resolver.ts | 2 +- .../src/core/doc-renderer/controller.ts | 18 +- .../backend/server/src/core/doc/options.ts | 2 +- .../__tests__/projection-checker.spec.ts | 135 ++ .../entitlement/__tests__/projection.spec.ts | 434 ++++++ .../entitlement/__tests__/service.spec.ts | 508 +++++++ .../server/src/core/entitlement/index.ts | 23 + .../core/entitlement/projection-checker.ts | 290 ++++ .../server/src/core/entitlement/projection.ts | 465 ++++++ .../server/src/core/entitlement/service.ts | 1037 +++++++++++++ .../backend/server/src/core/features/index.ts | 2 + .../server/src/core/features/resolver.ts | 62 +- .../backend/server/src/core/features/types.ts | 16 +- .../server/src/core/notification/resolver.ts | 4 +- .../__snapshots__/actions.spec.ts.md | 8 + .../__snapshots__/actions.spec.ts.snap | Bin 1615 -> 1638 bytes .../src/core/permission/__tests__/doc.spec.ts | 114 +- .../core/permission/__tests__/docs.spec.ts | 382 +++++ .../core/permission/__tests__/policy.spec.ts | 275 ++-- .../core/permission/__tests__/service.spec.ts | 1313 +++++++++++++++++ .../permission/__tests__/workspace.spec.ts | 86 +- .../server/src/core/permission/builder.ts | 138 +- .../server/src/core/permission/config.ts | 31 + .../src/core/permission/context-loader.ts | 463 ++++++ .../server/src/core/permission/context.ts | 125 ++ .../server/src/core/permission/controller.ts | 53 - .../server/src/core/permission/diagnostic.ts | 326 ++++ .../backend/server/src/core/permission/doc.ts | 75 - .../server/src/core/permission/index.ts | 35 +- .../server/src/core/permission/policy.ts | 360 +---- .../src/core/permission/projection-checker.ts | 166 +++ .../server/src/core/permission/service.ts | 317 ++++ .../src/core/permission/sql-predicate.ts | 306 ++++ .../server/src/core/permission/types.ts | 1 + .../server/src/core/permission/workspace.ts | 210 --- .../src/core/quota/__tests__/state.spec.ts | 495 +++++++ .../backend/server/src/core/quota/index.ts | 1 + .../backend/server/src/core/quota/realtime.ts | 147 ++ .../server/src/core/quota/service.module.ts | 9 +- .../backend/server/src/core/quota/service.ts | 214 ++- .../backend/server/src/core/quota/state.ts | 413 ++++++ .../backend/server/src/core/quota/utils.ts | 6 +- .../core/realtime/__tests__/registry.spec.ts | 108 +- .../backend/server/src/core/realtime/index.ts | 2 + .../backend/server/src/core/realtime/rooms.ts | 8 + .../backend/server/src/core/sync/gateway.ts | 6 +- .../server/src/core/workspaces/controller.ts | 15 +- .../src/core/workspaces/resolvers/admin.ts | 66 +- .../src/core/workspaces/resolvers/blob.ts | 15 +- .../src/core/workspaces/resolvers/doc.ts | 50 +- .../src/core/workspaces/resolvers/history.ts | 4 +- .../src/core/workspaces/resolvers/member.ts | 54 +- .../core/workspaces/resolvers/workspace.ts | 4 +- .../server/src/core/workspaces/service.ts | 8 +- ...00000000-backfill-permission-projection.ts | 28 + ...0000000-backfill-entitlement-projection.ts | 35 + .../server/src/data/migrations/index.ts | 2 + .../src/models/__tests__/comment.spec.ts | 20 +- packages/backend/server/src/models/base.ts | 9 + .../backend/server/src/models/doc-user.ts | 142 +- packages/backend/server/src/models/doc.ts | 60 +- packages/backend/server/src/models/index.ts | 19 + .../src/models/permission-projection.ts | 555 +++++++ .../server/src/models/permission-write.ts | 742 ++++++++++ .../src/models/workspace-runtime-state.ts | 223 +++ .../src/models/workspace-user-compat.ts | 389 +++++ .../server/src/models/workspace-user.ts | 489 +++--- .../backend/server/src/models/workspace.ts | 64 +- .../server/src/plugins/calendar/resolver.ts | 8 +- .../server/src/plugins/copilot/byok/policy.ts | 42 +- .../src/plugins/copilot/byok/resolver.ts | 4 +- .../src/plugins/copilot/context/realtime.ts | 4 +- .../src/plugins/copilot/context/resolver.ts | 20 +- .../src/plugins/copilot/conversation/inbox.ts | 10 +- .../plugins/copilot/conversation/policy.ts | 12 +- .../src/plugins/copilot/mcp/provider.ts | 4 +- .../server/src/plugins/copilot/resolver.ts | 6 +- .../runtime/hosts/capability-policy-host.ts | 20 +- .../plugins/copilot/runtime/task-policy.ts | 13 +- .../plugins/copilot/runtime/tool-runtime.ts | 4 +- .../src/plugins/copilot/tools/blob-read.ts | 4 +- .../copilot/tools/doc-keyword-search.ts | 4 +- .../src/plugins/copilot/tools/doc-read.ts | 4 +- .../copilot/tools/doc-semantic-search.ts | 4 +- .../src/plugins/copilot/tools/doc-write.ts | 8 +- .../plugins/copilot/transcript/realtime.ts | 4 +- .../plugins/copilot/transcript/resolver.ts | 4 +- .../src/plugins/copilot/workspace/resolver.ts | 6 +- .../plugins/indexer/__tests__/service.spec.ts | 61 + .../server/src/plugins/indexer/index.ts | 3 +- .../server/src/plugins/indexer/resolver.ts | 157 +- .../server/src/plugins/indexer/service.ts | 18 + .../server/src/plugins/license/index.ts | 9 +- .../server/src/plugins/license/resolver.ts | 4 +- .../server/src/plugins/license/service.ts | 458 +++--- .../server/src/plugins/payment/cron.ts | 10 +- .../server/src/plugins/payment/event.ts | 87 +- .../server/src/plugins/payment/index.ts | 2 + .../src/plugins/payment/manager/user.ts | 35 +- .../src/plugins/payment/manager/workspace.ts | 13 +- .../server/src/plugins/payment/resolver.ts | 4 +- .../src/plugins/payment/revenuecat/webhook.ts | 17 +- packages/backend/server/src/schema.gql | 8 +- packages/common/graphql/src/schema.ts | 30 +- packages/common/realtime/src/index.ts | 46 +- packages/frontend/admin/src/config.json | 12 + .../admin/src/modules/settings/config.ts | 1 - .../workspaces/components/workspace-panel.tsx | 34 +- .../AdminUpdateWorkspaceInput.graphql.swift | 7 - .../modules/cloud/entities/user-quota.spec.ts | 126 ++ .../src/modules/cloud/entities/user-quota.ts | 166 ++- .../frontend/core/src/modules/cloud/index.ts | 2 +- .../src/modules/cloud/stores/user-quota.ts | 22 +- .../permissions/entities/permission.ts | 7 +- .../src/modules/quota/entities/quota.spec.ts | 130 ++ .../core/src/modules/quota/entities/quota.ts | 150 +- .../frontend/core/src/modules/quota/index.ts | 3 +- .../core/src/modules/quota/stores/quota.ts | 23 +- .../modules/workspace-engine/impls/cloud.ts | 16 +- tests/kit/src/utils/cloud.ts | 18 +- 150 files changed, 13463 insertions(+), 2458 deletions(-) create mode 100644 packages/backend/server/src/__tests__/models/permission-projection.spec.ts create mode 100644 packages/backend/server/src/__tests__/payment/event.spec.ts create mode 100644 packages/backend/server/src/core/entitlement/__tests__/projection-checker.spec.ts create mode 100644 packages/backend/server/src/core/entitlement/__tests__/projection.spec.ts create mode 100644 packages/backend/server/src/core/entitlement/__tests__/service.spec.ts create mode 100644 packages/backend/server/src/core/entitlement/index.ts create mode 100644 packages/backend/server/src/core/entitlement/projection-checker.ts create mode 100644 packages/backend/server/src/core/entitlement/projection.ts create mode 100644 packages/backend/server/src/core/entitlement/service.ts create mode 100644 packages/backend/server/src/core/permission/__tests__/service.spec.ts create mode 100644 packages/backend/server/src/core/permission/config.ts create mode 100644 packages/backend/server/src/core/permission/context-loader.ts create mode 100644 packages/backend/server/src/core/permission/context.ts delete mode 100644 packages/backend/server/src/core/permission/controller.ts create mode 100644 packages/backend/server/src/core/permission/diagnostic.ts delete mode 100644 packages/backend/server/src/core/permission/doc.ts create mode 100644 packages/backend/server/src/core/permission/projection-checker.ts create mode 100644 packages/backend/server/src/core/permission/service.ts create mode 100644 packages/backend/server/src/core/permission/sql-predicate.ts delete mode 100644 packages/backend/server/src/core/permission/workspace.ts create mode 100644 packages/backend/server/src/core/quota/__tests__/state.spec.ts create mode 100644 packages/backend/server/src/core/quota/realtime.ts create mode 100644 packages/backend/server/src/core/quota/state.ts create mode 100644 packages/backend/server/src/data/migrations/1765500000000-backfill-permission-projection.ts create mode 100644 packages/backend/server/src/data/migrations/1765600000000-backfill-entitlement-projection.ts create mode 100644 packages/backend/server/src/models/permission-projection.ts create mode 100644 packages/backend/server/src/models/permission-write.ts create mode 100644 packages/backend/server/src/models/workspace-runtime-state.ts create mode 100644 packages/backend/server/src/models/workspace-user-compat.ts create mode 100644 packages/frontend/core/src/modules/cloud/entities/user-quota.spec.ts create mode 100644 packages/frontend/core/src/modules/quota/entities/quota.spec.ts diff --git a/.docker/selfhost/schema.json b/.docker/selfhost/schema.json index 68cafd5d2a..55ee949d4f 100644 --- a/.docker/selfhost/schema.json +++ b/.docker/selfhost/schema.json @@ -300,6 +300,22 @@ } } }, + "permission": { + "type": "object", + "description": "Configuration for permission module", + "properties": { + "readModel": { + "type": "string", + "description": "Permission data source for Rust evaluation\n@default \"projection\"\n@environment `AFFINE_PERMISSION_READ_MODEL`", + "default": "projection" + }, + "fallbackLegacyLoader": { + "type": "boolean", + "description": "Fallback from projection loader to legacy loader when projection input loading fails\n@default false\n@environment `AFFINE_PERMISSION_FALLBACK_LEGACY_LOADER`", + "default": false + } + } + }, "storages": { "type": "object", "description": "Configuration for storages module", diff --git a/packages/backend/native/src/permission/candidates.rs b/packages/backend/native/src/permission/candidates.rs index ef55ad02ef..78353778ba 100644 --- a/packages/backend/native/src/permission/candidates.rs +++ b/packages/backend/native/src/permission/candidates.rs @@ -84,6 +84,10 @@ fn restricted_decision(input: &PermissionEvaluationInputV1, action: &str) -> Vec return Vec::new(); } + if input.legacy_compat_mode && input.subject.allow_local && input.workspace.local { + return Vec::new(); + } + let mut restrictions = Vec::new(); if !input.runtime.known { restrictions.push(PermissionDecisionRestrictionV1 { diff --git a/packages/backend/native/src/permission/evaluator.rs b/packages/backend/native/src/permission/evaluator.rs index 63efd78bfb..fe4af7578b 100644 --- a/packages/backend/native/src/permission/evaluator.rs +++ b/packages/backend/native/src/permission/evaluator.rs @@ -347,9 +347,12 @@ mod tests { local: true, ..Default::default() }; + input.runtime.known = false; + input.runtime.stale = true; input.workspace_actions = vec!["Workspace.Delete".to_string()]; let output = evaluate_permission(input).unwrap(); assert!(decision(&output.workspace.decisions, "Workspace.Delete").allowed); + assert!(decision(&output.docs[0].decisions, "Doc.Update").allowed); } #[test] diff --git a/packages/backend/server/src/__tests__/copilot/byok.spec.ts b/packages/backend/server/src/__tests__/copilot/byok.spec.ts index 2a09c71c09..67148e6f3a 100644 --- a/packages/backend/server/src/__tests__/copilot/byok.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/byok.spec.ts @@ -5,6 +5,7 @@ import ava, { type ExecutionContext, type TestFn } from 'ava'; import Sinon from 'sinon'; import { Cache, CryptoHelper } from '../../base'; +import { EntitlementService } from '../../core/entitlement'; import { Models, WorkspaceRole } from '../../models'; import { CopilotAccessPolicy } from '../../plugins/copilot/access'; import { ByokService } from '../../plugins/copilot/byok'; @@ -14,6 +15,11 @@ import { ByokKeyTestStatus, ByokProvider, } from '../../plugins/copilot/byok/types'; +import { + SubscriptionPlan, + SubscriptionRecurring, + SubscriptionStatus, +} from '../../plugins/payment/types'; import { createTestingModule, type TestingModule } from '../utils'; interface Context { @@ -24,11 +30,18 @@ interface Context { byok: ByokService; crypto: CryptoHelper; cache: Cache; + entitlement: EntitlementService; } -const test = ava as TestFn; +const test = ava.serial as TestFn; +const originalNamespace = globalThis.env.NAMESPACE; +const originalDeploymentType = globalThis.env.DEPLOYMENT_TYPE; test.before(async t => { + Object.assign(globalThis.env, { + NAMESPACE: 'dev', + DEPLOYMENT_TYPE: 'affine', + }); const module = await createTestingModule(); t.context.module = module; t.context.models = module.get(Models); @@ -37,6 +50,7 @@ test.before(async t => { t.context.byok = module.get(ByokService); t.context.crypto = module.get(CryptoHelper); t.context.cache = module.get(Cache); + t.context.entitlement = module.get(EntitlementService); }); test.beforeEach(async t => { @@ -45,6 +59,10 @@ test.beforeEach(async t => { test.after.always(async t => { await t.context.module.close(); + Object.assign(globalThis.env, { + NAMESPACE: originalNamespace, + DEPLOYMENT_TYPE: originalDeploymentType, + }); }); async function createUserWorkspace(t: ExecutionContext) { @@ -59,6 +77,73 @@ function workspaceHash(workspaceId: string) { return createHash('sha256').update(workspaceId).digest('hex').slice(0, 12); } +async function grantUserPlan( + t: ExecutionContext, + userId: string, + feature: ByokUserPlanFeature = 'pro_plan_v1' +) { + if (feature === 'unlimited_copilot') { + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: userId, + plan: SubscriptionPlan.AI, + recurring: SubscriptionRecurring.Monthly, + status: SubscriptionStatus.Active, + }); + return; + } + + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: userId, + plan: SubscriptionPlan.Pro, + recurring: + feature === 'lifetime_pro_plan_v1' + ? SubscriptionRecurring.Lifetime + : SubscriptionRecurring.Monthly, + status: SubscriptionStatus.Active, + }); +} + +async function revokeUserPlan( + t: ExecutionContext, + userId: string, + feature: ByokUserPlanFeature = 'pro_plan_v1' +) { + if (feature === 'unlimited_copilot') { + await t.context.entitlement.revokeCloudSubscription({ + targetId: userId, + plan: SubscriptionPlan.AI, + }); + return; + } + + await t.context.entitlement.revokeCloudSubscription({ + targetId: userId, + plan: SubscriptionPlan.Pro, + }); +} + +async function grantTeamPlan( + t: ExecutionContext, + workspaceId: string +) { + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: workspaceId, + plan: SubscriptionPlan.Team, + recurring: SubscriptionRecurring.Yearly, + status: SubscriptionStatus.Active, + }); +} + +async function revokeTeamPlan( + t: ExecutionContext, + workspaceId: string +) { + await t.context.entitlement.revokeCloudSubscription({ + targetId: workspaceId, + plan: SubscriptionPlan.Team, + }); +} + type ByokMatrixCase = { name: string; role: WorkspaceRole; @@ -110,25 +195,13 @@ async function createByokMatrixWorkspace( ); } if (input.team) { - await t.context.models.workspaceFeature.add( - workspace.id, - 'team_plan_v1', - 'test' - ); + await grantTeamPlan(t, workspace.id); } if (input.ownerPlan) { - await t.context.models.userFeature.add( - owner.id, - input.ownerPlanFeature ?? 'pro_plan_v1', - 'test' - ); + await grantUserPlan(t, owner.id, input.ownerPlanFeature); } if (input.actorPlan && actor.id !== owner.id) { - await t.context.models.userFeature.add( - actor.id, - input.actorPlanFeature ?? 'pro_plan_v1', - 'test' - ); + await grantUserPlan(t, actor.id, input.actorPlanFeature); } return { owner, actor, workspace }; @@ -252,7 +325,7 @@ for (const matrixCase of byokManagementMatrix) { test('byok service persists encrypted server keys and never returns plaintext', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const primary = await t.context.byok.upsertConfig({ workspaceId: workspace.id, @@ -325,7 +398,7 @@ test('byok service persists encrypted server keys and never returns plaintext', test('byok service preserves server key fields during partial updates', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const key = await t.context.byok.upsertConfig({ workspaceId: workspace.id, @@ -381,7 +454,7 @@ test('byok service preserves server key fields during partial updates', async t test('local leases are short lived and do not persist keys to server configs', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const before = Date.now(); const lease = await t.context.byok.createLocalLease({ @@ -486,7 +559,7 @@ test('local leases persist normalized custom endpoints', async t => { ).get(() => true); t.teardown(() => customEndpointSupported.restore()); const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const lease = await t.context.byok.createLocalLease({ workspaceId: workspace.id, @@ -659,13 +732,10 @@ for (const matrixCase of byokProfileAvailabilityMatrix) { } if (matrixCase.revokeOwnerPlan) { - await t.context.models.userFeature.remove(owner.id, 'pro_plan_v1'); + await revokeUserPlan(t, owner.id); } if (matrixCase.revokeTeam) { - await t.context.models.workspaceFeature.remove( - workspace.id, - 'team_plan_v1' - ); + await revokeTeamPlan(t, workspace.id); } if (matrixCase.demoteActor) { await t.context.models.workspaceUser.set( @@ -695,7 +765,7 @@ test('BYOK profile availability: local-only workspace does not resolve BYOK prof const user = await t.context.models.user.create({ email: `${randomUUID()}@affine.pro`, }); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const profiles = await t.context.byok.getProfiles({ workspaceId: randomUUID(), @@ -707,7 +777,7 @@ test('BYOK profile availability: local-only workspace does not resolve BYOK prof test('test key failure disables a saved key and success restores it', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const key = await t.context.byok.upsertConfig({ workspaceId: workspace.id, userId: user.id, @@ -778,7 +848,7 @@ test('test key failure disables a saved key and success restores it', async t => test('local key test does not mutate saved server config', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const key = await t.context.byok.upsertConfig({ workspaceId: workspace.id, userId: user.id, @@ -817,7 +887,7 @@ test('local key test does not mutate saved server config', async t => { test('Gemini key test sends key in header and returns safe failure message', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const fetch = Sinon.stub(globalThis, 'fetch').resolves( new Response( @@ -852,7 +922,7 @@ test('Gemini key test sends key in header and returns safe failure message', asy test('FAL key test uses read-only platform API probe endpoint', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const fetch = Sinon.stub(globalThis, 'fetch').resolves( new Response('{}', { status: 200 }) @@ -877,7 +947,7 @@ test('FAL key test uses read-only platform API probe endpoint', async t => { test('provider test failures do not return raw provider response body', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const cases = [ { body: 'authorization: Bearer token=a+b%2F==', @@ -925,7 +995,7 @@ test('provider test failures do not return raw provider response body', async t test('dispatch failure disables server BYOK key by provider id', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const key = await t.context.byok.upsertConfig({ workspaceId: workspace.id, userId: user.id, @@ -956,7 +1026,7 @@ test('dispatch failure disables server BYOK key by provider id', async t => { test('dispatch accounting ignores provider ids from another workspace hash', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const otherWorkspace = await t.context.models.workspace.create(user.id); const key = await t.context.byok.upsertConfig({ workspaceId: workspace.id, @@ -996,7 +1066,7 @@ test('dispatch accounting ignores provider ids from another workspace hash', asy test('effective profiles use local lease before server keys and skip disabled keys', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const serverKey = await t.context.byok.upsertConfig({ workspaceId: workspace.id, userId: user.id, @@ -1067,7 +1137,7 @@ test('effective profiles use local lease before server keys and skip disabled ke test('capability warnings match server Gemini background coverage', async t => { const { user, workspace } = await createUserWorkspace(t); - await t.context.models.userFeature.add(user.id, 'pro_plan_v1', 'test'); + await grantUserPlan(t, user.id); const emptySettings = await t.context.byok.getSettings(workspace.id, user.id); t.deepEqual( diff --git a/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts b/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts index f402c06cf6..e96d1642f5 100644 --- a/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts +++ b/packages/backend/server/src/__tests__/copilot/copilot.e2e.ts @@ -732,7 +732,7 @@ test('should be able to chat with special image model', async t => { promptName ); const messageId = await createCopilotMessage(app, sessionId, 'some-tag', [ - `https://example.com/${promptName}.jpg`, + smallestPng, ]); const ret3 = await chatWithImages(app, sessionId, messageId); t.is( diff --git a/packages/backend/server/src/__tests__/copilot/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot/copilot.spec.ts index 2dea559157..1edd3b5773 100644 --- a/packages/backend/server/src/__tests__/copilot/copilot.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/copilot.spec.ts @@ -17,6 +17,7 @@ import { import { ConfigModule } from '../../base/config'; import { AuthService } from '../../core/auth'; import { QuotaModule } from '../../core/quota'; +import { QuotaStateService } from '../../core/quota/state'; import { StorageModule, WorkspaceBlobStorage } from '../../core/storage'; import { ContextCategories, @@ -101,6 +102,7 @@ type Context = { actionBridge: ActionRuntimeBridge; cronJobs: CopilotCronJobs; subscription: SubscriptionService; + quotaState: QuotaStateService; }; const buildTurn = ( @@ -199,6 +201,7 @@ test.before(async t => { const workspaceEmbedding = module.get(CopilotWorkspaceService); const cronJobs = module.get(CopilotCronJobs); const subscription = module.get(SubscriptionService); + const quotaState = module.get(QuotaStateService); t.context.module = module; t.context.auth = auth; @@ -225,6 +228,7 @@ test.before(async t => { t.context.workspaceEmbedding = workspaceEmbedding; t.context.cronJobs = cronJobs; t.context.subscription = subscription; + t.context.quotaState = quotaState; await module.initTestingDB(); }); @@ -2172,7 +2176,7 @@ test('model selection policy should resolve requested optional models consistent }); test('capability policy host should gate pro model requests by subscription status', async t => { - const { subscription, module } = t.context; + const { quotaState, subscription, module } = t.context; const capabilityPolicy = module.get(CapabilityPolicyHost); const mockStatus = (status?: SubscriptionStatus) => { @@ -2181,6 +2185,10 @@ test('capability policy host should gate pro model requests by subscription stat // @ts-expect-error mock getSubscription: async () => (status ? { status } : null), })); + Sinon.stub(quotaState, 'reconcileUserQuotaState').resolves({ + plan: status === SubscriptionStatus.Active ? 'pro' : 'free', + flags: {}, + } as Awaited>); }; // payment disabled -> allow requested if in optional; pro not blocked diff --git a/packages/backend/server/src/__tests__/copilot/tool-call-loop.spec.ts b/packages/backend/server/src/__tests__/copilot/tool-call-loop.spec.ts index c15388126e..614795293e 100644 --- a/packages/backend/server/src/__tests__/copilot/tool-call-loop.spec.ts +++ b/packages/backend/server/src/__tests__/copilot/tool-call-loop.spec.ts @@ -3,7 +3,7 @@ import test from 'ava'; import { z } from 'zod'; import type { DocReader } from '../../core/doc'; -import type { AccessController } from '../../core/permission'; +import type { PermissionAccess } from '../../core/permission'; import type { Models } from '../../models'; import { LlmRequest, @@ -404,7 +404,7 @@ test('doc_read should return specific sync errors for unavailable docs', async t user: () => ({ workspace: () => ({ doc: () => ({ can: async () => true }) }), }), - } as unknown as AccessController; + } as unknown as PermissionAccess; for (const testCase of cases) { let docReaderCalled = false; @@ -447,7 +447,7 @@ test('document search tools should return sync error for local workspace', async docs: async () => [], }), }), - } as unknown as AccessController; + } as unknown as PermissionAccess; const models = { workspace: { @@ -510,7 +510,7 @@ test('doc_semantic_search should return empty array when nothing matches', async docs: async () => [], }), }), - } as unknown as AccessController; + } as unknown as PermissionAccess; const models = { workspace: { @@ -542,7 +542,7 @@ test('doc_semantic_search should pass BYOK route context into embedding matches' docs: async () => [], }), }), - } as unknown as AccessController; + } as unknown as PermissionAccess; const models = { workspace: { @@ -595,7 +595,7 @@ test('blob_read should return explicit error when attachment context is missing' }), }), }), - } as unknown as AccessController; + } as unknown as PermissionAccess; const blobTool = createBlobReadTool( buildBlobContentGetter(ac, null).bind(null, { diff --git a/packages/backend/server/src/__tests__/doc/history.spec.ts b/packages/backend/server/src/__tests__/doc/history.spec.ts index 7aa15002c2..dfd9e7f7d5 100644 --- a/packages/backend/server/src/__tests__/doc/history.spec.ts +++ b/packages/backend/server/src/__tests__/doc/history.spec.ts @@ -57,6 +57,21 @@ function getSnapshot(timestamp: number = Date.now()): DocRecord { }; } +test('history max age converts quota seconds to milliseconds', async t => { + Sinon.restore(); + const options = m.get(DocStorageOptions); + // @ts-expect-error private service boundary is asserted here + Sinon.stub(options.quota, 'getWorkspaceQuota').resolves({ + name: 'Pro', + blobLimit: 1, + storageQuota: 1, + historyPeriod: 30, + memberLimit: 1, + }); + + t.is(await options.historyMaxAge('1'), 30_000); +}); + test('should create doc history if never created before', async t => { // @ts-expect-error private method Sinon.stub(adapter, 'lastDocHistory').resolves(null); diff --git a/packages/backend/server/src/__tests__/e2e/comment/resolver.spec.ts b/packages/backend/server/src/__tests__/e2e/comment/resolver.spec.ts index 1aea976b36..9eb91b00c5 100644 --- a/packages/backend/server/src/__tests__/e2e/comment/resolver.spec.ts +++ b/packages/backend/server/src/__tests__/e2e/comment/resolver.spec.ts @@ -273,16 +273,64 @@ e2e('should update comment work', async t => { t.truthy(result.updateComment); }); -e2e('should update comment failed by another user', async t => { +e2e('should update comment work by doc Editor', async t => { const docId = randomUUID(); + await app.create(Mockers.DocUser, { + workspaceId: teamWorkspace.id, + docId, + userId: member.id, + type: DocRole.Editor, + }); await app.login(owner); - const createResult = await app.gql({ query: createCommentMutation, variables: { input: { - workspaceId: workspace.id, + workspaceId: teamWorkspace.id, + docId, + docMode: DocMode.page, + docTitle: 'test', + content: { + type: 'paragraph', + content: [{ type: 'text', text: 'test' }], + }, + }, + }, + }); + + await app.login(member); + const result = await app.gql({ + query: updateCommentMutation, + variables: { + input: { + id: createResult.createComment.id, + content: { + type: 'paragraph', + content: [{ type: 'text', text: 'test update' }], + }, + }, + }, + }); + + t.truthy(result.updateComment); +}); + +e2e('should update comment failed without update permission', async t => { + const docId = randomUUID(); + await app.create(Mockers.DocUser, { + workspaceId: teamWorkspace.id, + docId, + userId: member.id, + type: DocRole.Reader, + }); + + await app.login(owner); + const createResult = await app.gql({ + query: createCommentMutation, + variables: { + input: { + workspaceId: teamWorkspace.id, docId, docMode: DocMode.page, docTitle: 'test', @@ -1145,15 +1193,79 @@ e2e('should update reply work when user is reply owner', async t => { t.truthy(result.updateReply); }); -e2e('should update reply failed when user is not reply owner', async t => { +e2e('should update reply work by doc Editor', async t => { const docId = randomUUID(); + await app.create(Mockers.DocUser, { + workspaceId: teamWorkspace.id, + docId, + userId: member.id, + type: DocRole.Editor, + }); await app.login(owner); const createResult = await app.gql({ query: createCommentMutation, variables: { input: { - workspaceId: workspace.id, + workspaceId: teamWorkspace.id, + docId, + docMode: DocMode.page, + docTitle: 'test', + content: { + type: 'paragraph', + content: [{ type: 'text', text: 'test' }], + }, + }, + }, + }); + + const createReplyResult = await app.gql({ + query: createReplyMutation, + variables: { + input: { + commentId: createResult.createComment.id, + docMode: DocMode.page, + docTitle: 'test', + content: { + type: 'paragraph', + content: [{ type: 'text', text: 'test' }], + }, + }, + }, + }); + + await app.login(member); + const result = await app.gql({ + query: updateReplyMutation, + variables: { + input: { + id: createReplyResult.createReply.id, + content: { + type: 'paragraph', + content: [{ type: 'text', text: 'test update' }], + }, + }, + }, + }); + + t.truthy(result.updateReply); +}); + +e2e('should update reply failed without update permission', async t => { + const docId = randomUUID(); + await app.create(Mockers.DocUser, { + workspaceId: teamWorkspace.id, + docId, + userId: member.id, + type: DocRole.Reader, + }); + + await app.login(owner); + const createResult = await app.gql({ + query: createCommentMutation, + variables: { + input: { + workspaceId: teamWorkspace.id, docId, docMode: DocMode.page, docTitle: 'test', diff --git a/packages/backend/server/src/__tests__/e2e/doc-renderer/controller.spec.ts b/packages/backend/server/src/__tests__/e2e/doc-renderer/controller.spec.ts index 4622ec2955..458b6d4fa1 100644 --- a/packages/backend/server/src/__tests__/e2e/doc-renderer/controller.spec.ts +++ b/packages/backend/server/src/__tests__/e2e/doc-renderer/controller.spec.ts @@ -28,37 +28,43 @@ e2e('should render doc share page with apple-itunes-app meta tag', async t => { ); }); -e2e( +e2e.serial( 'should render doc share page without apple-itunes-app meta tag when selfhosted', async t => { + const previousDeploymentType = globalThis.env.DEPLOYMENT_TYPE; // @ts-expect-error override globalThis.env.DEPLOYMENT_TYPE = 'selfhosted'; - await using app = await createApp(); + try { + await using app = await createApp(); - const owner = await app.signup(); - const workspace = await app.create(Mockers.Workspace, { - owner, - }); + const owner = await app.signup(); + const workspace = await app.create(Mockers.Workspace, { + owner, + }); - const docSnapshot = await app.create(Mockers.DocSnapshot, { - workspaceId: workspace.id, - user: owner, - }); - // set public to true - await app.create(Mockers.DocMeta, { - workspaceId: workspace.id, - docId: docSnapshot.id, - public: true, - }); + const docSnapshot = await app.create(Mockers.DocSnapshot, { + workspaceId: workspace.id, + user: owner, + }); + // set public to true + await app.create(Mockers.DocMeta, { + workspaceId: workspace.id, + docId: docSnapshot.id, + public: true, + }); - const res = await app - .GET(`/workspace/${workspace.id}/${docSnapshot.id}`) - .expect(200) - .expect('Content-Type', 'text/html; charset=utf-8'); + const res = await app + .GET(`/workspace/${workspace.id}/${docSnapshot.id}`) + .expect(200) + .expect('Content-Type', 'text/html; charset=utf-8'); - t.notRegex( - res.text, - // - ); + t.notRegex( + res.text, + // + ); + } finally { + // @ts-expect-error restore mutable test env singleton + globalThis.env.DEPLOYMENT_TYPE = previousDeploymentType; + } } ); diff --git a/packages/backend/server/src/__tests__/e2e/doc/resolver.spec.ts b/packages/backend/server/src/__tests__/e2e/doc/resolver.spec.ts index 9ef6bd71b2..93ddb70685 100644 --- a/packages/backend/server/src/__tests__/e2e/doc/resolver.spec.ts +++ b/packages/backend/server/src/__tests__/e2e/doc/resolver.spec.ts @@ -69,6 +69,64 @@ e2e('should get recently updated docs', async t => { t.is(recentlyUpdatedDocs.edges[2].node.title, doc1.title); }); +e2e('should filter recently updated docs by doc read permission', async t => { + const owner = await app.signup(); + const member = await app.createUser(); + await app.login(member); + + await app.switchUser(owner); + const workspace = await app.create(Mockers.Workspace, { + owner: { id: owner.id }, + }); + await app.create(Mockers.WorkspaceUser, { + workspaceId: workspace.id, + userId: member.id, + type: WorkspaceRole.Collaborator, + }); + + const privateSnapshot = await app.create(Mockers.DocSnapshot, { + workspaceId: workspace.id, + user: owner, + }); + await app.create(Mockers.DocMeta, { + workspaceId: workspace.id, + docId: privateSnapshot.id, + title: 'private-doc', + defaultRole: DocRole.None, + }); + + const publicSnapshot = await app.create(Mockers.DocSnapshot, { + workspaceId: workspace.id, + user: owner, + }); + const publicDoc = await app.create(Mockers.DocMeta, { + workspaceId: workspace.id, + docId: publicSnapshot.id, + title: 'public-doc', + defaultRole: DocRole.None, + public: true, + }); + + await app.switchUser(member); + const { + workspace: { recentlyUpdatedDocs }, + } = await app.gql({ + query: getRecentlyUpdatedDocsQuery, + variables: { + workspaceId: workspace.id, + pagination: { + first: 10, + }, + }, + }); + + t.is(recentlyUpdatedDocs.totalCount, 1); + t.deepEqual( + recentlyUpdatedDocs.edges.map(edge => edge.node.id), + [publicDoc.docId] + ); +}); + e2e( 'should get doc with public attribute when doc snapshot not exists', async t => { diff --git a/packages/backend/server/src/__tests__/e2e/storage/r2-proxy.spec.ts b/packages/backend/server/src/__tests__/e2e/storage/r2-proxy.spec.ts index 8f98fbc733..402fa0d1ca 100644 --- a/packages/backend/server/src/__tests__/e2e/storage/r2-proxy.spec.ts +++ b/packages/backend/server/src/__tests__/e2e/storage/r2-proxy.spec.ts @@ -15,9 +15,18 @@ import { R2StorageProvider, } from '../../../base/storage/providers/r2'; import { SIGNED_URL_EXPIRED } from '../../../base/storage/providers/utils'; -import { WorkspaceBlobStorage } from '../../../core/storage'; +import { EntitlementService } from '../../../core/entitlement'; +import { + CommentAttachmentStorage, + WorkspaceBlobStorage, +} from '../../../core/storage'; import { MULTIPART_THRESHOLD } from '../../../core/storage/constants'; import { R2UploadController } from '../../../core/storage/r2-proxy'; +import { + SubscriptionPlan, + SubscriptionRecurring, + SubscriptionStatus, +} from '../../../plugins/payment/types'; import { app, e2e, Mockers } from '../test'; class MockR2Provider extends R2StorageProvider { @@ -160,6 +169,8 @@ async function setBlobStorage(storage: StorageProviderConfig) { configFactory.override({ storages: { blob: { storage } } }); const blobStorage = app.get(WorkspaceBlobStorage); await blobStorage.onConfigInit(); + const commentAttachmentStorage = app.get(CommentAttachmentStorage); + await commentAttachmentStorage.onConfigInit(); const controller = app.get(R2UploadController); // reset cached provider in controller (controller as any).provider = null; @@ -245,7 +256,13 @@ async function getBlobUploadPartUrl( } async function setupWorkspace() { - const owner = await app.signup({ feature: 'pro_plan_v1' }); + const owner = await app.signup(); + await app.get(EntitlementService).upsertFromCloudSubscription({ + targetId: owner.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Monthly, + status: SubscriptionStatus.Active, + }); const workspace = await app.create(Mockers.Workspace, { owner }); return { owner, workspace }; } @@ -435,7 +452,13 @@ e2e( e2e( 'should still fallback to graphql when provider does not support presign', async t => { - await setBlobStorage(defaultBlobStorage); + await setBlobStorage({ + provider: 'fs', + bucket: 'test-fallback-bucket', + config: { + path: '/tmp/affine-r2-proxy-test', + }, + }); const { workspace } = await setupWorkspace(); const buffer = Buffer.from('graph'); diff --git a/packages/backend/server/src/__tests__/e2e/workspace/controller.spec.ts b/packages/backend/server/src/__tests__/e2e/workspace/controller.spec.ts index 6fdc6017dc..09d57cdde3 100644 --- a/packages/backend/server/src/__tests__/e2e/workspace/controller.spec.ts +++ b/packages/backend/server/src/__tests__/e2e/workspace/controller.spec.ts @@ -1,6 +1,11 @@ import { randomUUID } from 'node:crypto'; import { mock } from 'node:test'; +import { + Config, + ConfigFactory, + type StorageProviderConfig, +} from '../../../base'; import { CommentAttachmentStorage } from '../../../core/storage'; import { Mockers } from '../../mocks'; import { app, e2e } from '../test'; @@ -21,6 +26,11 @@ e2e.afterEach.always(() => { mock.reset(); }); +async function useCommentAttachmentBlobStorage(storage: StorageProviderConfig) { + app.get(ConfigFactory).override({ storages: { blob: { storage } } }); + await app.get(CommentAttachmentStorage).onConfigInit(); +} + // #region comment attachment e2e( @@ -61,35 +71,50 @@ e2e( } ); -e2e('should get comment attachment body', async t => { +e2e.serial('should get comment attachment body', async t => { + const defaultBlobStorage = structuredClone( + app.get(Config).storages.blob.storage + ); + await useCommentAttachmentBlobStorage({ + provider: 'fs', + bucket: 'test-comment-attachment', + config: { + path: '/tmp/affine-test-comment-attachment', + }, + }); + const { owner, workspace } = await createWorkspace(); await app.login(owner); - const docId = randomUUID(); - const key = randomUUID(); - const attachment = app.get(CommentAttachmentStorage); - await attachment.put( - workspace.id, - docId, - key, - 'test.txt', - Buffer.from('test'), - owner.id - ); + try { + const docId = randomUUID(); + const key = randomUUID(); + const attachment = app.get(CommentAttachmentStorage); + await attachment.put( + workspace.id, + docId, + key, + 'test.txt', + Buffer.from('test'), + owner.id + ); - const res = await app.GET( - `/api/workspaces/${workspace.id}/docs/${docId}/comment-attachments/${key}` - ); + const res = await app.GET( + `/api/workspaces/${workspace.id}/docs/${docId}/comment-attachments/${key}` + ); - t.is(res.status, 200); - t.is(res.headers['content-type'], 'text/plain'); - t.is(res.headers['content-length'], '4'); - t.is(res.headers['cache-control'], 'private, max-age=2592000, immutable'); - t.regex( - res.headers['last-modified'], - /^\w{3}, \d{2} \w{3} \d{4} \d{2}:\d{2}:\d{2} GMT$/ - ); - t.is(res.text, 'test'); + t.is(res.status, 200); + t.is(res.headers['content-type'], 'text/plain'); + t.is(res.headers['content-length'], '4'); + t.is(res.headers['cache-control'], 'private, max-age=2592000, immutable'); + t.regex( + res.headers['last-modified'], + /^\w{3}, \d{2} \w{3} \d{4} \d{2}:\d{2}:\d{2} GMT$/ + ); + t.is(res.text, 'test'); + } finally { + await useCommentAttachmentBlobStorage(defaultBlobStorage); + } }); e2e('should get comment attachment redirect url', async t => { diff --git a/packages/backend/server/src/__tests__/e2e/workspace/member.spec.ts b/packages/backend/server/src/__tests__/e2e/workspace/member.spec.ts index 2b35ad004d..fea38d2611 100644 --- a/packages/backend/server/src/__tests__/e2e/workspace/member.spec.ts +++ b/packages/backend/server/src/__tests__/e2e/workspace/member.spec.ts @@ -1,28 +1,32 @@ +import { randomUUID } from 'node:crypto'; + import { acceptInviteByInviteIdMutation, approveWorkspaceTeamMemberMutation, createInviteLinkMutation, - deleteBlobMutation, getInviteInfoQuery, getMembersByWorkspaceIdQuery, inviteByEmailsMutation, leaveWorkspaceMutation, - releaseDeletedBlobsMutation, revokeMemberPermissionMutation, WorkspaceInviteLinkExpireTime, WorkspaceMemberStatus, } from '@affine/graphql'; import { faker } from '@faker-js/faker'; +import { EntitlementService } from '../../../core/entitlement'; +import { WorkspacePolicyService } from '../../../core/permission'; import { Models } from '../../../models'; -import { FeatureConfigs } from '../../../models/common/feature'; import { SubscriptionPlan, SubscriptionRecurring, + SubscriptionStatus, } from '../../../plugins/payment/types'; import { Mockers } from '../../mocks'; import { app, e2e } from '../test'; +const TWO_BILLION_BYTES = 2_000_000_000; + async function createWorkspace() { const owner = await app.create(Mockers.User); const workspace = await app.create(Mockers.Workspace, { @@ -35,6 +39,23 @@ async function createWorkspace() { }; } +async function grantTeamPlan(workspaceId: string, quantity: number) { + await app.get(EntitlementService).upsertFromCloudSubscription({ + targetId: workspaceId, + plan: SubscriptionPlan.Team, + recurring: SubscriptionRecurring.Yearly, + status: SubscriptionStatus.Active, + quantity, + }); +} + +async function revokeTeamPlan(workspaceId: string) { + await app.get(EntitlementService).revokeCloudSubscription({ + targetId: workspaceId, + plan: SubscriptionPlan.Team, + }); +} + e2e('should invite a user', async t => { const { owner, workspace } = await createWorkspace(); const u2 = await app.create(Mockers.User); @@ -91,19 +112,16 @@ e2e('should invite a user', async t => { e2e('should re-check seat when accepting an email invitation', async t => { const { owner, workspace } = await createWorkspace(); const member = await app.create(Mockers.User); - await app.create(Mockers.TeamWorkspace, { - id: workspace.id, - quantity: 4, - }); + await grantTeamPlan(workspace.id, 12); - await app.create(Mockers.WorkspaceUser, { - workspaceId: workspace.id, - userId: (await app.create(Mockers.User)).id, - }); - await app.create(Mockers.WorkspaceUser, { - workspaceId: workspace.id, - userId: (await app.create(Mockers.User)).id, - }); + await Promise.all( + Array.from({ length: 10 }).map(async () => { + await app.create(Mockers.WorkspaceUser, { + workspaceId: workspace.id, + userId: (await app.create(Mockers.User)).id, + }); + }) + ); await app.login(owner); const invite = await app.gql({ @@ -116,10 +134,10 @@ e2e('should re-check seat when accepting an email invitation', async t => { await app.eventBus.emitAsync('workspace.members.allocateSeats', { workspaceId: workspace.id, - quantity: 4, + quantity: 12, }); - await app.models.workspaceFeature.remove(workspace.id, 'team_plan_v1'); + await revokeTeamPlan(workspace.id); await app.login(member); await t.throwsAsync( @@ -147,24 +165,6 @@ e2e.serial( async t => { const { owner, workspace } = await createWorkspace(); const member = await app.create(Mockers.User); - const freeStorageQuota = FeatureConfigs.free_plan_v1.configs.storageQuota; - const lifetimeStorageQuota = - FeatureConfigs.lifetime_pro_plan_v1.configs.storageQuota; - - FeatureConfigs.free_plan_v1.configs.storageQuota = 1; - FeatureConfigs.lifetime_pro_plan_v1.configs.storageQuota = 2; - t.teardown(() => { - FeatureConfigs.free_plan_v1.configs.storageQuota = freeStorageQuota; - FeatureConfigs.lifetime_pro_plan_v1.configs.storageQuota = - lifetimeStorageQuota; - }); - - await app.models.userFeature.switchQuota( - owner.id, - 'lifetime_pro_plan_v1', - 'test setup' - ); - await app.login(owner); const invite = await app.gql({ query: inviteByEmailsMutation, @@ -174,26 +174,26 @@ e2e.serial( }, }); - await app.models.blob.upsert({ - workspaceId: workspace.id, - key: 'overflow-blob', - mime: 'application/octet-stream', - size: 2, - status: 'completed', - uploadId: null, - }); - - await app.eventBus.emitAsync('user.subscription.canceled', { - userId: owner.id, - plan: SubscriptionPlan.Pro, - recurring: SubscriptionRecurring.Lifetime, - }); + const overflowBlobKeys = Array.from( + { length: 6 }, + (_, index) => `overflow-blob-${index}` + ); + await Promise.all( + overflowBlobKeys.map(key => + app.models.blob.upsert({ + workspaceId: workspace.id, + key, + mime: 'application/octet-stream', + size: TWO_BILLION_BYTES, + status: 'completed', + uploadId: null, + }) + ) + ); t.true( - await app.models.workspaceFeature.has( - workspace.id, - 'quota_exceeded_readonly_workspace_v1' - ) + (await app.get(WorkspacePolicyService).getWorkspaceState(workspace.id)) + .isReadonly ); await app.login(member); @@ -216,26 +216,13 @@ e2e.serial( t.is(pendingInvite.status, WorkspaceMemberStatus.Pending); await app.login(owner); - await app.gql({ - query: deleteBlobMutation, - variables: { - workspaceId: workspace.id, - key: 'overflow-blob', - permanently: false, - }, - }); - await app.gql({ - query: releaseDeletedBlobsMutation, - variables: { - workspaceId: workspace.id, - }, - }); + for (const key of overflowBlobKeys) { + await app.models.blob.delete(workspace.id, key, true); + } t.false( - await app.models.workspaceFeature.has( - workspace.id, - 'quota_exceeded_readonly_workspace_v1' - ) + (await app.get(WorkspacePolicyService).getWorkspaceState(workspace.id)) + .isReadonly ); await app.login(member); @@ -596,10 +583,7 @@ e2e( 'should invite by link and send review request notification over quota limit', async t => { const { owner, workspace } = await createWorkspace(); - await app.create(Mockers.TeamWorkspace, { - id: workspace.id, - quantity: 3, - }); + await grantTeamPlan(workspace.id, 3); await app.login(owner); const { createInviteLink } = await app.gql({ @@ -639,10 +623,7 @@ e2e( name: faker.internet.displayName({ firstName: 'Lucy' }), }); const user2 = await app.create(Mockers.User, { - email: faker.internet.email({ - firstName: 'Jeanne', - lastName: 'Doe', - }), + email: `jeanne_doe.${randomUUID()}@affine.pro`, }); await app.create(Mockers.WorkspaceUser, { workspaceId: workspace.id, diff --git a/packages/backend/server/src/__tests__/e2e/workspace/team.spec.ts b/packages/backend/server/src/__tests__/e2e/workspace/team.spec.ts index 37ee910f8d..5b862d086f 100644 --- a/packages/backend/server/src/__tests__/e2e/workspace/team.spec.ts +++ b/packages/backend/server/src/__tests__/e2e/workspace/team.spec.ts @@ -6,6 +6,7 @@ import { revokePublicPageMutation, WorkspaceMemberStatus, } from '@affine/graphql'; +import { PrismaClient } from '@prisma/client'; import { QuotaService } from '../../../core/quota/service'; import { WorkspaceRole } from '../../../models'; @@ -98,7 +99,31 @@ const revokeMember = async (workspaceId: string, userId: string) => { return revokeMember; }; -e2e('should set new invited users to AllocatingSeat', async t => { +const cancelTeamWorkspace = async (workspaceId: string) => { + const db = app.get(PrismaClient); + await db.entitlement.updateMany({ + where: { + targetType: 'workspace', + targetId: workspaceId, + plan: 'team', + }, + data: { status: 'revoked' }, + }); + await db.subscription.updateMany({ + where: { + targetId: workspaceId, + plan: SubscriptionPlan.Team, + }, + data: { status: 'canceled' }, + }); + await app.eventBus.emitAsync('workspace.subscription.canceled', { + workspaceId, + plan: SubscriptionPlan.Team, + recurring: SubscriptionRecurring.Monthly, + }); +}; + +e2e('should set new invited users to waiting-seat status', async t => { const { owner, workspace } = await createTeamWorkspace(); await app.login(owner); @@ -117,7 +142,7 @@ e2e('should set new invited users to AllocatingSeat', async t => { const invitationInfo = await getInvitationInfo( result.inviteMembers[0].inviteId! ); - t.is(invitationInfo.status, WorkspaceMemberStatus.AllocatingSeat); + t.is(invitationInfo.status, WorkspaceMemberStatus.NeedMoreSeat); }); e2e('should allocate seats', async t => { @@ -151,11 +176,11 @@ e2e('should allocate seats', async t => { }); t.is( - members.find(m => m.user.id === u1.id)?.status, + members.find(m => m.user?.id === u1.id)?.status, WorkspaceMemberStatus.Pending ); t.is( - members.find(m => m.user.id === u2.id)?.status, + members.find(m => m.user?.id === u2.id)?.status, WorkspaceMemberStatus.Accepted ); @@ -201,11 +226,11 @@ e2e('should set all rests to NeedMoreSeat', async t => { }); t.is( - members.find(m => m.user.id === u2.id)?.status, + members.find(m => m.user?.id === u2.id)?.status, WorkspaceMemberStatus.NeedMoreSeat ); t.is( - members.find(m => m.user.id === u3.id)?.status, + members.find(m => m.user?.id === u3.id)?.status, WorkspaceMemberStatus.NeedMoreSeat ); }); @@ -237,11 +262,7 @@ e2e( status: WorkspaceMemberStatus.UnderReview, }); - await app.eventBus.emitAsync('workspace.subscription.canceled', { - workspaceId: workspace.id, - plan: SubscriptionPlan.Team, - recurring: SubscriptionRecurring.Monthly, - }); + await cancelTeamWorkspace(workspace.id); const [members] = await app.models.workspaceUser.paginate(workspace.id, { first: 20, @@ -265,11 +286,7 @@ e2e( async t => { const { workspace, owner, admin } = await createTeamWorkspace(); - await app.eventBus.emitAsync('workspace.subscription.canceled', { - workspaceId: workspace.id, - plan: SubscriptionPlan.Team, - recurring: SubscriptionRecurring.Monthly, - }); + await cancelTeamWorkspace(workspace.id); t.false(await app.models.workspace.isTeamWorkspace(workspace.id)); t.false( @@ -306,11 +323,7 @@ e2e( await app.login(owner); await publishDoc(workspace.id, 'published-doc'); - await app.eventBus.emitAsync('workspace.subscription.canceled', { - workspaceId: workspace.id, - plan: SubscriptionPlan.Team, - recurring: SubscriptionRecurring.Monthly, - }); + await cancelTeamWorkspace(workspace.id); t.false(await app.models.workspace.isTeamWorkspace(workspace.id)); t.true( @@ -325,7 +338,7 @@ e2e( ); await t.throwsAsync(publishDoc(workspace.id, 'blocked-doc')); - await t.notThrowsAsync(revokePublicDoc(workspace.id, 'published-doc')); + await t.throwsAsync(revokePublicDoc(workspace.id, 'published-doc')); const quota = await app .get(QuotaService) diff --git a/packages/backend/server/src/__tests__/mocks/team-workspace.mock.ts b/packages/backend/server/src/__tests__/mocks/team-workspace.mock.ts index ad7b79c68c..f48d9e75da 100644 --- a/packages/backend/server/src/__tests__/mocks/team-workspace.mock.ts +++ b/packages/backend/server/src/__tests__/mocks/team-workspace.mock.ts @@ -27,6 +27,16 @@ export class MockTeamWorkspace extends Mocker< quantity, }, }); + await this.db.entitlement.create({ + data: { + targetType: 'workspace', + targetId: id, + source: 'cloud_subscription', + plan: 'team', + status: 'active', + quantity, + }, + }); await this.db.workspaceFeature.create({ data: { diff --git a/packages/backend/server/src/__tests__/mocks/workspace.mock.ts b/packages/backend/server/src/__tests__/mocks/workspace.mock.ts index 5a95b78604..c23061eba9 100644 --- a/packages/backend/server/src/__tests__/mocks/workspace.mock.ts +++ b/packages/backend/server/src/__tests__/mocks/workspace.mock.ts @@ -45,6 +45,55 @@ export class MockWorkspace extends Mocker { : undefined, }, }); + const runtimeStateColumns = await this.db.$queryRaw< + Array<{ exists: boolean }> + >` + SELECT EXISTS ( + SELECT 1 + FROM information_schema.columns + WHERE table_name = 'workspace_runtime_states' + AND column_name = 'known' + ) AS "exists" + `; + if (runtimeStateColumns[0]?.exists) { + await this.db.$executeRaw` + INSERT INTO workspace_runtime_states ( + workspace_id, + known, + readonly, + readonly_reasons, + last_reconciled_at, + stale_after, + updated_at + ) + VALUES (${workspace.id}, true, false, ARRAY[]::TEXT[], now(), NULL, now()) + ON CONFLICT (workspace_id) + DO UPDATE SET + known = true, + readonly = false, + readonly_reasons = ARRAY[]::TEXT[], + last_reconciled_at = now(), + stale_after = NULL, + updated_at = now() + `; + } else { + await this.db.$executeRaw` + INSERT INTO workspace_runtime_states ( + workspace_id, + readonly, + readonly_reasons, + stale_at, + updated_at + ) + VALUES (${workspace.id}, false, ARRAY[]::TEXT[], NULL, now()) + ON CONFLICT (workspace_id) + DO UPDATE SET + readonly = false, + readonly_reasons = ARRAY[]::TEXT[], + stale_at = NULL, + updated_at = now() + `; + } // create a rootDoc snapshot if (snapshot) { diff --git a/packages/backend/server/src/__tests__/models/doc-user.spec.ts b/packages/backend/server/src/__tests__/models/doc-user.spec.ts index 9717140786..b61ea181df 100644 --- a/packages/backend/server/src/__tests__/models/doc-user.spec.ts +++ b/packages/backend/server/src/__tests__/models/doc-user.spec.ts @@ -73,6 +73,24 @@ test('should set doc user role', async t => { t.is(role?.type, DocRole.Manager); }); +test('should batch update existing doc user roles', async t => { + const workspace = await create(); + const user = await models.user.create({ email: 'u1@affine.pro' }); + const docId = 'fake-doc-id'; + + await models.docUser.set(workspace.id, docId, user.id, DocRole.Reader); + const count = await models.docUser.batchSetUserRoles( + workspace.id, + docId, + [user.id], + DocRole.Editor + ); + const role = await models.docUser.get(workspace.id, docId, user.id); + + t.is(count, 1); + t.is(role?.type, DocRole.Editor); +}); + test('should not allow setting doc owner through setDocUserRole', async t => { const workspace = await create(); const user = await models.user.create({ email: 'u1@affine.pro' }); @@ -96,6 +114,23 @@ test('should delete doc user role', async t => { t.is(role, null); }); +test('should delete doc grants by user id', async t => { + const workspace = await create(); + const user = await models.user.create({ email: 'u1@affine.pro' }); + const docId = 'fake-doc-id'; + + await models.docUser.set(workspace.id, docId, user.id, DocRole.Manager); + await models.docUser.deleteByUserId(user.id); + + t.is(await models.docUser.get(workspace.id, docId, user.id), null); + t.is( + await db.docGrant.count({ + where: { principalType: 'user', principalId: user.id }, + }), + 0 + ); +}); + test('should paginate doc user roles', async t => { const workspace = await create(); const docId = 'fake-doc-id'; diff --git a/packages/backend/server/src/__tests__/models/feature-user.spec.ts b/packages/backend/server/src/__tests__/models/feature-user.spec.ts index 30f9a5525c..bd2d6b8913 100644 --- a/packages/backend/server/src/__tests__/models/feature-user.spec.ts +++ b/packages/backend/server/src/__tests__/models/feature-user.spec.ts @@ -1,12 +1,16 @@ import { User } from '@prisma/client'; import ava, { TestFn } from 'ava'; +import { AdminFeatureManagementResolver } from '../../core/features/resolver'; +import { AvailableUserFeatureConfig } from '../../core/features/types'; import { FeatureType, Models, UserFeatureModel, UserModel } from '../../models'; +import { Feature } from '../../models/common/feature'; import { createTestingModule, TestingModule } from '../utils'; interface Context { module: TestingModule; model: UserFeatureModel; + resolver: AdminFeatureManagementResolver; u1: User; } @@ -16,6 +20,7 @@ test.before(async t => { const module = await createTestingModule({}); t.context.model = module.get(UserFeatureModel); + t.context.resolver = module.get(AdminFeatureManagementResolver); t.context.module = module; }); @@ -31,6 +36,21 @@ test.after(async t => { await t.context.module.close(); }); +test('configurable user features exclude commercial projection features', t => { + const config = new AvailableUserFeatureConfig(); + + t.false(config.availableUserFeatures().has(Feature.UnlimitedCopilot)); + t.false(config.configurableUserFeatures().has(Feature.UnlimitedCopilot)); +}); + +test('admin feature resolver rejects commercial projection features', async t => { + await t.throwsAsync( + t.context.resolver.updateUserFeatures(t.context.u1.id, [Feature.ProPlan]), + { message: /not configurable/ } + ); + t.deepEqual(await t.context.model.list(t.context.u1.id), []); +}); + test('should get null if user feature not found', async t => { const { model, u1 } = t.context; const userFeature = await model.get(u1.id, 'ai_early_access'); @@ -39,12 +59,14 @@ test('should get null if user feature not found', async t => { test('should get user feature', async t => { const { model, u1 } = t.context; + await model.add(u1.id, 'free_plan_v1', 'legacy projection'); const userFeature = await model.get(u1.id, 'free_plan_v1'); t.is(userFeature?.name, 'free_plan_v1'); }); test('should get user quota', async t => { const { model, u1 } = t.context; + await model.add(u1.id, 'free_plan_v1', 'legacy projection'); const userQuota = await model.getQuota(u1.id); t.snapshot(userQuota?.configs, 'free plan'); }); @@ -52,6 +74,7 @@ test('should get user quota', async t => { test('should list user features', async t => { const { model, u1 } = t.context; + await model.add(u1.id, 'free_plan_v1', 'legacy projection'); t.like(await model.list(u1.id), ['free_plan_v1']); }); @@ -68,6 +91,7 @@ test('should list user features by type', async t => { test('should directly test user feature existence', async t => { const { model, u1 } = t.context; + await model.add(u1.id, 'free_plan_v1', 'legacy projection'); t.true(await model.has(u1.id, 'free_plan_v1')); t.false(await model.has(u1.id, 'ai_early_access')); }); @@ -112,6 +136,7 @@ test('should switch user quota', async t => { test('should not switch user quota if the new quota is the same as the current one', async t => { const { model, u1 } = t.context; + await model.add(u1.id, 'free_plan_v1', 'legacy projection'); await model.switchQuota(u1.id, 'free_plan_v1', 'test not switch'); // @ts-expect-error private @@ -135,6 +160,7 @@ test('should use pro plan as free for selfhost instance', async t => { registered: true, }); + await models.userFeature.add(u1.id, 'free_plan_v1', 'legacy projection'); const quota = await models.userFeature.getQuota(u1.id); t.snapshot( quota?.configs, diff --git a/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts b/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts index 4e03b369c0..0fca894891 100644 --- a/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts +++ b/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts @@ -1,6 +1,7 @@ import { Workspace } from '@prisma/client'; import ava, { TestFn } from 'ava'; +import { AdminWorkspaceResolver } from '../../core/workspaces/resolvers/admin'; import { FeatureType, UserModel, @@ -12,6 +13,7 @@ import { createTestingModule, type TestingModule } from '../utils'; interface Context { module: TestingModule; model: WorkspaceFeatureModel; + resolver: AdminWorkspaceResolver; ws: Workspace; } @@ -21,6 +23,7 @@ test.before(async t => { const module = await createTestingModule({}); t.context.model = module.get(WorkspaceFeatureModel); + t.context.resolver = module.get(AdminWorkspaceResolver); t.context.module = module; }); @@ -44,6 +47,17 @@ test('should get null if workspace feature not found', async t => { t.is(userFeature, null); }); +test('admin workspace update changes workspace flags', async t => { + await t.context.resolver.adminUpdateWorkspace({ + id: t.context.ws.id, + name: 'updated', + }); + t.is( + (await t.context.module.get(WorkspaceModel).get(t.context.ws.id))?.name, + 'updated' + ); +}); + test('should directly test workspace feature existence', async t => { const { model, ws } = t.context; diff --git a/packages/backend/server/src/__tests__/models/permission-projection.spec.ts b/packages/backend/server/src/__tests__/models/permission-projection.spec.ts new file mode 100644 index 0000000000..186be3f08b --- /dev/null +++ b/packages/backend/server/src/__tests__/models/permission-projection.spec.ts @@ -0,0 +1,587 @@ +import { readFileSync } from 'node:fs'; +import { join } from 'node:path'; + +import { PrismaClient } from '@prisma/client'; +import test from 'ava'; + +import { PermissionProjectionChecker } from '../../core/permission/projection-checker'; +import { + DocRole, + PERMISSION_PROJECTION_TRIGGER_ERROR_CATEGORIES, + PermissionProjectionModel, + permissionProjectionTriggerErrorCategory, + WorkspaceMemberStatus, + WorkspaceRole, +} from '../../models'; +import { createModule } from '../create-module'; +import { Mockers } from '../mocks'; + +const module = await createModule({}); +const db = module.get(PrismaClient); + +test.after.always(async () => { + await module.close(); +}); + +class TestPermissionProjectionModel extends PermissionProjectionModel { + constructor(private readonly fakeDb: unknown) { + super(); + } + + protected override get db() { + return this.fakeDb as never; + } +} + +let appliedPermissionProjectionTriggerFunctionUpdates = false; +async function applyPermissionProjectionTriggerFunctionUpdates() { + if (appliedPermissionProjectionTriggerFunctionUpdates) { + return; + } + const migration = readFileSync( + join( + process.cwd(), + 'migrations/20260512133700_workspace_runtime_states/migration.sql' + ), + 'utf8' + ); + for (const name of [ + 'affine_permission_project_new_workspace_member', + 'affine_permission_project_new_workspace_invitation', + 'affine_permission_project_new_doc_access_policy', + 'affine_permission_project_new_doc_grant', + ]) { + const sql = migration.match( + new RegExp( + `CREATE OR REPLACE FUNCTION ${name}\\(\\)[\\s\\S]*?END\\n\\$\\$;` + ) + )?.[0]; + if (!sql) { + throw new Error(`Missing migration function ${name}`); + } + await db.$executeRawUnsafe(sql); + } + appliedPermissionProjectionTriggerFunctionUpdates = true; +} + +async function hasCurrentWorkspaceInvitationColumns() { + const rows = await db.$queryRaw<{ columnName: string }[]>` + SELECT column_name AS "columnName" + FROM information_schema.columns + WHERE table_name = 'workspace_invitations' + AND column_name IN ('requested_role', 'status', 'kind') + `; + return rows.length === 3; +} + +test('PermissionProjectionModel checker returns mismatch and dirty-row counts', async t => { + const queryResults = [ + [{ count: 1n }], + [{ count: 2n }], + [{ count: 3n }], + [{ count: 4n }], + [{ count: 5n }], + [{ count: 6n }], + [{ count: 7n }], + [{ count: 8n }], + [{ count: 9n }], + [{ count: 10n }], + [ + { category: 'legacy_doc_external_row', count: 11n }, + { category: 'doc_default_owner', count: 12n }, + ], + ]; + const model = new TestPermissionProjectionModel({ + $queryRaw: async () => queryResults.shift(), + }); + + t.deepEqual(await model.checkLegacyProjection(), { + oldWorkspacePolicyMismatch: 1, + oldAcceptedMemberMismatch: 2, + extraProjectedMember: 3, + oldInvitationMismatch: 4, + extraProjectedInvitation: 5, + oldDocGrantMismatch: 6, + extraProjectedDocGrant: 7, + oldDocPolicyMismatch: 8, + extraProjectedDocPolicy: 9, + runtimeStateMissing: 0, + runtimeStateMismatch: 0, + ownerConflict: 10, + oldNewDecisionMismatch: 0, + invalidLegacyRows: { + legacy_doc_external_row: 11, + doc_default_owner: 12, + }, + }); +}); + +test('PermissionProjectionModel backfill runs as a single transaction', async t => { + const executed: unknown[] = []; + const model = new TestPermissionProjectionModel({ + $transaction: async (callback: (tx: unknown) => Promise) => { + await callback({ + $executeRaw: async (query: unknown) => { + executed.push(query); + }, + }); + }, + }); + + await model.backfillLegacyProjection(); + + t.is(executed.length, 10); +}); + +test('PermissionProjectionModel exposes stable trigger metric categories', t => { + t.deepEqual(PERMISSION_PROJECTION_TRIGGER_ERROR_CATEGORIES, [ + 'owner_conflict', + 'invalid_legacy_role', + 'foreign_key_missing', + 'projection_recursion_guard_missing', + 'unknown', + ]); +}); + +test('permission projection migration uses non-recursive origin guard', t => { + const migration = readFileSync( + join( + process.cwd(), + 'migrations/20260512133700_workspace_runtime_states/migration.sql' + ), + 'utf8' + ); + const guardBody = migration.match( + /CREATE OR REPLACE FUNCTION affine_permission_should_project_from_legacy\(\)[\s\S]*?END\n\$\$;/ + )?.[0]; + + t.truthy(guardBody); + t.true( + guardBody?.includes('IF NOT affine_permission_projection_enabled() THEN') + ); + t.false( + guardBody?.includes('IF NOT affine_permission_should_project_from_legacy()') + ); + t.truthy( + migration.match( + /CREATE OR REPLACE FUNCTION affine_permission_should_project_from_new\(\)[\s\S]*?IF NOT affine_permission_projection_enabled\(\) THEN[\s\S]*?END\n\$\$;/ + ) + ); +}); + +test('permission projection trigger maps legacy workspace permission rows', async t => { + const workspace = await module.create(Mockers.Workspace); + const [admin, pending] = await module.create(Mockers.User, 2); + + await db.workspaceUserRole.createMany({ + data: [ + { + workspaceId: workspace.id, + userId: admin.id, + type: WorkspaceRole.Admin, + status: WorkspaceMemberStatus.Accepted, + }, + { + workspaceId: workspace.id, + userId: pending.id, + type: WorkspaceRole.Collaborator, + status: WorkspaceMemberStatus.Pending, + }, + ], + }); + + const member = await db.workspaceMember.findFirstOrThrow({ + where: { + workspaceId: workspace.id, + userId: admin.id, + state: 'active', + }, + }); + const invitation = await db.workspaceInvitation.findUniqueOrThrow({ + where: { + workspaceId_inviteeUserId: { + workspaceId: workspace.id, + inviteeUserId: pending.id, + }, + }, + }); + + t.is(member.role, 'admin'); + t.is(invitation.requestedRole, 'member'); + t.is(invitation.status, 'pending'); +}); + +test('permission projection trigger maps legacy doc policy rows', async t => { + const workspace = await module.create(Mockers.Workspace); + + await db.workspaceDoc.create({ + data: { + workspaceId: workspace.id, + docId: 'public-doc', + public: true, + defaultRole: DocRole.Reader, + }, + }); + + const policy = await db.docAccessPolicy.findUniqueOrThrow({ + where: { + workspaceId_docId: { + workspaceId: workspace.id, + docId: 'public-doc', + }, + }, + }); + + t.is(policy.visibility, 'public'); + t.is(policy.publicRole, 'external'); + t.is(policy.memberDefaultRole, 'reader'); +}); + +async function hasDocGrantLegacyProjectionColumns() { + const rows = await db.$queryRaw<{ columnName: string }[]>` + SELECT column_name AS "columnName" + FROM information_schema.columns + WHERE table_name = 'doc_grants' + AND column_name IN ( + 'legacy_workspace_id', + 'legacy_doc_id', + 'legacy_user_id' + ) + `; + return rows.length === 3; +} + +test('permission projection trigger maps legacy doc grants and drops dirty rows', async t => { + if (!(await hasDocGrantLegacyProjectionColumns())) { + t.false( + Boolean(process.env.CI), + 'current local test database predates doc_grants legacy columns' + ); + return; + } + + const workspace = await module.create(Mockers.Workspace); + const user = await module.create(Mockers.User); + + await db.workspaceDocUserRole.createMany({ + data: [ + { + workspaceId: workspace.id, + docId: 'valid-grant', + userId: user.id, + type: DocRole.Reader, + }, + { + workspaceId: workspace.id, + docId: 'dirty-external', + userId: user.id, + type: DocRole.External, + }, + { + workspaceId: workspace.id, + docId: 'dirty-none', + userId: user.id, + type: DocRole.None, + }, + ], + }); + + const grants = await db.docGrant.findMany({ + where: { + workspaceId: workspace.id, + principalId: user.id, + }, + orderBy: { + docId: 'asc', + }, + }); + + t.deepEqual( + grants.map(grant => [grant.docId, grant.role]), + [['valid-grant', 'reader']] + ); +}); + +test('permission projection trigger clears legacy row for non-active new workspace member states', async t => { + await applyPermissionProjectionTriggerFunctionUpdates(); + const workspace = await module.create(Mockers.Workspace); + const user = await module.create(Mockers.User); + + const member = await db.workspaceMember.create({ + data: { + workspaceId: workspace.id, + userId: user.id, + role: 'member', + state: 'active', + }, + }); + + t.truthy( + await db.workspaceUserRole.findUnique({ + where: { + workspaceId_userId: { + workspaceId: workspace.id, + userId: user.id, + }, + }, + }) + ); + + await db.workspaceMember.update({ + where: { id: member.id }, + data: { state: 'suspended' }, + }); + + t.is( + await db.workspaceUserRole.findUnique({ + where: { + workspaceId_userId: { + workspaceId: workspace.id, + userId: user.id, + }, + }, + }), + null + ); +}); + +test('permission projection trigger clears legacy row for terminal new invitation statuses', async t => { + if (!(await hasCurrentWorkspaceInvitationColumns())) { + t.false( + Boolean(process.env.CI), + 'current local test database predates workspace invitation projection columns' + ); + return; + } + await applyPermissionProjectionTriggerFunctionUpdates(); + const workspace = await module.create(Mockers.Workspace); + const user = await module.create(Mockers.User); + + const [invitation] = await db.$queryRaw<{ id: string }[]>` + INSERT INTO workspace_invitations ( + workspace_id, + invitee_user_id, + requested_role, + status, + kind + ) + VALUES ( + ${workspace.id}, + ${user.id}, + 'member', + 'pending', + 'email' + ) + RETURNING id + `; + + t.is( + ( + await db.workspaceUserRole.findUniqueOrThrow({ + where: { + workspaceId_userId: { + workspaceId: workspace.id, + userId: user.id, + }, + }, + }) + ).status, + 'Pending' + ); + + await db.$executeRaw` + UPDATE workspace_invitations + SET status = 'declined' + WHERE id = ${invitation.id} + `; + + t.is( + await db.workspaceUserRole.findUnique({ + where: { + workspaceId_userId: { + workspaceId: workspace.id, + userId: user.id, + }, + }, + }), + null + ); +}); + +test('permission projection trigger preserves doc metadata when new doc policy is deleted', async t => { + await applyPermissionProjectionTriggerFunctionUpdates(); + const workspace = await module.create(Mockers.Workspace); + + await db.workspaceDoc.create({ + data: { + workspaceId: workspace.id, + docId: 'metadata-doc', + public: true, + defaultRole: DocRole.Reader, + mode: 1, + blocked: true, + title: 'Title', + summary: 'Summary', + publishedAt: new Date('2026-01-01T00:00:00Z'), + }, + }); + + await db.docAccessPolicy.delete({ + where: { + workspaceId_docId: { + workspaceId: workspace.id, + docId: 'metadata-doc', + }, + }, + }); + + const doc = await db.workspaceDoc.findUniqueOrThrow({ + where: { + workspaceId_docId: { + workspaceId: workspace.id, + docId: 'metadata-doc', + }, + }, + }); + + t.is(doc.public, false); + t.is(doc.defaultRole, DocRole.Manager); + t.is(doc.publishedAt, null); + t.is(doc.mode, 1); + t.is(doc.blocked, true); + t.is(doc.title, 'Title'); + t.is(doc.summary, 'Summary'); +}); + +test('permission projection trigger ignores group doc grants on legacy projection', async t => { + await applyPermissionProjectionTriggerFunctionUpdates(); + const workspace = await module.create(Mockers.Workspace); + const user = await module.create(Mockers.User); + + await db.docGrant.create({ + data: { + workspaceId: workspace.id, + docId: 'group-doc', + principalType: 'user', + principalId: user.id, + role: 'reader', + }, + }); + await db.docGrant.create({ + data: { + workspaceId: workspace.id, + docId: 'group-doc', + principalType: 'group', + principalId: user.id, + role: 'manager', + }, + }); + await db.docGrant.delete({ + where: { + workspaceId_docId_principalType_principalId: { + workspaceId: workspace.id, + docId: 'group-doc', + principalType: 'group', + principalId: user.id, + }, + }, + }); + + const legacyGrant = await db.workspaceDocUserRole.findUniqueOrThrow({ + where: { + workspaceId_docId_userId: { + workspaceId: workspace.id, + docId: 'group-doc', + userId: user.id, + }, + }, + }); + + t.is(legacyGrant.type, DocRole.Reader); +}); + +test('PermissionProjectionModel parses trigger error metric category', t => { + t.is( + permissionProjectionTriggerErrorCategory( + new Error('permission_projection_error:owner_conflict:duplicate owner') + ), + 'owner_conflict' + ); + t.is( + permissionProjectionTriggerErrorCategory( + new Error('permission_projection_error:unexpected:nope') + ), + 'unknown' + ); + t.is(permissionProjectionTriggerErrorCategory(new Error('other')), null); +}); + +test('PermissionProjectionChecker reports old/new loader decision mismatches', async t => { + const checker = new PermissionProjectionChecker( + { + workspace: { + findMany: async () => [], + }, + $queryRaw: async () => [ + { + category: 'active_member_doc', + workspaceId: 'w1', + docId: 'doc1', + userId: 'u1', + workspaceActions: null, + docActions: ['Doc.Read'], + }, + { + category: 'explicit_doc_grant', + workspaceId: 'w1', + docId: 'doc2', + userId: 'u1', + workspaceActions: null, + docActions: ['Doc.Read'], + }, + { + category: 'workspace_invitation', + workspaceId: 'w1', + docId: null, + userId: 'u2', + workspaceActions: ['Workspace.Read'], + docActions: null, + }, + ], + } as never, + { + permissionProjection: { + checkLegacyProjection: async () => ({}), + }, + } as never, + { + load: async (input: { docs?: [{ docId: string }] }) => ({ + version: 1, + workspace: { marker: 'legacy' }, + docs: input.docs + ? [{ docId: input.docs[0].docId, marker: 'legacy' }] + : [], + }), + loadFromNewTables: async (input: { docs?: [{ docId: string }] }) => ({ + version: 1, + workspace: { marker: input.docs ? 'legacy' : 'projection' }, + docs: input.docs + ? [ + { + docId: input.docs[0].docId, + marker: + input.docs[0].docId === 'doc1' ? 'legacy' : 'projection', + }, + ] + : [], + }), + } as never, + { + evaluate: (input: unknown) => input, + } as never + ); + + t.deepEqual(await checker.checkLegacyProjection(), { + oldNewDecisionMismatch: 2, + }); +}); diff --git a/packages/backend/server/src/__tests__/models/workspace-user.spec.ts b/packages/backend/server/src/__tests__/models/workspace-user.spec.ts index 4c3a3a36c3..4f77c4c23b 100644 --- a/packages/backend/server/src/__tests__/models/workspace-user.spec.ts +++ b/packages/backend/server/src/__tests__/models/workspace-user.spec.ts @@ -151,6 +151,22 @@ test('should not get inactive workspace role', async t => { t.is(role, null); }); +test('should not activate a missing workspace invitation', async t => { + const workspace = await module.create(Mockers.Workspace); + const user = await module.create(Mockers.User); + + await t.throwsAsync( + models.workspaceUser.setStatus( + workspace.id, + user.id, + WorkspaceMemberStatus.Accepted + ), + { message: 'Cannot activate a missing workspace invitation.' } + ); + + t.is(await models.workspaceUser.get(workspace.id, user.id), null); +}); + test('should update user role', async t => { const workspace = await module.create(Mockers.Workspace); const user = await module.create(Mockers.User); @@ -215,6 +231,114 @@ test('should delete workspace user role', async t => { t.is(role, null); }); +test('should delete legacy-only external workspace user role', async t => { + const workspace = await module.create(Mockers.Workspace); + const u1 = await module.create(Mockers.User); + + await models.workspaceUser.set(workspace.id, u1.id, WorkspaceRole.External, { + status: WorkspaceMemberStatus.Accepted, + }); + + t.truthy(await models.workspaceUser.get(workspace.id, u1.id)); + + await models.workspaceUser.delete(workspace.id, u1.id); + + t.is(await models.workspaceUser.get(workspace.id, u1.id), null); +}); + +test('should convert existing workspace user role to legacy-only external role', async t => { + const workspace = await module.create(Mockers.Workspace); + const u1 = await module.create(Mockers.User); + + await models.workspaceUser.set( + workspace.id, + u1.id, + WorkspaceRole.Collaborator, + { + status: WorkspaceMemberStatus.Accepted, + } + ); + await models.workspaceUser.set(workspace.id, u1.id, WorkspaceRole.External, { + status: WorkspaceMemberStatus.Accepted, + }); + + const role = await models.workspaceUser.get(workspace.id, u1.id); + t.is(role?.type, WorkspaceRole.External); + t.is( + await db.workspaceMember.count({ + where: { + workspaceId: workspace.id, + userId: u1.id, + state: 'active', + }, + }), + 0 + ); +}); + +test('should backfill legacy permission id for new workspace member writes', async t => { + const workspace = await module.create(Mockers.Workspace); + const u1 = await module.create(Mockers.User); + + await models.workspaceUser.set( + workspace.id, + u1.id, + WorkspaceRole.Collaborator, + { + status: WorkspaceMemberStatus.Accepted, + } + ); + + const legacyRole = await db.workspaceUserRole.findUniqueOrThrow({ + where: { + workspaceId_userId: { + workspaceId: workspace.id, + userId: u1.id, + }, + }, + }); + const member = await db.workspaceMember.findFirstOrThrow({ + where: { + workspaceId: workspace.id, + userId: u1.id, + state: 'active', + }, + }); + + t.is(member.legacyPermissionId, legacyRole.id); +}); + +test('should backfill legacy permission id for new workspace invitation writes', async t => { + const workspace = await module.create(Mockers.Workspace); + const u1 = await module.create(Mockers.User); + + await models.workspaceUser.set( + workspace.id, + u1.id, + WorkspaceRole.Collaborator, + { + status: WorkspaceMemberStatus.Pending, + } + ); + + const legacyRole = await db.workspaceUserRole.findUniqueOrThrow({ + where: { + workspaceId_userId: { + workspaceId: workspace.id, + userId: u1.id, + }, + }, + }); + const invitation = await db.workspaceInvitation.findFirstOrThrow({ + where: { + workspaceId: workspace.id, + inviteeUserId: u1.id, + }, + }); + + t.is(invitation.legacyPermissionId, legacyRole.id); +}); + test('should get user workspace roles with filter', async t => { const ws1 = await module.create(Mockers.Workspace); const ws2 = await module.create(Mockers.Workspace); diff --git a/packages/backend/server/src/__tests__/payment/event.spec.ts b/packages/backend/server/src/__tests__/payment/event.spec.ts new file mode 100644 index 0000000000..742a2dcdde --- /dev/null +++ b/packages/backend/server/src/__tests__/payment/event.spec.ts @@ -0,0 +1,204 @@ +import { PrismaClient } from '@prisma/client'; +import ava, { TestFn } from 'ava'; + +import { CryptoHelper, EventBus } from '../../base'; +import { EntitlementService } from '../../core/entitlement'; +import { WorkspacePolicyService } from '../../core/permission'; +import { QuotaStateService } from '../../core/quota/state'; +import { WorkspaceService } from '../../core/workspaces'; +import { Models } from '../../models'; +import { LicenseService } from '../../plugins/license/service'; +import { PaymentEventHandlers } from '../../plugins/payment/event'; +import { + SubscriptionPlan, + SubscriptionRecurring, + SubscriptionVariant, +} from '../../plugins/payment/types'; + +type Context = Record; + +const test = ava as TestFn; + +test('workspace subscription activation only sends upgrade notification', async t => { + const events: Array<{ name: string; payload: unknown }> = []; + let reconciled = false; + const handler = new PaymentEventHandlers( + { + isTeamWorkspace: async () => true, + sendTeamWorkspaceUpgradedEmail: async () => {}, + } as unknown as WorkspaceService, + { + reconcileWorkspaceQuotaState: async () => { + reconciled = true; + }, + } as unknown as WorkspacePolicyService, + { + reconcileWorkspaceQuotaState: async () => ({ seatLimit: 7 }), + } as unknown as QuotaStateService, + { + emit: (name: string, payload: unknown) => events.push({ name, payload }), + } as unknown as EventBus + ); + + await handler.onWorkspaceSubscriptionUpdated({ + workspaceId: 'ws', + plan: SubscriptionPlan.Team, + recurring: SubscriptionRecurring.Yearly, + quantity: 999, + }); + + t.deepEqual(events, []); + t.false(reconciled); +}); + +test('workspace entitlement change allocates seats from effective quota state', async t => { + const events: Array<{ name: string; payload: unknown }> = []; + const handler = new PaymentEventHandlers( + {} as unknown as WorkspaceService, + {} as unknown as WorkspacePolicyService, + { + reconcileWorkspaceQuotaState: async () => ({ + plan: 'team', + seatLimit: 7, + }), + } as unknown as QuotaStateService, + { + emit: (name: string, payload: unknown) => events.push({ name, payload }), + } as unknown as EventBus + ); + + await handler.onEntitlementChanged({ + targetType: 'workspace', + targetId: 'ws', + }); + + t.deepEqual(events, [ + { + name: 'workspace.members.allocateSeats', + payload: { workspaceId: 'ws', quantity: 7 }, + }, + ]); +}); + +test('onetime selfhost license seat allocation ignores projected license quantity', async t => { + const events: Array<{ name: string; payload: unknown }> = []; + const service = new LicenseService( + { + installedLicense: { + findUnique: async () => ({ + key: 'license-key', + workspaceId: 'ws', + quantity: 999, + recurring: SubscriptionRecurring.Yearly, + variant: SubscriptionVariant.Onetime, + }), + }, + } as unknown as PrismaClient, + { + emit: (name: string, payload: unknown) => events.push({ name, payload }), + } as unknown as EventBus, + {} as unknown as Models, + {} as unknown as CryptoHelper, + {} as unknown as WorkspacePolicyService, + {} as unknown as EntitlementService, + { + reconcileWorkspaceQuotaState: async () => ({ seatLimit: 4 }), + } as unknown as QuotaStateService + ); + + await service.updateTeamSeats({ + workspaceId: 'ws', + } as Events['workspace.members.updated']); + + t.deepEqual(events, [ + { + name: 'workspace.members.allocateSeats', + payload: { workspaceId: 'ws', quantity: 4 }, + }, + ]); +}); + +test('recurring selfhost license activation returns activation projection without remote health recheck', async t => { + const events: Array<{ name: string; payload: unknown }> = []; + const affineProRequests: string[] = []; + const upserts: unknown[] = []; + const entitlements: unknown[] = []; + const expiresAt = Date.now() + 30 * 24 * 60 * 60 * 1000; + const service = new LicenseService( + { + installedLicense: { + findUnique: async () => null, + upsert: async (input: unknown) => { + upserts.push(input); + return { + workspaceId: 'ws', + key: 'license-key', + quantity: 3, + recurring: SubscriptionRecurring.Monthly, + variant: null, + }; + }, + }, + } as unknown as PrismaClient, + { + emit: (name: string, payload: unknown) => events.push({ name, payload }), + } as unknown as EventBus, + {} as unknown as Models, + {} as unknown as CryptoHelper, + {} as unknown as WorkspacePolicyService, + { + upsertFromValidatedSelfhostLicense: async (input: unknown) => { + entitlements.push(input); + }, + } as unknown as EntitlementService, + {} as unknown as QuotaStateService + ); + + ( + service as unknown as { + fetchAffinePro: (path: string) => Promise<{ + plan: SubscriptionPlan; + recurring: SubscriptionRecurring; + quantity: number; + endAt: number; + res: Response; + }>; + } + ).fetchAffinePro = async (path: string) => { + affineProRequests.push(path); + return { + plan: SubscriptionPlan.SelfHostedTeam, + recurring: SubscriptionRecurring.Monthly, + quantity: 3, + endAt: expiresAt, + res: new Response(null, { + headers: { + 'x-next-validate-key': 'next-validate-key', + }, + }), + }; + }; + + const license = await service.activateTeamLicense('ws', 'license-key'); + + t.like(license, { + workspaceId: 'ws', + key: 'license-key', + quantity: 3, + recurring: SubscriptionRecurring.Monthly, + }); + t.is(entitlements.length, 1); + t.is(upserts.length, 1); + t.deepEqual(affineProRequests, ['/api/team/licenses/license-key/activate']); + t.deepEqual(events, [ + { + name: 'workspace.subscription.activated', + payload: { + workspaceId: 'ws', + plan: SubscriptionPlan.SelfHostedTeam, + recurring: SubscriptionRecurring.Monthly, + quantity: 3, + }, + }, + ]); +}); diff --git a/packages/backend/server/src/__tests__/storage/blob-upload-cleanup.spec.ts b/packages/backend/server/src/__tests__/storage/blob-upload-cleanup.spec.ts index 7fb71e9a9e..5bb790c373 100644 --- a/packages/backend/server/src/__tests__/storage/blob-upload-cleanup.spec.ts +++ b/packages/backend/server/src/__tests__/storage/blob-upload-cleanup.spec.ts @@ -86,7 +86,10 @@ test('should cleanup expired pending blobs', async t => { ], }); - const abortSpy = Sinon.spy(t.context.storage, 'abortMultipartUpload'); + const abortSpy = Sinon.stub( + t.context.storage, + 'abortMultipartUpload' + ).resolves(); const deleteSpy = Sinon.spy(t.context.storage, 'delete'); t.teardown(() => { abortSpy.restore(); diff --git a/packages/backend/server/src/__tests__/worker.e2e.ts b/packages/backend/server/src/__tests__/worker.e2e.ts index b847af0421..0b20652740 100644 --- a/packages/backend/server/src/__tests__/worker.e2e.ts +++ b/packages/backend/server/src/__tests__/worker.e2e.ts @@ -9,7 +9,7 @@ import type { TestingApp } from './utils'; type TestContext = { app: TestingApp; }; -const test = ava as TestFn; +const test = ava.serial as TestFn; let safeFetchStub: Sinon.SinonStub | undefined; let safeFetchHandler: diff --git a/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts b/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts index 4698413623..93d3035cc4 100644 --- a/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts +++ b/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts @@ -3,7 +3,8 @@ import { createHash } from 'node:crypto'; import test from 'ava'; import Sinon from 'sinon'; -import { Config, StorageProviderFactory } from '../../base'; +import { Config, ConfigFactory, StorageProviderFactory } from '../../base'; +import { QuotaStateService } from '../../core/quota/state'; import { WorkspaceBlobStorage } from '../../core/storage/wrappers/blob'; import { BlobModel, WorkspaceFeatureModel } from '../../models'; import { @@ -35,6 +36,18 @@ let model: WorkspaceFeatureModel; test.before(async () => { app = await createTestingApp(); model = app.get(WorkspaceFeatureModel); + app.get(ConfigFactory).override({ + storages: { + blob: { + storage: { + provider: 'fs', + bucket: 'test', + config: { path: '/tmp/affine-test-storage' }, + }, + }, + }, + }); + await app.get(WorkspaceBlobStorage).onConfigInit(); }); test.beforeEach(async () => { @@ -45,6 +58,26 @@ test.after.always(async () => { await app.close(); }); +async function withRestrictedWorkspaceQuota(workspaceId: string) { + const quotaState = app.get(QuotaStateService); + const blobModel = app.get(BlobModel); + const base = await quotaState.reconcileWorkspaceQuotaState(workspaceId); + return Sinon.stub(quotaState, 'reconcileWorkspaceQuotaState').callsFake( + async id => { + if (id !== workspaceId) { + return base; + } + + return { + ...base, + blobLimit: BigInt(RESTRICTED_QUOTA.blobLimit), + storageQuota: BigInt(RESTRICTED_QUOTA.storageQuota), + usedStorageQuota: BigInt(await blobModel.totalSize(workspaceId)), + }; + } + ); +} + test('should set blobs', async t => { await app.signupV1('u1@affine.pro'); @@ -233,7 +266,8 @@ test('should reject blob exceeded limit', async t => { await app.signupV1('u1@affine.pro'); const workspace1 = await createWorkspace(app); - await model.add(workspace1.id, 'team_plan_v1', 'test', RESTRICTED_QUOTA); + const quotaStub = await withRestrictedWorkspaceQuota(workspace1.id); + t.teardown(() => quotaStub.restore()); const buffer1 = Buffer.from( Array.from({ length: RESTRICTED_QUOTA.blobLimit + 1 }, () => 0) @@ -247,7 +281,8 @@ test('should reject blob exceeded storage quota', async t => { await app.signupV1('u1@affine.pro'); const workspace = await createWorkspace(app); - await model.add(workspace.id, 'team_plan_v1', 'test', RESTRICTED_QUOTA); + const quotaStub = await withRestrictedWorkspaceQuota(workspace.id); + t.teardown(() => quotaStub.restore()); const buffer = Buffer.from(Array.from({ length: OneMB }, () => 0)); diff --git a/packages/backend/server/src/__tests__/workspace/controller.spec.ts b/packages/backend/server/src/__tests__/workspace/controller.spec.ts index 9cc28456b6..c336b4fdf8 100644 --- a/packages/backend/server/src/__tests__/workspace/controller.spec.ts +++ b/packages/backend/server/src/__tests__/workspace/controller.spec.ts @@ -7,7 +7,9 @@ import Sinon from 'sinon'; import supertest from 'supertest'; import { applyUpdate, Doc as YDoc, Map as YMap } from 'yjs'; +import { ConfigFactory } from '../../base'; import { PgWorkspaceDocStorageAdapter } from '../../core/doc'; +import { PermissionReadModel } from '../../core/permission/config'; import { WorkspaceBlobStorage } from '../../core/storage'; import { Models, PublicDocMode, WorkspaceRole } from '../../models'; import { @@ -152,6 +154,31 @@ test('should be able to get private workspace with public pages', async t => { t.is(res.text, 'blob'); }); +test('should be able to get private workspace with public pages using new permission model', async t => { + const { app, storage } = t.context; + const config = app.get(ConfigFactory); + + config.override({ + permission: { + readModel: PermissionReadModel.Projection, + }, + }); + try { + storage.get.resolves(blob()); + const res = await app.GET('/api/workspaces/private/blobs/test'); + + t.is(res.status, HttpStatus.OK); + t.is(res.get('content-type'), 'text/plain'); + t.is(res.text, 'blob'); + } finally { + config.override({ + permission: { + readModel: PermissionReadModel.Legacy, + }, + }); + } +}); + test('should not be able to get private workspace with no public pages', async t => { const { app } = t.context; diff --git a/packages/backend/server/src/base/metrics/metrics.ts b/packages/backend/server/src/base/metrics/metrics.ts index a830c208a2..59858754da 100644 --- a/packages/backend/server/src/base/metrics/metrics.ts +++ b/packages/backend/server/src/base/metrics/metrics.ts @@ -62,6 +62,7 @@ export type KnownMetricScopes = | 'queue' | 'storage' | 'process' + | 'permission' | 'workspace'; const metricCreators: MetricCreators = { diff --git a/packages/backend/server/src/core/comment/realtime.ts b/packages/backend/server/src/core/comment/realtime.ts index 49d67b42f9..901b3c890a 100644 --- a/packages/backend/server/src/core/comment/realtime.ts +++ b/packages/backend/server/src/core/comment/realtime.ts @@ -2,7 +2,7 @@ import { Injectable, OnModuleInit } from '@nestjs/common'; import { z } from 'zod'; import { decodeWithJson, encodeWithJson } from '../../base/graphql'; -import { AccessController } from '../permission'; +import { PermissionAccess } from '../permission'; import { realtimeCommentRoom, RealtimePublisher, @@ -20,7 +20,7 @@ export function commentRoom(workspaceId: string, docId: string) { export class CommentRealtimeProvider implements OnModuleInit { constructor( private readonly service: CommentService, - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly registry: RealtimeRegistry ) {} diff --git a/packages/backend/server/src/core/comment/resolver.ts b/packages/backend/server/src/core/comment/resolver.ts index ba5dc547a2..1bb0f35313 100644 --- a/packages/backend/server/src/core/comment/resolver.ts +++ b/packages/backend/server/src/core/comment/resolver.ts @@ -25,7 +25,7 @@ import { import { Comment, DocMode, Models, Reply } from '../../models'; import { CurrentUser } from '../auth/session'; import { ServerFeature, ServerService } from '../config'; -import { AccessController, DocAction } from '../permission'; +import { DocAction, PermissionAccess } from '../permission'; import { RealtimePublisher } from '../realtime'; import { CommentAttachmentStorage } from '../storage'; import { UserType } from '../user'; @@ -54,7 +54,7 @@ export interface CommentCursor { export class CommentResolver { constructor( private readonly service: CommentService, - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly commentAttachmentStorage: CommentAttachmentStorage, private readonly queue: JobQueue, private readonly models: Models, @@ -469,11 +469,7 @@ export class CommentResolver { private async assertPermission( me: UserType, - item: { - workspaceId: string; - docId: string; - userId?: string; - }, + item: { workspaceId: string; docId: string; userId?: string }, action: DocAction ) { // the owner of the comment/reply can update, delete, resolve it diff --git a/packages/backend/server/src/core/config/resolver.ts b/packages/backend/server/src/core/config/resolver.ts index 2c989615a2..53fba591d1 100644 --- a/packages/backend/server/src/core/config/resolver.ts +++ b/packages/backend/server/src/core/config/resolver.ts @@ -173,7 +173,7 @@ export class ServerFeatureConfigResolver extends AvailableUserFeatureConfig { description: 'Workspace features available for admin configuration', }) availableWorkspaceFeatures(): WorkspaceFeatureName[] { - return ['unlimited_workspace', 'team_plan_v1']; + return []; } } diff --git a/packages/backend/server/src/core/doc-renderer/controller.ts b/packages/backend/server/src/core/doc-renderer/controller.ts index 69f0da2bac..fca8bd3441 100644 --- a/packages/backend/server/src/core/doc-renderer/controller.ts +++ b/packages/backend/server/src/core/doc-renderer/controller.ts @@ -11,7 +11,7 @@ import { Models } from '../../models'; import { htmlSanitize } from '../../native'; import { Public } from '../auth'; import { DocReader } from '../doc'; -import { WorkspacePolicyService } from '../permission'; +import { PermissionService } from '../permission'; interface RenderOptions { title: string; @@ -61,7 +61,7 @@ export class DocRendererController { private readonly doc: DocReader, private readonly models: Models, private readonly config: Config, - private readonly policy: WorkspacePolicyService + private readonly permission: PermissionService ) { this.webAssets = this.readHtmlAssets(join(env.projectRoot, 'static')); this.mobileAssets = this.readHtmlAssets( @@ -99,10 +99,11 @@ export class DocRendererController { req.accepts().some(t => markdownType.has(t.toLowerCase())) ) { try { - const canReadMarkdown = await this.policy.canReadSharedDoc( + const canReadMarkdown = await this.permission.canDoc({ workspaceId, - sub - ); + docId: sub, + action: 'Doc.Read', + }); if (!canReadMarkdown) { res.status(404).end(); return; @@ -162,7 +163,7 @@ export class DocRendererController { workspaceId: string, docId: string ): Promise { - if (await this.policy.canPreviewDoc(workspaceId, docId)) { + if (await this.permission.canPreviewDoc({ workspaceId, docId })) { return this.doc.getDocContent(workspaceId, docId); } @@ -172,8 +173,9 @@ export class DocRendererController { private async getWorkspaceContent( workspaceId: string ): Promise { - const canPreviewWorkspace = - await this.policy.canPreviewWorkspace(workspaceId); + const canPreviewWorkspace = await this.permission.canPreviewWorkspace({ + workspaceId, + }); if (!canPreviewWorkspace) return null; const workspaceContent = await this.doc.getWorkspaceContent(workspaceId); diff --git a/packages/backend/server/src/core/doc/options.ts b/packages/backend/server/src/core/doc/options.ts index f4377c04a3..3fc0fcc3bb 100644 --- a/packages/backend/server/src/core/doc/options.ts +++ b/packages/backend/server/src/core/doc/options.ts @@ -73,7 +73,7 @@ export class DocStorageOptions implements IDocStorageOptions { historyMaxAge = async (spaceId: string) => { const quota = await this.quota.getWorkspaceQuota(spaceId); - return quota.historyPeriod; + return quota.historyPeriod * 1000; }; historyMinInterval = (_spaceId: string) => { diff --git a/packages/backend/server/src/core/entitlement/__tests__/projection-checker.spec.ts b/packages/backend/server/src/core/entitlement/__tests__/projection-checker.spec.ts new file mode 100644 index 0000000000..a717a118a2 --- /dev/null +++ b/packages/backend/server/src/core/entitlement/__tests__/projection-checker.spec.ts @@ -0,0 +1,135 @@ +import { randomUUID } from 'node:crypto'; + +import { PrismaClient } from '@prisma/client'; +import ava, { TestFn } from 'ava'; + +import { + createTestingModule, + type TestingModule, +} from '../../../__tests__/utils'; +import { Models } from '../../../models'; +import { + SubscriptionPlan, + SubscriptionRecurring, + SubscriptionStatus, +} from '../../../plugins/payment/types'; +import { + EntitlementModule, + EntitlementProjectionChecker, + EntitlementService, +} from '../index'; + +interface Context { + module: TestingModule; + db: PrismaClient; + models: Models; + entitlement: EntitlementService; + checker: EntitlementProjectionChecker; +} + +const test = ava as TestFn; + +test.before(async t => { + const module = await createTestingModule({ imports: [EntitlementModule] }); + t.context.module = module; + t.context.db = module.get(PrismaClient); + t.context.models = module.get(Models); + t.context.entitlement = module.get(EntitlementService); + t.context.checker = module.get(EntitlementProjectionChecker); +}); + +test.beforeEach(async t => { + await t.context.module.initTestingDB(); +}); + +test.after.always(async t => { + await t.context.module.close(); +}); + +test('checker distinguishes valid projection from dirty legacy features', async t => { + const cleanUser = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: cleanUser.id, + plan: 'pro', + recurring: SubscriptionRecurring.Monthly, + status: 'active', + }); + + const dirtyUser = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + await t.context.models.userFeature.add( + dirtyUser.id, + 'pro_plan_v1', + 'dirty legacy feature' + ); + + const report = await t.context.checker.checkEntitlementProjection(); + + t.is(report.dirtyLegacyUserFeatures, 1); + t.is(report.missingUserFeatureProjection, 0); +}); + +test('checker reports missing legacy projection and stale state', async t => { + const user = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: user.id, + plan: 'pro', + recurring: SubscriptionRecurring.Monthly, + status: 'active', + }); + await t.context.db.subscription.delete({ + where: { targetId_plan: { targetId: user.id, plan: 'pro' } }, + }); + await t.context.db.effectiveUserQuotaState.update({ + where: { userId: user.id }, + data: { + staleAfter: new Date('2020-01-01T00:00:00Z'), + }, + }); + + const report = await t.context.checker.checkEntitlementProjection(); + + t.is(report.cloudSubscriptionProjectionMissing, 1); + t.is(report.staleEffectiveUserState, 1); +}); + +test('checker reports legal legacy facts missing entitlements', async t => { + const user = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + await t.context.db.subscription.create({ + data: { + targetId: user.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Monthly, + status: SubscriptionStatus.Active, + start: new Date(), + }, + }); + + const owner = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + const workspace = await t.context.models.workspace.create(owner.id); + await t.context.db.installedLicense.create({ + data: { + key: 'legacy-verifiable-key', + workspaceId: workspace.id, + quantity: 5, + recurring: SubscriptionRecurring.Yearly, + validateKey: 'validate-key', + validatedAt: new Date(), + license: Buffer.from('raw-license'), + }, + }); + + const report = await t.context.checker.checkEntitlementProjection(); + + t.is(report.cloudSubscriptionEntitlementMissing, 1); + t.is(report.selfhostLicenseEntitlementMissing, 1); +}); diff --git a/packages/backend/server/src/core/entitlement/__tests__/projection.spec.ts b/packages/backend/server/src/core/entitlement/__tests__/projection.spec.ts new file mode 100644 index 0000000000..6b424b1506 --- /dev/null +++ b/packages/backend/server/src/core/entitlement/__tests__/projection.spec.ts @@ -0,0 +1,434 @@ +import { randomUUID } from 'node:crypto'; + +import { PrismaClient } from '@prisma/client'; +import ava, { TestFn } from 'ava'; + +import { + createTestingModule, + type TestingModule, +} from '../../../__tests__/utils'; +import { Models } from '../../../models'; +import { + SubscriptionPlan, + SubscriptionRecurring, + SubscriptionStatus, +} from '../../../plugins/payment/types'; +import { EntitlementModule, EntitlementService } from '../index'; +import { LegacyEntitlementProjectionService } from '../projection'; + +interface Context { + module: TestingModule; + db: PrismaClient; + models: Models; + entitlement: EntitlementService; + projection: LegacyEntitlementProjectionService; +} + +const test = ava as TestFn; + +test.before(async t => { + const module = await createTestingModule({ imports: [EntitlementModule] }); + t.context.module = module; + t.context.db = module.get(PrismaClient); + t.context.models = module.get(Models); + t.context.entitlement = module.get(EntitlementService); + t.context.projection = module.get(LegacyEntitlementProjectionService); +}); + +test.beforeEach(async t => { + await t.context.module.initTestingDB(); +}); + +test.after.always(async t => { + await t.context.module.close(); +}); + +test('projects user entitlement to legacy user features and subscriptions', async t => { + const user = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Yearly, + status: 'active', + }); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.AI, + recurring: SubscriptionRecurring.Monthly, + status: 'active', + }); + + t.true(await t.context.models.userFeature.has(user.id, 'pro_plan_v1')); + t.true(await t.context.models.userFeature.has(user.id, 'unlimited_copilot')); + t.like( + await t.context.db.subscription.findUnique({ + where: { + targetId_plan: { targetId: user.id, plan: SubscriptionPlan.Pro }, + }, + }), + { + recurring: SubscriptionRecurring.Yearly, + status: 'active', + } + ); + + await t.context.entitlement.revokeCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.AI, + }); + t.false(await t.context.models.userFeature.has(user.id, 'unlimited_copilot')); +}); + +test('projects workspace entitlement and readonly state to legacy workspace features', async t => { + const owner = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + const workspace = await t.context.models.workspace.create(owner.id); + + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: workspace.id, + plan: SubscriptionPlan.Team, + recurring: SubscriptionRecurring.Yearly, + status: 'active', + quantity: 8, + }); + + const teamFeature = await t.context.models.workspaceFeature.get( + workspace.id, + 'team_plan_v1' + ); + t.is(teamFeature?.configs.memberLimit, 8); + + await t.context.db.effectiveWorkspaceQuotaState.upsert({ + where: { + workspaceId: workspace.id, + }, + create: { + workspaceId: workspace.id, + plan: 'free', + ownerUserId: owner.id, + usesOwnerQuota: true, + seatLimit: 3, + memberCount: 4, + overcapacityMemberCount: 1, + blobLimit: BigInt(10), + storageQuota: BigInt(10), + usedStorageQuota: BigInt(1), + historyPeriodSeconds: 7, + readonly: true, + readonlyReasons: ['member_overflow'], + known: true, + stale: false, + }, + update: { + plan: 'free', + ownerUserId: owner.id, + usesOwnerQuota: true, + seatLimit: 3, + memberCount: 4, + overcapacityMemberCount: 1, + blobLimit: BigInt(10), + storageQuota: BigInt(10), + usedStorageQuota: BigInt(1), + historyPeriodSeconds: 7, + readonly: true, + readonlyReasons: ['member_overflow'], + known: true, + stale: false, + }, + }); + await t.context.projection.onWorkspaceQuotaStateChanged({ + workspaceId: workspace.id, + }); + + t.true( + await t.context.models.workspaceFeature.has( + workspace.id, + 'quota_exceeded_readonly_workspace_v1' + ) + ); +}); + +test('installed license scanner never trusts quantity without raw license', async t => { + const owner = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + const workspace = await t.context.models.workspace.create(owner.id); + + await t.context.db.installedLicense.create({ + data: { + key: 'legacy-key', + workspaceId: workspace.id, + quantity: 100, + recurring: SubscriptionRecurring.Yearly, + validateKey: '', + validatedAt: new Date(), + }, + }); + + await t.context.projection.scanInstalledLicenses(); + + const entitlement = await t.context.db.entitlement.findFirst({ + where: { + source: 'selfhost_license', + subjectId: 'legacy-key', + }, + }); + t.is(entitlement?.status, 'needs_reupload'); + t.is(entitlement?.quantity, null); +}); + +test.serial( + 'selfhosted legacy projection ignores unknown entitlements', + async t => { + const previousDeploymentType = globalThis.env.DEPLOYMENT_TYPE; + // @ts-expect-error test mutates env singleton for deployment-specific projection semantics + globalThis.env.DEPLOYMENT_TYPE = 'selfhosted'; + try { + const user = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + await t.context.db.entitlement.create({ + data: { + targetType: 'user', + targetId: user.id, + source: 'cloud_subscription', + plan: 'ai', + status: 'active', + subjectId: `forged-ai:${user.id}`, + }, + }); + + await t.context.projection.onEntitlementChanged({ + targetType: 'user', + targetId: user.id, + }); + + t.false( + await t.context.models.userFeature.has(user.id, 'unlimited_copilot') + ); + t.is( + await t.context.db.subscription.count({ where: { targetId: user.id } }), + 0 + ); + } finally { + // @ts-expect-error restore mutable test env singleton + globalThis.env.DEPLOYMENT_TYPE = previousDeploymentType; + } + } +); + +test('backfill marks selfhost team subscriptions as needing license revalidation', async t => { + await t.context.db.subscription.create({ + data: { + targetId: 'license-key-target', + plan: SubscriptionPlan.SelfHostedTeam, + recurring: SubscriptionRecurring.Yearly, + status: SubscriptionStatus.Active, + start: new Date(), + }, + }); + + await t.context.projection.backfillEntitlementsAndQuotaStates(); + + t.like( + await t.context.db.entitlement.findFirstOrThrow({ + where: { + source: 'selfhost_license', + subjectId: 'license-key-target', + }, + }), + { + targetType: 'instance', + targetId: 'license-key-target', + plan: 'selfhost_team', + status: 'needs_reupload', + } + ); +}); + +test('key based selfhost entitlements without raw payload need reupload', async t => { + const owner = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + const workspace = await t.context.models.workspace.create(owner.id); + + await t.context.entitlement.upsertFromSelfhostLicense({ + workspaceId: workspace.id, + licenseKey: 'remote-key', + recurring: SubscriptionRecurring.Yearly, + quantity: 5, + validateKey: 'validate-key', + expiresAt: new Date(Date.now() + 3600_000), + }); + + await t.context.projection.scanInstalledLicenses(); + + t.like( + await t.context.db.entitlement.findFirstOrThrow({ + where: { source: 'selfhost_license', subjectId: 'remote-key' }, + }), + { status: 'needs_reupload', quantity: null } + ); +}); + +test('revoked selfhost entitlement removes installed license projection', async t => { + const owner = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + const workspace = await t.context.models.workspace.create(owner.id); + + await t.context.db.entitlement.create({ + data: { + targetType: 'workspace', + targetId: workspace.id, + source: 'selfhost_license', + plan: 'selfhost_team', + status: 'active', + subjectId: 'revoked-key', + quantity: 5, + signedPayload: Buffer.from('signed-license-payload'), + metadata: { + recurring: SubscriptionRecurring.Yearly, + validateKey: 'validate-key', + }, + expiresAt: new Date(Date.now() + 3600_000), + validatedAt: new Date(), + }, + }); + await t.context.db.installedLicense.create({ + data: { + key: 'revoked-key', + workspaceId: workspace.id, + quantity: 5, + recurring: SubscriptionRecurring.Yearly, + validateKey: 'validate-key', + validatedAt: new Date(), + license: Buffer.from('signed-license-payload'), + }, + }); + + await t.context.entitlement.revokeBySubject( + 'selfhost_license', + 'revoked-key' + ); + + t.falsy( + await t.context.db.installedLicense.findUnique({ + where: { workspaceId: workspace.id }, + }) + ); +}); + +test('installed license projection uses explicit entitlement status priority', async t => { + const owner = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + const workspace = await t.context.models.workspace.create(owner.id); + + await t.context.db.entitlement.createMany({ + data: [ + { + targetType: 'workspace', + targetId: workspace.id, + source: 'selfhost_license', + plan: 'selfhost_team', + status: 'expired', + subjectId: 'expired-key', + quantity: 5, + metadata: { + recurring: SubscriptionRecurring.Yearly, + validateKey: 'expired-validate-key', + }, + expiresAt: new Date(Date.now() - 3600_000), + validatedAt: new Date(), + }, + { + targetType: 'workspace', + targetId: workspace.id, + source: 'selfhost_license', + plan: 'selfhost_team', + status: 'grace', + subjectId: 'grace-key', + quantity: 6, + metadata: { + recurring: SubscriptionRecurring.Yearly, + validateKey: 'grace-validate-key', + }, + expiresAt: new Date(Date.now() - 1800_000), + graceUntil: new Date(Date.now() + 3600_000), + validatedAt: new Date(), + }, + ], + }); + + await t.context.projection.onEntitlementChanged({ + targetType: 'workspace', + targetId: workspace.id, + }); + + const installedLicense = + await t.context.db.installedLicense.findUniqueOrThrow({ + where: { workspaceId: workspace.id }, + }); + t.is(installedLicense.key, 'grace-key'); + t.is(installedLicense.quantity, 6); + t.is(installedLicense.validateKey, 'grace-validate-key'); +}); + +test.serial( + 'selfhosted projection does not trust non-null signed payload', + async t => { + const previousDeploymentType = globalThis.env.DEPLOYMENT_TYPE; + // @ts-expect-error test mutates env singleton for deployment-specific projection semantics + globalThis.env.DEPLOYMENT_TYPE = 'selfhosted'; + try { + const owner = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + const workspace = await t.context.models.workspace.create(owner.id); + + await t.context.db.entitlement.create({ + data: { + targetType: 'workspace', + targetId: workspace.id, + source: 'selfhost_license', + plan: 'selfhost_team', + status: 'active', + subjectId: 'forged-key', + quantity: 100, + signedPayload: Buffer.from('not-a-valid-license'), + metadata: { + recurring: SubscriptionRecurring.Yearly, + validateKey: 'validate-key', + }, + expiresAt: new Date(Date.now() + 3600_000), + validatedAt: new Date(), + }, + }); + + await t.context.projection.onEntitlementChanged({ + targetType: 'workspace', + targetId: workspace.id, + }); + + t.falsy( + await t.context.models.workspaceFeature.get( + workspace.id, + 'team_plan_v1' + ) + ); + t.falsy( + await t.context.db.installedLicense.findUnique({ + where: { workspaceId: workspace.id }, + }) + ); + } finally { + // @ts-expect-error restore mutable test env singleton + globalThis.env.DEPLOYMENT_TYPE = previousDeploymentType; + } + } +); diff --git a/packages/backend/server/src/core/entitlement/__tests__/service.spec.ts b/packages/backend/server/src/core/entitlement/__tests__/service.spec.ts new file mode 100644 index 0000000000..8214719df1 --- /dev/null +++ b/packages/backend/server/src/core/entitlement/__tests__/service.spec.ts @@ -0,0 +1,508 @@ +import { PrismaClient } from '@prisma/client'; +import ava, { TestFn } from 'ava'; + +import { + createTestingModule, + type TestingModule, +} from '../../../__tests__/utils'; +import { Models } from '../../../models'; +import { + SubscriptionPlan, + SubscriptionRecurring, + SubscriptionStatus, +} from '../../../plugins/payment/types'; +import { EntitlementModule } from '../index'; +import { EntitlementService } from '../service'; + +interface Context { + module: TestingModule; + db: PrismaClient; + models: Models; + service: EntitlementService; +} + +const test = ava as TestFn; + +test.before(async t => { + const module = await createTestingModule({ imports: [EntitlementModule] }); + t.context.module = module; + t.context.db = module.get(PrismaClient); + t.context.models = module.get(Models); + t.context.service = module.get(EntitlementService); +}); + +test.beforeEach(async t => { + await t.context.module.initTestingDB(); +}); + +test.after.always(async t => { + await t.context.module.close(); +}); + +test('upserts admin grant entitlement as commercial source of truth', async t => { + const owner = await t.context.models.user.create({ + email: 'admin-grant-owner@affine.pro', + }); + const workspace = await t.context.models.workspace.create(owner.id); + + const entitlement = await t.context.service.upsertAdminGrant({ + targetType: 'workspace', + targetId: workspace.id, + plan: 'team', + quantity: 6, + }); + const resolved = await t.context.service.resolveWorkspaceEntitlement( + workspace.id + ); + + t.is(entitlement.source, 'admin_grant'); + t.is(entitlement.plan, 'team'); + t.is(entitlement.quantity, 6); + t.is(resolved.plan, 'team'); + t.is(resolved.quota.seatLimit, 6); +}); + +test('admin grant replaces and revokes previous admin grant', async t => { + const user = await t.context.models.user.create({ + email: 'admin-grant-replace@affine.pro', + }); + + await t.context.service.upsertAdminGrant({ + targetType: 'user', + targetId: user.id, + plan: 'lifetime_pro', + }); + await t.context.service.upsertAdminGrant({ + targetType: 'user', + targetId: user.id, + plan: 'pro', + }); + + const [resolved, entitlements] = await Promise.all([ + t.context.service.resolveUserEntitlement(user.id), + t.context.db.entitlement.findMany({ + where: { source: 'admin_grant', targetId: user.id }, + }), + ]); + + t.is(resolved.plan, 'pro'); + t.is( + entitlements.filter(entitlement => entitlement.status === 'active').length, + 1 + ); + t.false( + entitlements.some( + entitlement => + entitlement.plan === 'lifetime_pro' && entitlement.status === 'active' + ) + ); + + await t.context.service.revokeAdminGrant('user', user.id); + t.is((await t.context.service.resolveUserEntitlement(user.id)).plan, 'free'); +}); + +test('admin grant rejects self-hosted commercial entitlement without writing', async t => { + const originalDeploymentType = globalThis.env.DEPLOYMENT_TYPE; + // @ts-expect-error test mutates env singleton for deployment-specific entitlement semantics + globalThis.env.DEPLOYMENT_TYPE = 'selfhosted'; + const owner = await t.context.models.user.create({ + email: 'admin-grant-selfhost@affine.pro', + }); + const workspace = await t.context.models.workspace.create(owner.id); + + try { + await t.throwsAsync( + t.context.service.upsertAdminGrant({ + targetType: 'workspace', + targetId: workspace.id, + plan: 'team', + quantity: 6, + }), + { message: /signed license/ } + ); + t.is( + await t.context.db.entitlement.count({ + where: { source: 'admin_grant', targetId: workspace.id }, + }), + 0 + ); + } finally { + // @ts-expect-error restore mutable test env singleton + globalThis.env.DEPLOYMENT_TYPE = originalDeploymentType; + } +}); + +test('admin grant rejects incompatible target plan without writing', async t => { + const user = await t.context.models.user.create({ + email: 'admin-grant-invalid@affine.pro', + }); + + await t.context.service.upsertAdminGrant({ + targetType: 'user', + targetId: user.id, + plan: 'pro', + }); + await t.throwsAsync( + t.context.service.upsertAdminGrant({ + targetType: 'user', + targetId: user.id, + plan: 'team', + quantity: 6, + }), + { message: /not configurable/ } + ); + + const active = await t.context.db.entitlement.findMany({ + where: { source: 'admin_grant', targetId: user.id, status: 'active' }, + }); + t.is(active.length, 1); + t.is(active[0].plan, 'pro'); +}); + +test('upserts cloud subscription entitlements without writing legacy features', async t => { + const proUser = await t.context.models.user.create({ + email: 'user-pro@affine.pro', + }); + const aiUser = await t.context.models.user.create({ + email: 'user-ai@affine.pro', + }); + const owner = await t.context.models.user.create({ + email: 'workspace-owner@affine.pro', + }); + const teamWorkspace = await t.context.models.workspace.create(owner.id); + const cases = [ + { + targetId: proUser.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Yearly, + status: 'active', + expected: { targetType: 'user', plan: 'pro', status: 'active' }, + }, + { + targetId: aiUser.id, + plan: SubscriptionPlan.AI, + recurring: SubscriptionRecurring.Monthly, + status: 'trialing', + expected: { targetType: 'user', plan: 'ai', status: 'active' }, + }, + { + targetId: teamWorkspace.id, + plan: SubscriptionPlan.Team, + recurring: SubscriptionRecurring.Yearly, + status: 'past_due', + quantity: 7, + expected: { targetType: 'workspace', plan: 'team', status: 'grace' }, + }, + ]; + + for (const item of cases) { + const entitlement = await t.context.service.upsertFromCloudSubscription({ + ...item, + subscriptionId: `${item.targetId}:${item.plan}`, + start: new Date('2026-05-14T00:00:00Z'), + }); + + t.like(entitlement, item.expected, item.targetId); + } + + t.is(await t.context.db.entitlement.count(), cases.length); +}); + +test('revokes cloud subscription entitlement by subject', async t => { + const user = await t.context.models.user.create({ + email: 'revoke-user@affine.pro', + }); + const entitlement = await t.context.service.upsertFromCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Monthly, + status: 'active', + subscriptionId: 'sub_1', + }); + + await t.context.service.revokeCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.Pro, + subscriptionId: 'sub_1', + }); + + const updated = await t.context.db.entitlement.findUnique({ + where: { id: entitlement.id }, + }); + t.is(updated?.status, 'revoked'); +}); + +test('revokes onetime or revenuecat entitlements using fallback subject', async t => { + const user = await t.context.models.user.create({ + email: 'fallback-user@affine.pro', + }); + const entitlement = await t.context.service.upsertFromCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Yearly, + status: 'active', + }); + + await t.context.service.revokeCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.Pro, + subscriptionId: 1, + }); + + const updated = await t.context.db.entitlement.findUnique({ + where: { id: entitlement.id }, + }); + t.is(updated?.status, 'revoked'); +}); + +test('resolves higher priority commercial entitlement over ai capability', async t => { + const user = await t.context.models.user.create({ + email: 'priority-user@affine.pro', + }); + await t.context.service.upsertFromCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Yearly, + status: 'active', + }); + await t.context.service.upsertFromCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.AI, + recurring: SubscriptionRecurring.Monthly, + status: 'active', + }); + + const resolved = await t.context.service.resolveUserEntitlement(user.id); + t.is(resolved.plan, 'pro'); + t.is(resolved.quota.storageQuota, 100 * 1024 * 1024 * 1024); +}); + +test('ignores expired active entitlements during best entitlement selection', async t => { + const user = await t.context.models.user.create({ + email: 'expired-user@affine.pro', + }); + const cases = [ + { + status: 'active', + subjectId: 'expired-subscription', + expiresAt: new Date('2020-01-01T00:00:00Z'), + }, + { + status: 'grace', + subjectId: 'open-ended-grace', + }, + ]; + + for (const item of cases) { + await t.context.db.entitlement.create({ + data: { + targetType: 'user', + targetId: user.id, + source: 'cloud_subscription', + plan: 'pro', + ...item, + }, + }); + } + + t.falsy(await t.context.service.getBestEntitlement('user', user.id)); + const resolved = await t.context.service.resolveUserEntitlement(user.id); + t.is(resolved.plan, 'free'); +}); + +test('selfhosted resolution ignores unsigned DB entitlements', async t => { + const previousDeploymentType = globalThis.env.DEPLOYMENT_TYPE; + // @ts-expect-error test mutates env singleton for deployment-specific trust boundary + globalThis.env.DEPLOYMENT_TYPE = 'selfhosted'; + try { + const user = await t.context.models.user.create({ + email: 'forged-user@affine.pro', + }); + const owner = await t.context.models.user.create({ + email: 'forged-workspace-owner@affine.pro', + }); + const workspace = await t.context.models.workspace.create(owner.id); + const cases = [ + { + targetType: 'user', + targetId: user.id, + source: 'cloud_subscription', + plan: 'ai', + quantity: null, + }, + { + targetType: 'workspace', + targetId: workspace.id, + source: 'cloud_subscription', + plan: 'team', + quantity: 100, + }, + { + targetType: 'workspace', + targetId: workspace.id, + source: 'selfhost_license', + plan: 'selfhost_team', + quantity: 100, + }, + ] as const; + + for (const item of cases) { + await t.context.db.entitlement.create({ + data: { + ...item, + status: 'active', + subjectId: `${item.source}:${item.plan}:${item.targetId}`, + quantity: item.quantity ?? undefined, + }, + }); + } + + t.falsy(await t.context.service.getBestEntitlement('user', user.id)); + t.falsy( + await t.context.service.getBestEntitlement('workspace', workspace.id) + ); + + const userResolved = await t.context.service.resolveUserEntitlement( + user.id + ); + const workspaceResolved = + await t.context.service.resolveWorkspaceEntitlement(workspace.id); + + t.is(userResolved.plan, 'selfhost_free'); + t.is(workspaceResolved.plan, 'selfhost_free'); + } finally { + // @ts-expect-error restore mutable test env singleton + globalThis.env.DEPLOYMENT_TYPE = previousDeploymentType; + } +}); + +test('cloud resolution lazily imports legacy subscriptions written after backfill', async t => { + const user = await t.context.models.user.create({ + email: 'legacy-subscription-user@affine.pro', + }); + await t.context.db.subscription.create({ + data: { + targetId: user.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Yearly, + status: SubscriptionStatus.Active, + quantity: 1, + start: new Date(), + }, + }); + + const userResolved = await t.context.service.resolveUserEntitlement(user.id); + const userEntitlement = await t.context.db.entitlement.findFirst({ + where: { + targetType: 'user', + targetId: user.id, + source: 'cloud_subscription', + plan: 'pro', + }, + }); + + t.is(userResolved.plan, 'pro'); + t.is(userEntitlement?.status, 'active'); + + const owner = await t.context.models.user.create({ + email: 'legacy-subscription-owner@affine.pro', + }); + const workspace = await t.context.models.workspace.create(owner.id); + await t.context.db.subscription.create({ + data: { + targetId: workspace.id, + plan: SubscriptionPlan.Team, + recurring: SubscriptionRecurring.Yearly, + status: SubscriptionStatus.Active, + quantity: 7, + start: new Date(), + }, + }); + + const workspaceResolved = await t.context.service.resolveWorkspaceEntitlement( + workspace.id + ); + + t.is(workspaceResolved.plan, 'team'); + t.is(workspaceResolved.quantity, 7); + t.is(workspaceResolved.quota.seatLimit, 7); + + await t.context.db.subscription.delete({ + where: { + targetId_plan: { targetId: user.id, plan: SubscriptionPlan.Pro }, + }, + }); + + const revokedResolved = await t.context.service.resolveUserEntitlement( + user.id + ); + const revokedEntitlement = await t.context.db.entitlement.findFirst({ + where: { + targetType: 'user', + targetId: user.id, + source: 'cloud_subscription', + plan: 'pro', + }, + }); + + t.is(revokedResolved.plan, 'free'); + t.is(revokedEntitlement?.status, 'revoked'); +}); + +test('cloud resolution revokes projected entitlements after legacy subscription deletion', async t => { + const user = await t.context.models.user.create({ + email: 'legacy-delete-user@affine.pro', + }); + const entitlement = await t.context.service.upsertFromCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Yearly, + status: SubscriptionStatus.Active, + }); + + await t.context.db.subscription.findUniqueOrThrow({ + where: { + targetId_plan: { targetId: user.id, plan: SubscriptionPlan.Pro }, + }, + }); + await t.context.db.subscription.delete({ + where: { + targetId_plan: { targetId: user.id, plan: SubscriptionPlan.Pro }, + }, + }); + + const resolved = await t.context.service.resolveUserEntitlement(user.id); + const updated = await t.context.db.entitlement.findUnique({ + where: { id: entitlement.id }, + }); + + t.is(resolved.plan, 'free'); + t.is(updated?.status, 'revoked'); +}); + +test('cloud resolution keeps projected string-subscription entitlements while legacy row exists', async t => { + const user = await t.context.models.user.create({ + email: 'string-subscription-user@affine.pro', + }); + const entitlement = await t.context.service.upsertFromCloudSubscription({ + targetId: user.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Yearly, + status: SubscriptionStatus.Active, + subscriptionId: 'sub_legacy_string', + }); + + await t.context.db.subscription.findUniqueOrThrow({ + where: { + targetId_plan: { targetId: user.id, plan: SubscriptionPlan.Pro }, + }, + }); + + const resolved = await t.context.service.resolveUserEntitlement(user.id); + const updated = await t.context.db.entitlement.findUnique({ + where: { id: entitlement.id }, + }); + + t.is(resolved.plan, 'pro'); + t.is(updated?.status, 'active'); +}); diff --git a/packages/backend/server/src/core/entitlement/index.ts b/packages/backend/server/src/core/entitlement/index.ts new file mode 100644 index 0000000000..6e8df528ac --- /dev/null +++ b/packages/backend/server/src/core/entitlement/index.ts @@ -0,0 +1,23 @@ +import { Module } from '@nestjs/common'; + +import { LegacyEntitlementProjectionService } from './projection'; +import { EntitlementProjectionChecker } from './projection-checker'; +import { EntitlementService } from './service'; + +@Module({ + providers: [ + EntitlementService, + LegacyEntitlementProjectionService, + EntitlementProjectionChecker, + ], + exports: [ + EntitlementService, + LegacyEntitlementProjectionService, + EntitlementProjectionChecker, + ], +}) +export class EntitlementModule {} + +export { EntitlementService }; +export { EntitlementProjectionChecker }; +export { LegacyEntitlementProjectionService }; diff --git a/packages/backend/server/src/core/entitlement/projection-checker.ts b/packages/backend/server/src/core/entitlement/projection-checker.ts new file mode 100644 index 0000000000..5cc3ac9e34 --- /dev/null +++ b/packages/backend/server/src/core/entitlement/projection-checker.ts @@ -0,0 +1,290 @@ +import { Injectable } from '@nestjs/common'; +import { PrismaClient } from '@prisma/client'; + +@Injectable() +export class EntitlementProjectionChecker { + constructor(private readonly db: PrismaClient) {} + + async checkEntitlementProjection() { + const now = new Date(); + const [ + missingEffectiveUserState, + missingEffectiveWorkspaceState, + staleEffectiveUserState, + staleEffectiveWorkspaceState, + cloudSubscriptionProjectionMissing, + selfhostLicenseProjectionMissing, + cloudSubscriptionEntitlementMissing, + selfhostLicenseEntitlementMissing, + dirtyLegacyUserFeatures, + dirtyLegacyWorkspaceFeatures, + missingUserFeatureProjection, + missingWorkspaceFeatureProjection, + ] = await Promise.all([ + this.db.user.count({ + where: { quotaState: null }, + }), + this.db.workspace.count({ + where: { quotaState: null }, + }), + this.db.effectiveUserQuotaState.count({ + where: { + OR: [{ stale: true }, { known: false }, { staleAfter: { lt: now } }], + }, + }), + this.db.effectiveWorkspaceQuotaState.count({ + where: { + OR: [{ stale: true }, { known: false }, { staleAfter: { lt: now } }], + }, + }), + this.cloudSubscriptionProjectionMissing(), + this.selfhostLicenseProjectionMissing(), + this.cloudSubscriptionEntitlementMissing(), + this.selfhostLicenseEntitlementMissing(), + this.dirtyLegacyUserFeatures(), + this.dirtyLegacyWorkspaceFeatures(), + this.missingUserFeatureProjection(), + this.missingWorkspaceFeatureProjection(), + ]); + + return { + missingEffectiveUserState, + missingEffectiveWorkspaceState, + staleEffectiveUserState, + staleEffectiveWorkspaceState, + cloudSubscriptionProjectionMissing, + selfhostLicenseProjectionMissing, + cloudSubscriptionEntitlementMissing, + selfhostLicenseEntitlementMissing, + dirtyLegacyUserFeatures, + dirtyLegacyWorkspaceFeatures, + missingUserFeatureProjection, + missingWorkspaceFeatureProjection, + }; + } + + private async cloudSubscriptionProjectionMissing() { + const legacyKeys = new Set( + ( + await this.db.subscription.findMany({ + where: { + status: { in: ['active', 'trialing', 'past_due'] }, + }, + select: { targetId: true, plan: true }, + }) + ).map(subscription => `${subscription.targetId}:${subscription.plan}`) + ); + const entitlements = await this.validEntitlements({ + source: 'cloud_subscription', + }); + + return entitlements.filter( + entitlement => + entitlement.targetId && + !legacyKeys.has( + `${entitlement.targetId}:${this.subscriptionPlan(entitlement.plan)}` + ) + ).length; + } + + private async selfhostLicenseProjectionMissing() { + const licenseKeys = new Set( + ( + await this.db.installedLicense.findMany({ + select: { key: true }, + }) + ).map(license => license.key) + ); + const entitlements = await this.validEntitlements({ + source: 'selfhost_license', + }); + + return entitlements.filter( + entitlement => + entitlement.subjectId && !licenseKeys.has(entitlement.subjectId) + ).length; + } + + private async cloudSubscriptionEntitlementMissing() { + const activeSubscriptions = await this.db.subscription.findMany({ + where: { + status: { in: ['active', 'trialing', 'past_due'] }, + }, + select: { targetId: true, plan: true }, + }); + const valid = new Set( + ( + await this.validEntitlements({ + source: 'cloud_subscription', + }) + ).map( + entitlement => + `${entitlement.targetId}:${this.subscriptionPlan(entitlement.plan)}` + ) + ); + + return activeSubscriptions.filter( + subscription => + !valid.has(`${subscription.targetId}:${subscription.plan}`) + ).length; + } + + private async selfhostLicenseEntitlementMissing() { + const licenses = await this.db.installedLicense.findMany({ + where: { + license: { not: null }, + }, + select: { key: true }, + }); + const validKeys = new Set( + ( + await this.validEntitlements({ + source: 'selfhost_license', + }) + ).flatMap(entitlement => entitlement.subjectId ?? []) + ); + + return licenses.filter(license => !validKeys.has(license.key)).length; + } + + private async dirtyLegacyUserFeatures() { + const rows = await this.db.userFeature.findMany({ + where: { + activated: true, + name: { + in: ['pro_plan_v1', 'lifetime_pro_plan_v1', 'unlimited_copilot'], + }, + }, + select: { + userId: true, + name: true, + }, + }); + + const valid = new Set( + ( + await this.validEntitlements({ + targetType: 'user', + plan: { in: ['pro', 'lifetime_pro', 'ai'] }, + }) + ).map(entitlement => `${entitlement.targetId}:${entitlement.plan}`) + ); + + return rows.filter(row => { + const plan = + row.name === 'lifetime_pro_plan_v1' + ? 'lifetime_pro' + : row.name === 'pro_plan_v1' + ? 'pro' + : 'ai'; + return !valid.has(`${row.userId}:${plan}`); + }).length; + } + + private async dirtyLegacyWorkspaceFeatures() { + const rows = await this.db.workspaceFeature.findMany({ + where: { + activated: true, + name: 'team_plan_v1', + }, + select: { workspaceId: true }, + }); + const validWorkspaceIds = new Set( + ( + await this.validEntitlements({ + targetType: 'workspace', + plan: { in: ['team', 'selfhost_team'] }, + }) + ).flatMap(entitlement => entitlement.targetId ?? []) + ); + + return rows.filter(row => !validWorkspaceIds.has(row.workspaceId)).length; + } + + private async missingUserFeatureProjection() { + const entitlements = await this.validEntitlements({ + targetType: 'user', + plan: { in: ['pro', 'lifetime_pro', 'ai'] }, + }); + const features = new Set( + ( + await this.db.userFeature.findMany({ + where: { + activated: true, + name: { + in: ['pro_plan_v1', 'lifetime_pro_plan_v1', 'unlimited_copilot'], + }, + }, + select: { userId: true, name: true }, + }) + ).map(feature => `${feature.userId}:${feature.name}`) + ); + + return entitlements.filter(entitlement => { + if (!entitlement.targetId) { + return false; + } + const feature = + entitlement.plan === 'lifetime_pro' + ? 'lifetime_pro_plan_v1' + : entitlement.plan === 'pro' + ? 'pro_plan_v1' + : 'unlimited_copilot'; + return !features.has(`${entitlement.targetId}:${feature}`); + }).length; + } + + private async missingWorkspaceFeatureProjection() { + const entitlements = await this.validEntitlements({ + targetType: 'workspace', + plan: { in: ['team', 'selfhost_team'] }, + }); + const featureWorkspaceIds = new Set( + ( + await this.db.workspaceFeature.findMany({ + where: { + activated: true, + name: 'team_plan_v1', + }, + select: { workspaceId: true }, + }) + ).map(feature => feature.workspaceId) + ); + + return entitlements.filter( + entitlement => + entitlement.targetId && !featureWorkspaceIds.has(entitlement.targetId) + ).length; + } + + private validEntitlements(where: Record) { + const now = new Date(); + return this.db.entitlement.findMany({ + where: { + ...where, + ...(where.source === 'selfhost_license' + ? { signedPayload: { not: null } } + : {}), + OR: [ + { + status: 'active', + OR: [{ expiresAt: null }, { expiresAt: { gt: now } }], + }, + { + status: 'grace', + graceUntil: { gt: now }, + }, + ], + }, + select: { + targetId: true, + subjectId: true, + plan: true, + }, + }); + } + + private subscriptionPlan(plan: string) { + return plan === 'lifetime_pro' ? 'pro' : plan; + } +} diff --git a/packages/backend/server/src/core/entitlement/projection.ts b/packages/backend/server/src/core/entitlement/projection.ts new file mode 100644 index 0000000000..d4d4624083 --- /dev/null +++ b/packages/backend/server/src/core/entitlement/projection.ts @@ -0,0 +1,465 @@ +import { Injectable } from '@nestjs/common'; +import { Entitlement, PrismaClient } from '@prisma/client'; + +import { OnEvent } from '../../base'; +import { Models } from '../../models'; +import { + SubscriptionPlan, + SubscriptionRecurring, + SubscriptionStatus, +} from '../../plugins/payment/types'; +import { EntitlementService } from './service'; + +type Metadata = { + provider?: string | null; + recurring?: string | null; + variant?: string | null; + subscriptionId?: string | number | null; + stripeSubscriptionId?: string | null; + validateKey?: string | null; + legacyProjected?: boolean; +}; + +@Injectable() +export class LegacyEntitlementProjectionService { + constructor( + private readonly db: PrismaClient, + private readonly models: Models, + private readonly entitlement: EntitlementService + ) {} + + @OnEvent('entitlement.changed') + async onEntitlementChanged({ + targetType, + targetId, + }: Events['entitlement.changed']) { + if (targetType === 'user') { + await this.projectCloudSubscriptions('user', targetId); + await this.projectUserFeatures(targetId); + } else if (targetType === 'workspace') { + await this.projectCloudSubscriptions('workspace', targetId); + await Promise.all([ + this.projectWorkspaceFeatures(targetId), + this.projectInstalledLicense(targetId), + ]); + } + } + + @OnEvent('workspace.quota_state.changed') + async onWorkspaceQuotaStateChanged({ + workspaceId, + }: Events['workspace.quota_state.changed']) { + await this.projectReadonlyFeature(workspaceId); + } + + async scanInstalledLicenses() { + const licenses = await this.db.installedLicense.findMany(); + + await Promise.all( + licenses.map(async license => + license.license + ? await this.entitlement.upsertFromSelfhostLicense({ + workspaceId: license.workspaceId, + licenseKey: license.key, + recurring: license.recurring, + quantity: license.quantity, + expiresAt: license.expiredAt, + validatedAt: license.validatedAt, + license: Buffer.from(license.license), + }) + : license.validateKey + ? await this.entitlement.upsertFromValidatedSelfhostLicense({ + workspaceId: license.workspaceId, + licenseKey: license.key, + recurring: license.recurring, + quantity: license.quantity, + expiresAt: license.expiredAt, + validatedAt: license.validatedAt, + validateKey: license.validateKey, + variant: license.variant, + }) + : await this.entitlement.markSelfhostLicenseNeedsReupload({ + workspaceId: license.workspaceId, + licenseKey: license.key, + reason: 'Installed license has no raw payload to verify.', + }) + ) + ); + } + + async backfillEntitlementsAndQuotaStates() { + const [subscriptions, users, workspaces] = await Promise.all([ + this.db.subscription.findMany(), + this.db.user.findMany({ select: { id: true } }), + this.db.workspace.findMany({ select: { id: true } }), + ]); + + for (const subscription of subscriptions) { + if (subscription.plan === SubscriptionPlan.SelfHostedTeam) { + await this.entitlement.markSelfhostLicenseNeedsReupload({ + licenseKey: subscription.targetId, + reason: + 'Historical self-hosted team subscription needs license activation or revalidation.', + }); + continue; + } + await this.entitlement.upsertFromCloudSubscription(subscription); + } + + await this.scanInstalledLicenses(); + + await Promise.all([ + ...users.map(user => + this.db.effectiveUserQuotaState.upsert({ + where: { userId: user.id }, + update: { stale: true }, + create: { + userId: user.id, + plan: 'free', + blobLimit: BigInt(0), + storageQuota: BigInt(0), + usedStorageQuota: BigInt(0), + historyPeriodSeconds: 0, + known: false, + stale: true, + }, + }) + ), + ...workspaces.map(workspace => + this.db.effectiveWorkspaceQuotaState.upsert({ + where: { workspaceId: workspace.id }, + update: { stale: true }, + create: { + workspaceId: workspace.id, + plan: 'free', + usesOwnerQuota: true, + seatLimit: 0, + memberCount: 0, + overcapacityMemberCount: 0, + blobLimit: BigInt(0), + storageQuota: BigInt(0), + usedStorageQuota: BigInt(0), + historyPeriodSeconds: 0, + known: false, + stale: true, + }, + }) + ), + ]); + } + + private async projectUserFeatures(userId: string) { + const entitlements = await this.activeEntitlements('user', userId); + const quotaEntitlement = entitlements.find(entitlement => + ['lifetime_pro', 'pro'].includes(entitlement.plan) + ); + + if (quotaEntitlement?.plan === 'lifetime_pro') { + await this.models.userFeature.switchQuota( + userId, + 'lifetime_pro_plan_v1', + 'legacy entitlement projection' + ); + } else if (quotaEntitlement?.plan === 'pro') { + await this.models.userFeature.switchQuota( + userId, + 'pro_plan_v1', + 'legacy entitlement projection' + ); + } else if ( + await this.hasActiveUserFeature(userId, [ + 'pro_plan_v1', + 'lifetime_pro_plan_v1', + ]) + ) { + await this.models.userFeature.switchQuota( + userId, + 'free_plan_v1', + 'legacy entitlement projection' + ); + } + + if (entitlements.some(entitlement => entitlement.plan === 'ai')) { + await this.models.userFeature.add( + userId, + 'unlimited_copilot', + 'legacy entitlement projection' + ); + } else { + await this.models.userFeature.remove(userId, 'unlimited_copilot'); + } + } + + private async projectWorkspaceFeatures(workspaceId: string) { + const [entitlement, resolved] = await Promise.all([ + this.entitlement.getBestEntitlement('workspace', workspaceId), + this.entitlement.resolveWorkspaceEntitlement(workspaceId), + ]); + + if ( + entitlement && + ['team', 'selfhost_team'].includes(resolved.plan) && + resolved.valid && + resolved.quota.seatLimit + ) { + await this.models.workspaceFeature.add( + workspaceId, + 'team_plan_v1', + 'legacy entitlement projection', + { + memberLimit: resolved.quota.seatLimit, + } + ); + } else { + await this.models.workspaceFeature.remove(workspaceId, 'team_plan_v1'); + } + } + + private async projectCloudSubscriptions( + targetType: 'user' | 'workspace', + targetId: string + ) { + if (env.selfhosted) return; + const entitlements = await this.db.entitlement.findMany({ + where: { + targetType, + targetId, + source: 'cloud_subscription', + }, + orderBy: { updatedAt: 'asc' }, + }); + + for (const entitlement of this.projectableCloudEntitlements(entitlements)) { + const metadata = entitlement.metadata as Metadata; + await this.db.subscription.upsert({ + where: { + targetId_plan: { + targetId, + plan: this.subscriptionPlan(entitlement.plan), + }, + }, + update: { + recurring: metadata.recurring ?? SubscriptionRecurring.Monthly, + variant: metadata.variant ?? null, + quantity: entitlement.quantity ?? 1, + stripeSubscriptionId: metadata.stripeSubscriptionId ?? null, + provider: this.provider(metadata.provider), + status: this.subscriptionStatus(entitlement.status), + start: entitlement.startsAt ?? entitlement.createdAt, + end: entitlement.expiresAt, + trialEnd: entitlement.graceUntil, + }, + create: { + targetId, + plan: this.subscriptionPlan(entitlement.plan), + recurring: metadata.recurring ?? SubscriptionRecurring.Monthly, + variant: metadata.variant ?? null, + quantity: entitlement.quantity ?? 1, + stripeSubscriptionId: metadata.stripeSubscriptionId ?? null, + provider: this.provider(metadata.provider), + status: this.subscriptionStatus(entitlement.status), + start: entitlement.startsAt ?? entitlement.createdAt, + end: entitlement.expiresAt, + trialEnd: entitlement.graceUntil, + }, + }); + if (!metadata.legacyProjected) { + await this.db.entitlement.update({ + where: { id: entitlement.id }, + data: { + metadata: { + ...metadata, + legacyProjected: true, + }, + }, + }); + } + } + } + + private *projectableCloudEntitlements(entitlements: Entitlement[]) { + const byPlan = new Map(); + + for (const entitlement of entitlements) { + const plan = this.subscriptionPlan(entitlement.plan); + const current = byPlan.get(plan); + + if ( + !current || + this.subscriptionProjectionPriority(entitlement) > + this.subscriptionProjectionPriority(current) + ) { + byPlan.set(plan, entitlement); + } + } + + yield* byPlan.values(); + } + + private subscriptionProjectionPriority(entitlement: { + status: string; + updatedAt: Date; + }) { + const statusPriority = + entitlement.status === 'active' || entitlement.status === 'grace' + ? 2 + : entitlement.status === 'expired' + ? 1 + : 0; + + return ( + statusPriority * 10_000_000_000_000 + entitlement.updatedAt.getTime() + ); + } + + private async projectInstalledLicense(workspaceId: string) { + const [entitlements, resolved] = await Promise.all([ + this.db.entitlement.findMany({ + where: { + targetType: 'workspace', + targetId: workspaceId, + source: 'selfhost_license', + }, + orderBy: [{ signedPayload: 'desc' }, { updatedAt: 'desc' }], + }), + this.entitlement.resolveWorkspaceEntitlement(workspaceId), + ]); + const entitlement = entitlements.sort( + (left, right) => + this.installedLicenseStatusPriority(right.status) - + this.installedLicenseStatusPriority(left.status) || + Number(!!right.signedPayload) - Number(!!left.signedPayload) || + right.updatedAt.getTime() - left.updatedAt.getTime() + )[0]; + + if (!entitlement) { + return; + } + + if ( + resolved.plan !== 'selfhost_team' || + !['active', 'grace', 'expired'].includes(resolved.status) + ) { + await this.db.installedLicense.deleteMany({ + where: { workspaceId }, + }); + return; + } + + const metadata = entitlement.metadata as Metadata; + const expiredAt = resolved.expiresAt + ? new Date(resolved.expiresAt) + : entitlement.expiresAt; + await this.db.installedLicense.upsert({ + where: { workspaceId }, + update: { + key: resolved.subjectId ?? entitlement.subjectId ?? entitlement.id, + quantity: resolved.quantity ?? 1, + recurring: + resolved.recurring ?? + metadata.recurring ?? + SubscriptionRecurring.Monthly, + variant: metadata.variant ?? null, + validateKey: metadata.validateKey ?? '', + validatedAt: entitlement.validatedAt ?? new Date(), + expiredAt, + license: entitlement.signedPayload + ? Buffer.from(entitlement.signedPayload) + : null, + }, + create: { + workspaceId, + key: resolved.subjectId ?? entitlement.subjectId ?? entitlement.id, + quantity: resolved.quantity ?? 1, + recurring: + resolved.recurring ?? + metadata.recurring ?? + SubscriptionRecurring.Monthly, + variant: metadata.variant ?? null, + validateKey: metadata.validateKey ?? '', + validatedAt: entitlement.validatedAt ?? new Date(), + expiredAt, + license: entitlement.signedPayload + ? Buffer.from(entitlement.signedPayload) + : null, + }, + }); + } + + private installedLicenseStatusPriority(status: string) { + if (status === 'active' || status === 'grace') { + return 3; + } + if (status === 'expired') { + return 2; + } + if (status === 'needs_reupload') { + return 1; + } + return 0; + } + + private async projectReadonlyFeature(workspaceId: string) { + const state = await this.db.effectiveWorkspaceQuotaState.findUnique({ + where: { + workspaceId, + }, + }); + + if (state?.readonly) { + await this.models.workspaceFeature.add( + workspaceId, + 'quota_exceeded_readonly_workspace_v1', + `legacy quota state projection: ${state.readonlyReasons.join(',')}` + ); + } else { + await this.models.workspaceFeature.remove( + workspaceId, + 'quota_exceeded_readonly_workspace_v1' + ); + } + } + + private async activeEntitlements( + targetType: 'user' | 'workspace', + targetId: string + ) { + return this.entitlement.getActiveEntitlements(targetType, targetId); + } + + private async hasActiveUserFeature(userId: string, names: string[]) { + const count = await this.db.userFeature.count({ + where: { + userId, + name: { in: names }, + activated: true, + }, + }); + + return count > 0; + } + + private subscriptionPlan(plan: string) { + if (plan === 'lifetime_pro') { + return SubscriptionPlan.Pro; + } + if (plan === 'selfhost_team') { + return SubscriptionPlan.SelfHostedTeam; + } + return plan; + } + + private subscriptionStatus(status: string) { + if (status === 'active') { + return SubscriptionStatus.Active; + } + if (status === 'grace') { + return SubscriptionStatus.PastDue; + } + return SubscriptionStatus.Canceled; + } + + private provider(provider: string | null | undefined) { + return provider === 'revenuecat' ? 'revenuecat' : 'stripe'; + } +} diff --git a/packages/backend/server/src/core/entitlement/service.ts b/packages/backend/server/src/core/entitlement/service.ts new file mode 100644 index 0000000000..383d142281 --- /dev/null +++ b/packages/backend/server/src/core/entitlement/service.ts @@ -0,0 +1,1037 @@ +import { Injectable } from '@nestjs/common'; +import { Entitlement, Prisma, PrismaClient } from '@prisma/client'; + +import { BadRequest, CryptoHelper, EventBus } from '../../base'; +import { resolveEntitlementV1 } from '../../native'; +import { + SubscriptionPlan, + SubscriptionRecurring, + SubscriptionStatus, + SubscriptionVariant, +} from '../../plugins/payment/types'; + +type TargetType = 'user' | 'workspace' | 'instance'; +type EntitlementStatus = + | 'active' + | 'grace' + | 'expired' + | 'revoked' + | 'needs_reupload'; + +export interface CloudSubscriptionEntitlementInput { + targetId: string; + plan: SubscriptionPlan | string; + recurring: SubscriptionRecurring | string; + status: string; + quantity?: number | null; + variant?: SubscriptionVariant | string | null; + provider?: string | null; + subscriptionId?: string | number | null; + stripeSubscriptionId?: string | null; + start?: Date | null; + end?: Date | null; + trialStart?: Date | null; + trialEnd?: Date | null; + canceledAt?: Date | null; +} + +export interface SelfhostLicenseEntitlementInput { + workspaceId: string; + licenseKey?: string; + recurring?: SubscriptionRecurring | string; + quantity?: number; + validateKey?: string | null; + variant?: SubscriptionVariant | string | null; + expiresAt?: Date | null; + validatedAt?: Date | null; + license?: Buffer | null; +} + +interface RemoteSelfhostLicense { + plan: string; + recurring: string; + quantity: number; + endAt: number; +} + +const REMOTE_SELFHOST_LICENSE_REVALIDATE_INTERVAL = 1000 * 60 * 10; +const REMOTE_SELFHOST_LICENSE_HEALTH_TIMEOUT = 10_000; + +declare global { + interface Events { + 'entitlement.changed': { + targetType: TargetType; + targetId: string; + }; + } +} + +@Injectable() +export class EntitlementService { + private readonly legacyCloudSubscriptionSyncs = new Map< + string, + Promise + >(); + private readonly remoteSelfhostLicenseVerifications = new Map< + string, + Promise + >(); + private readonly remoteSelfhostLicenseCache = new Map< + string, + { entitlement: Entitlement; verifiedUntil: number } + >(); + + constructor( + private readonly db: PrismaClient, + private readonly crypto: CryptoHelper, + private readonly event: EventBus + ) {} + + getUserEntitlements(userId: string) { + return this.db.entitlement.findMany({ + where: { targetType: 'user', targetId: userId }, + orderBy: { updatedAt: 'desc' }, + }); + } + + getWorkspaceEntitlements(workspaceId: string) { + return this.db.entitlement.findMany({ + where: { targetType: 'workspace', targetId: workspaceId }, + orderBy: { updatedAt: 'desc' }, + }); + } + + resolveUserEntitlement(userId: string) { + return this.resolveBestEntitlement('user', userId); + } + + resolveWorkspaceEntitlement(workspaceId: string) { + return this.resolveBestEntitlement('workspace', workspaceId); + } + + async getBestEntitlement(targetType: TargetType, targetId: string) { + await this.syncLegacyCloudSubscriptionEntitlements(targetType, targetId); + const entitlements = await this.db.entitlement.findMany({ + where: { + targetType, + targetId, + plan: targetType === 'user' ? { not: 'ai' } : undefined, + AND: [ + this.validSelfhostEntitlementWhere(), + this.validEntitlementWhereForTarget(targetType, new Date()), + ], + }, + orderBy: { updatedAt: 'desc' }, + }); + + const sorted = entitlements.sort( + (a, b) => this.priority(b) - this.priority(a) + ); + if (!env.selfhosted || targetType !== 'workspace') { + return sorted[0]; + } + + for (const entitlement of sorted) { + if (entitlement.signedPayload) { + return entitlement; + } + + const verified = await this.verifyRemoteSelfhostLicense(entitlement); + if (verified) { + return verified; + } + } + + return; + } + + async getActiveEntitlements(targetType: TargetType, targetId: string) { + await this.syncLegacyCloudSubscriptionEntitlements(targetType, targetId); + return this.db.entitlement.findMany({ + where: { + targetType, + targetId, + AND: [ + this.validSelfhostEntitlementWhere(), + this.validEntitlementWhereForTarget(targetType, new Date()), + ], + }, + orderBy: { updatedAt: 'desc' }, + }); + } + + async upsertFromCloudSubscription( + input: CloudSubscriptionEntitlementInput, + options: { emit?: boolean; legacySync?: boolean } = {} + ) { + const emit = options.emit ?? true; + const targetType = this.targetTypeForPlan(input.plan); + const subjectId = this.cloudSubjectId(input); + const status = this.statusFromSubscription(input.status); + const entitlement = await this.findBySubject( + 'cloud_subscription', + subjectId + ); + + const data = { + targetType, + targetId: input.targetId, + source: 'cloud_subscription', + plan: this.entitlementPlan(input.plan, input.recurring), + status, + subjectId, + quantity: + targetType === 'workspace' + ? this.normalizedQuantity(input.quantity) + : undefined, + metadata: { + provider: input.provider ?? 'stripe', + recurring: input.recurring, + variant: input.variant ?? null, + subscriptionId: input.subscriptionId ?? null, + stripeSubscriptionId: input.stripeSubscriptionId ?? null, + legacySync: options.legacySync ?? false, + }, + startsAt: input.start ?? undefined, + expiresAt: input.end ?? undefined, + graceUntil: + status === 'grace' ? (input.trialEnd ?? input.end ?? new Date()) : null, + validatedAt: new Date(), + }; + + if (entitlement) { + const updated = await this.db.entitlement.update({ + where: { id: entitlement.id }, + data, + }); + if (emit) { + await this.emitEntitlementChanged(updated); + } + return updated; + } + + const created = await this.db.entitlement.create({ data }); + if (emit) { + await this.emitEntitlementChanged(created); + } + return created; + } + + async upsertAdminGrant(input: { + targetType: Exclude; + targetId: string; + plan: string; + quantity?: number | null; + }) { + this.assertAdminGrantInput(input.targetType, input.plan); + if (env.selfhosted) { + throw new BadRequest( + 'Self-hosted commercial entitlements require a signed license.' + ); + } + const quantity = + input.targetType === 'workspace' + ? this.normalizedQuantity(input.quantity) + : undefined; + resolveEntitlementV1({ + deploymentType: 'cloud', + targetType: input.targetType, + targetId: input.targetId, + plan: input.plan, + quantity, + now: new Date().toISOString(), + }); + + const subjectId = this.adminGrantSubjectId( + input.targetType, + input.targetId + ); + const data = { + targetType: input.targetType, + targetId: input.targetId, + source: 'admin_grant', + plan: input.plan, + status: 'active', + subjectId, + quantity, + metadata: {}, + validatedAt: new Date(), + }; + const existing = await this.findBySubject('admin_grant', subjectId); + const entitlement = existing + ? await this.db.entitlement.update({ + where: { id: existing.id }, + data, + }) + : await this.db.entitlement.create({ data }); + const replaced = await this.db.entitlement.findMany({ + where: { + source: 'admin_grant', + targetType: input.targetType, + targetId: input.targetId, + status: { in: ['active', 'grace'] }, + id: { not: entitlement.id }, + }, + }); + await this.db.entitlement.updateMany({ + where: { id: { in: replaced.map(entitlement => entitlement.id) } }, + data: { status: 'revoked' }, + }); + await this.emitEntitlementChanged(entitlement); + return entitlement; + } + + async revokeAdminGrant( + targetType: Exclude, + targetId: string + ) { + const entitlements = await this.db.entitlement.findMany({ + where: { + source: 'admin_grant', + targetType, + targetId, + status: { in: ['active', 'grace'] }, + }, + }); + await this.db.entitlement.updateMany({ + where: { id: { in: entitlements.map(entitlement => entitlement.id) } }, + data: { status: 'revoked' }, + }); + await Promise.all( + entitlements.map(entitlement => this.emitEntitlementChanged(entitlement)) + ); + } + + async syncLegacyCloudSubscriptionEntitlements( + targetType: TargetType, + targetId: string + ) { + if (env.selfhosted || targetType === 'instance') { + return; + } + + const key = `${targetType}:${targetId}`; + const existing = this.legacyCloudSubscriptionSyncs.get(key); + if (existing) { + return existing; + } + + const task = this.doSyncLegacyCloudSubscriptionEntitlements( + targetType, + targetId + ) + .then(changed => { + if (changed) { + this.event.emit('entitlement.changed', { targetType, targetId }); + } + }) + .finally(() => { + this.legacyCloudSubscriptionSyncs.delete(key); + }); + this.legacyCloudSubscriptionSyncs.set(key, task); + return task; + } + + async upsertFromSelfhostLicense(input: SelfhostLicenseEntitlementInput) { + const resolved = input.license + ? resolveEntitlementV1({ + deploymentType: 'selfhosted', + targetType: 'workspace', + targetId: input.workspaceId, + signedPayload: input.license, + publicKey: this.crypto.AFFiNEProPublicKey?.toString(), + licenseAesKey: this.crypto.AFFiNEProLicenseAESKey?.toString('hex'), + now: new Date().toISOString(), + }) + : null; + const valid = resolved?.valid === true; + const subjectId = resolved?.subjectId ?? input.licenseKey; + if (!subjectId) { + throw new Error('selfhost license key is required'); + } + const entitlement = await this.findBySubject('selfhost_license', subjectId); + + const data = { + targetType: 'workspace', + targetId: input.workspaceId, + source: 'selfhost_license', + plan: 'selfhost_team', + status: valid ? 'active' : ('needs_reupload' as EntitlementStatus), + subjectId, + quantity: valid ? resolved.quantity : undefined, + signedPayload: input.license ?? undefined, + metadata: { + recurring: resolved?.recurring ?? input.recurring, + validateKey: input.validateKey ?? '', + variant: input.variant ?? null, + errorCode: resolved?.errorCode ?? (valid ? null : 'needs_reupload'), + errorMessage: + resolved?.errorMessage ?? + (valid ? null : 'Self-hosted license needs raw payload to verify.'), + }, + expiresAt: + input.expiresAt ?? + (resolved?.expiresAt ? new Date(resolved.expiresAt) : undefined), + validatedAt: input.validatedAt ?? new Date(), + }; + + if (entitlement) { + const updated = await this.db.entitlement.update({ + where: { id: entitlement.id }, + data, + }); + await this.emitEntitlementChanged(updated); + return updated; + } + + const created = await this.db.entitlement.create({ data }); + await this.emitEntitlementChanged(created); + return created; + } + + async upsertFromValidatedSelfhostLicense( + input: Omit & { + licenseKey: string; + quantity: number; + } + ) { + const entitlement = await this.findBySubject( + 'selfhost_license', + input.licenseKey + ); + const data = { + targetType: 'workspace', + targetId: input.workspaceId, + source: 'selfhost_license', + plan: 'selfhost_team', + status: 'active' as EntitlementStatus, + subjectId: input.licenseKey, + quantity: this.normalizedQuantity(input.quantity), + signedPayload: null, + metadata: { + recurring: input.recurring, + validateKey: input.validateKey ?? '', + variant: input.variant ?? null, + remoteValidated: true, + }, + expiresAt: input.expiresAt ?? undefined, + validatedAt: input.validatedAt ?? new Date(), + }; + + if (entitlement) { + const updated = await this.db.entitlement.update({ + where: { id: entitlement.id }, + data, + }); + await this.emitEntitlementChanged(updated); + return updated; + } + + const created = await this.db.entitlement.create({ data }); + await this.emitEntitlementChanged(created); + return created; + } + + async markSelfhostLicenseNeedsReupload(input: { + workspaceId?: string; + licenseKey: string; + reason: string; + }) { + const entitlement = await this.findBySubject( + 'selfhost_license', + input.licenseKey + ); + const targetType = input.workspaceId ? 'workspace' : 'instance'; + const data = { + targetType, + targetId: input.workspaceId ?? input.licenseKey, + source: 'selfhost_license', + plan: 'selfhost_team', + status: 'needs_reupload' as EntitlementStatus, + subjectId: input.licenseKey, + quantity: null, + metadata: { + errorCode: 'needs_reupload', + errorMessage: input.reason, + }, + validatedAt: new Date(), + }; + + if (entitlement) { + const updated = await this.db.entitlement.update({ + where: { id: entitlement.id }, + data, + }); + await this.emitEntitlementChanged(updated); + return updated; + } + + const created = await this.db.entitlement.create({ data }); + await this.emitEntitlementChanged(created); + return created; + } + + async revokeBySubject(source: string, subjectId: string) { + const entitlements = await this.db.entitlement.findMany({ + where: { source, subjectId, status: { in: ['active', 'grace'] } }, + }); + await this.db.entitlement.updateMany({ + where: { source, subjectId, status: { in: ['active', 'grace'] } }, + data: { status: 'revoked' }, + }); + await Promise.all( + entitlements.map(entitlement => this.emitEntitlementChanged(entitlement)) + ); + } + + async revokeCloudSubscription(input: { + targetId: string; + plan: SubscriptionPlan | string; + subscriptionId?: string | number | null; + stripeSubscriptionId?: string | null; + }) { + await this.db.subscription.updateMany({ + where: { + targetId: input.targetId, + plan: input.plan, + }, + data: { + status: SubscriptionStatus.Canceled, + end: new Date(), + }, + }); + await this.revokeBySubject( + 'cloud_subscription', + this.cloudSubjectId(input) + ); + if ( + !input.stripeSubscriptionId && + typeof input.subscriptionId !== 'string' + ) { + await this.revokeCloudSubscriptionByLegacyTarget(input); + } + } + + async builtinFree(targetType: TargetType) { + return resolveEntitlementV1({ + deploymentType: env.selfhosted ? 'selfhosted' : 'cloud', + targetType, + now: new Date().toISOString(), + }); + } + + private async resolveBestEntitlement( + targetType: TargetType, + targetId: string + ) { + const entitlement = await this.getBestEntitlement(targetType, targetId); + + if (!entitlement) { + return this.builtinFree(targetType); + } + + const deploymentType = + env.selfhosted && + entitlement.source === 'selfhost_license' && + !entitlement.signedPayload + ? 'cloud' + : env.selfhosted + ? 'selfhosted' + : 'cloud'; + + try { + return resolveEntitlementV1({ + deploymentType, + targetType, + targetId, + plan: entitlement.plan, + quantity: entitlement.quantity ?? undefined, + signedPayload: entitlement.signedPayload + ? Buffer.from(entitlement.signedPayload) + : undefined, + publicKey: this.crypto.AFFiNEProPublicKey?.toString(), + licenseAesKey: this.crypto.AFFiNEProLicenseAESKey?.toString('hex'), + now: new Date().toISOString(), + }); + } catch (e) { + if ( + env.selfhosted && + entitlement.source === 'selfhost_license' && + entitlement.plan === 'selfhost_team' + ) { + await this.markRemoteSelfhostLicenseNeedsReupload( + entitlement, + e instanceof Error ? e.message : 'Invalid self-hosted license.' + ); + return this.builtinFree(targetType); + } + + throw e; + } + } + + private findBySubject(source: string, subjectId: string) { + return this.db.entitlement.findFirst({ + where: { source, subjectId }, + orderBy: { updatedAt: 'desc' }, + }); + } + + private adminGrantSubjectId(targetType: TargetType, targetId: string) { + return `admin_grant:${targetType}:${targetId}`; + } + + private assertAdminGrantInput(targetType: string, plan: string) { + const plans = + targetType === 'user' + ? ['pro', 'lifetime_pro', 'ai'] + : targetType === 'workspace' + ? ['team'] + : null; + if (!plans?.includes(plan)) { + throw new BadRequest( + `Admin grant plan ${plan} is not configurable for ${targetType}` + ); + } + } + + private async doSyncLegacyCloudSubscriptionEntitlements( + targetType: Exclude, + targetId: string + ) { + let changed = false; + const legacyPlans = + targetType === 'user' + ? [SubscriptionPlan.Pro, SubscriptionPlan.AI] + : [SubscriptionPlan.Team]; + const entitlementPlans = + targetType === 'user' ? ['pro', 'lifetime_pro', 'ai'] : ['team']; + const subscriptions = await this.db.subscription.findMany({ + where: { + targetId, + plan: { in: legacyPlans }, + }, + orderBy: { updatedAt: 'asc' }, + }); + const legacySubjects = new Set( + subscriptions.map(subscription => this.cloudSubjectId(subscription)) + ); + const legacySubscriptionPlans = new Set( + subscriptions.map(subscription => subscription.plan) + ); + + for (const subscription of subscriptions) { + const before = await this.findBySubject( + 'cloud_subscription', + this.cloudSubjectId(subscription) + ); + const entitlement = await this.upsertFromCloudSubscription(subscription, { + emit: false, + legacySync: true, + }); + changed = + changed || + !before || + before.targetType !== entitlement.targetType || + before.targetId !== entitlement.targetId || + before.plan !== entitlement.plan || + before.status !== entitlement.status || + before.quantity !== entitlement.quantity || + before.expiresAt?.getTime() !== entitlement.expiresAt?.getTime() || + before.graceUntil?.getTime() !== entitlement.graceUntil?.getTime(); + } + + const staleEntitlements = await this.db.entitlement.findMany({ + where: { + targetType, + targetId, + source: 'cloud_subscription', + plan: { in: entitlementPlans }, + status: { in: ['active', 'grace'] }, + OR: [ + { metadata: { path: ['legacySync'], equals: true } }, + { metadata: { path: ['legacyProjected'], equals: true } }, + ], + }, + }); + const staleIds = staleEntitlements + .filter( + entitlement => + !legacySubjects.has(entitlement.subjectId ?? '') && + !legacySubscriptionPlans.has( + this.legacySubscriptionPlan(entitlement.plan) + ) + ) + .map(entitlement => entitlement.id); + + if (staleIds.length) { + await this.db.entitlement.updateMany({ + where: { id: { in: staleIds } }, + data: { status: 'revoked' }, + }); + changed = true; + } + + return changed; + } + + private cloudSubjectId( + input: Pick< + CloudSubscriptionEntitlementInput, + 'targetId' | 'plan' | 'subscriptionId' | 'stripeSubscriptionId' + > + ) { + return ( + input.stripeSubscriptionId ?? + (typeof input.subscriptionId === 'string' + ? input.subscriptionId + : `${input.targetId}:${input.plan}`) + ); + } + + private legacySubscriptionPlan(plan: string) { + if (plan === 'pro' || plan === 'lifetime_pro') { + return SubscriptionPlan.Pro; + } + if (plan === 'ai') { + return SubscriptionPlan.AI; + } + if (plan === 'team') { + return SubscriptionPlan.Team; + } + return plan; + } + + private async revokeCloudSubscriptionByLegacyTarget(input: { + targetId: string; + plan: SubscriptionPlan | string; + }) { + const targetType = this.targetTypeForPlan(input.plan); + const plans = + input.plan === SubscriptionPlan.Pro + ? ['pro', 'lifetime_pro'] + : [this.entitlementPlan(input.plan, SubscriptionRecurring.Monthly)]; + const entitlements = await this.db.entitlement.findMany({ + where: { + targetType, + targetId: input.targetId, + source: 'cloud_subscription', + plan: { in: plans }, + status: { in: ['active', 'grace'] }, + }, + }); + if (!entitlements.length) { + return; + } + + await this.db.entitlement.updateMany({ + where: { id: { in: entitlements.map(entitlement => entitlement.id) } }, + data: { status: 'revoked' }, + }); + await Promise.all( + entitlements.map(entitlement => this.emitEntitlementChanged(entitlement)) + ); + } + + private priority(entitlement: { status: string; plan: string }) { + const statusPriority = + entitlement.status === 'active' + ? 200 + : entitlement.status === 'grace' + ? 100 + : 0; + const planPriority = + entitlement.plan === 'team' || entitlement.plan === 'selfhost_team' + ? 40 + : entitlement.plan === 'lifetime_pro' + ? 30 + : entitlement.plan === 'pro' + ? 20 + : entitlement.plan === 'ai' + ? 10 + : 0; + + return statusPriority + planPriority; + } + + private validEntitlementWhere(now: Date) { + return { + OR: [ + { + status: 'active', + OR: [{ expiresAt: null }, { expiresAt: { gt: now } }], + }, + { status: 'grace', graceUntil: { gt: now } }, + ], + }; + } + + private validEntitlementWhereForTarget( + targetType: TargetType, + now: Date + ): Prisma.EntitlementWhereInput { + if (!env.selfhosted || targetType !== 'workspace') { + return this.validEntitlementWhere(now); + } + + return { + OR: [ + this.validEntitlementWhere(now), + { + status: 'active', + source: 'selfhost_license', + plan: 'selfhost_team', + signedPayload: null, + }, + ], + }; + } + + private validSelfhostEntitlementWhere(): Prisma.EntitlementWhereInput { + if (!env.selfhosted) return {}; + + return { + source: 'selfhost_license', + plan: 'selfhost_team', + targetType: 'workspace', + }; + } + + private async verifyRemoteSelfhostLicense(entitlement: Entitlement) { + const existing = this.remoteSelfhostLicenseVerifications.get( + entitlement.id + ); + if (existing) { + return existing; + } + + const task = this.doVerifyRemoteSelfhostLicense(entitlement).finally(() => { + this.remoteSelfhostLicenseVerifications.delete(entitlement.id); + }); + this.remoteSelfhostLicenseVerifications.set(entitlement.id, task); + return task; + } + + private async doVerifyRemoteSelfhostLicense(entitlement: Entitlement) { + if ( + entitlement.source !== 'selfhost_license' || + entitlement.plan !== 'selfhost_team' || + entitlement.targetType !== 'workspace' || + !entitlement.targetId || + !entitlement.subjectId + ) { + return null; + } + + const metadata = entitlement.metadata as { + validateKey?: string | null; + variant?: string | null; + }; + if (!metadata.validateKey) { + await this.markRemoteSelfhostLicenseNeedsReupload( + entitlement, + 'Missing remote validation key.' + ); + return null; + } + const cached = this.remoteSelfhostLicenseCache.get(entitlement.id); + if ( + cached && + cached.verifiedUntil > Date.now() && + cached.entitlement.expiresAt && + cached.entitlement.expiresAt > new Date() + ) { + return cached.entitlement; + } + + const endpoint = + process.env.AFFINE_PRO_SERVER_ENDPOINT ?? 'https://app.affine.pro'; + const signal = AbortSignal.timeout(REMOTE_SELFHOST_LICENSE_HEALTH_TIMEOUT); + try { + const res = await fetch( + `${endpoint}/api/team/licenses/${entitlement.subjectId}/health`, + { + signal, + headers: { + 'Content-Type': 'application/json', + 'x-validate-key': metadata.validateKey, + }, + } + ); + if (!res.ok) { + if (res.status >= 500) { + return this.remoteSelfhostFallbackEntitlement(entitlement); + } + + await this.markRemoteSelfhostLicenseNeedsReupload( + entitlement, + `Remote license health check failed: ${res.status}` + ); + return null; + } + + const payload = (await res + .json() + .catch(() => null)) as RemoteSelfhostLicense | null; + if (!payload) { + return this.remoteSelfhostFallbackEntitlement(entitlement); + } + const expiresAt = this.remoteSelfhostLicenseExpiresAt(payload.endAt); + if ( + payload.plan !== SubscriptionPlan.SelfHostedTeam || + payload.quantity < 1 || + !expiresAt + ) { + await this.markRemoteSelfhostLicenseNeedsReupload( + entitlement, + 'Remote license health payload is invalid.' + ); + return null; + } + + const validateKey = + res.headers.get('x-next-validate-key') ?? metadata.validateKey; + const [updated] = await Promise.all([ + this.db.entitlement.update({ + where: { id: entitlement.id }, + data: { + status: 'active', + quantity: this.normalizedQuantity(payload.quantity), + metadata: { + ...metadata, + recurring: payload.recurring, + validateKey, + remoteValidated: true, + errorCode: null, + errorMessage: null, + }, + expiresAt, + validatedAt: new Date(), + }, + }), + this.db.installedLicense + .updateMany({ + where: { key: entitlement.subjectId }, + data: { + quantity: this.normalizedQuantity(payload.quantity), + recurring: payload.recurring, + validateKey, + validatedAt: new Date(), + expiredAt: expiresAt, + }, + }) + .catch(() => null), + ]); + this.event.emit('entitlement.changed', { + targetType: 'workspace', + targetId: entitlement.targetId, + }); + this.remoteSelfhostLicenseCache.set(entitlement.id, { + entitlement: updated, + verifiedUntil: Date.now() + REMOTE_SELFHOST_LICENSE_REVALIDATE_INTERVAL, + }); + return updated; + } catch { + return this.remoteSelfhostFallbackEntitlement(entitlement); + } + } + + private remoteSelfhostFallbackEntitlement(entitlement: Entitlement) { + const cached = this.remoteSelfhostLicenseCache.get(entitlement.id); + if ( + !cached || + cached.verifiedUntil <= Date.now() || + !cached.entitlement.expiresAt || + cached.entitlement.expiresAt <= new Date() + ) { + return null; + } + + return cached.entitlement; + } + + private remoteSelfhostLicenseExpiresAt(endAt: unknown) { + const expiresAt = new Date(endAt as string | number | Date); + if (!Number.isFinite(expiresAt.getTime()) || expiresAt <= new Date()) { + return null; + } + return expiresAt; + } + + private async markRemoteSelfhostLicenseNeedsReupload( + entitlement: Entitlement, + reason: string + ) { + this.remoteSelfhostLicenseCache.delete(entitlement.id); + await this.db.entitlement.update({ + where: { id: entitlement.id }, + data: { + status: 'needs_reupload', + metadata: { + ...(entitlement.metadata as Record), + errorCode: 'needs_reupload', + errorMessage: reason, + }, + validatedAt: new Date(), + }, + }); + if (entitlement.targetId) { + this.event.emit('entitlement.changed', { + targetType: 'workspace', + targetId: entitlement.targetId, + }); + } + } + + private async emitEntitlementChanged(entitlement: Entitlement) { + if (!entitlement.targetId) { + return; + } + await this.event.emitAsync('entitlement.changed', { + targetType: entitlement.targetType as TargetType, + targetId: entitlement.targetId, + }); + } + + private targetTypeForPlan(plan: SubscriptionPlan | string): TargetType { + return plan === SubscriptionPlan.Team ? 'workspace' : 'user'; + } + + private entitlementPlan( + plan: SubscriptionPlan | string, + recurring: SubscriptionRecurring | string + ) { + if (plan === SubscriptionPlan.Pro) { + return recurring === SubscriptionRecurring.Lifetime + ? 'lifetime_pro' + : 'pro'; + } + if (plan === SubscriptionPlan.SelfHostedTeam) { + return 'selfhost_team'; + } + return plan; + } + + private statusFromSubscription(status: string): EntitlementStatus { + if ( + status === SubscriptionStatus.Active || + status === SubscriptionStatus.Trialing + ) { + return 'active'; + } + if (status === SubscriptionStatus.PastDue) { + return 'grace'; + } + if (status === SubscriptionStatus.Canceled) { + return 'revoked'; + } + return 'expired'; + } + + private normalizedQuantity(quantity: number | null | undefined) { + if (!quantity || quantity < 1) { + return 1; + } + return Math.min(quantity, 100000); + } +} diff --git a/packages/backend/server/src/core/features/index.ts b/packages/backend/server/src/core/features/index.ts index a36aa08e00..cb9ff5da91 100644 --- a/packages/backend/server/src/core/features/index.ts +++ b/packages/backend/server/src/core/features/index.ts @@ -1,5 +1,6 @@ import { Module } from '@nestjs/common'; +import { EntitlementModule } from '../entitlement'; import { AdminFeatureManagementResolver, UserFeatureResolver, @@ -7,6 +8,7 @@ import { import { EarlyAccessType, FeatureService } from './service'; @Module({ + imports: [EntitlementModule], providers: [ UserFeatureResolver, AdminFeatureManagementResolver, diff --git a/packages/backend/server/src/core/features/resolver.ts b/packages/backend/server/src/core/features/resolver.ts index 90d8973080..87ce795a45 100644 --- a/packages/backend/server/src/core/features/resolver.ts +++ b/packages/backend/server/src/core/features/resolver.ts @@ -1,5 +1,6 @@ import { Args, + Int, Mutation, Parent, registerEnumType, @@ -8,13 +9,10 @@ import { } from '@nestjs/graphql'; import { difference } from 'lodash-es'; -import { - Feature, - Models, - type UserFeatureName, - type WorkspaceFeatureName, -} from '../../models'; +import { BadRequest } from '../../base'; +import { Feature, Models, type UserFeatureName } from '../../models'; import { Admin } from '../common'; +import { EntitlementService } from '../entitlement'; import { UserType } from '../user/types'; import { AvailableUserFeatureConfig } from './types'; @@ -42,7 +40,10 @@ export class UserFeatureResolver extends AvailableUserFeatureConfig { @Admin() @Resolver(() => Boolean) export class AdminFeatureManagementResolver extends AvailableUserFeatureConfig { - constructor(private readonly models: Models) { + constructor( + private readonly models: Models, + private readonly entitlement: EntitlementService + ) { super(); } @@ -55,16 +56,20 @@ export class AdminFeatureManagementResolver extends AvailableUserFeatureConfig { features: UserFeatureName[] ) { const configurableUserFeatures = this.configurableUserFeatures(); + const unsupported = features.filter( + feature => !configurableUserFeatures.has(feature) + ); + if (unsupported.length) { + throw new BadRequest( + `User feature ${unsupported.join(', ')} is not configurable` + ); + } const removed = difference(Array.from(configurableUserFeatures), features); await Promise.all( - features.map(async feature => { - if (configurableUserFeatures.has(feature)) { - return this.models.userFeature.add(id, feature, 'admin panel'); - } else { - return; - } - }) + features.map(feature => + this.models.userFeature.add(id, feature, 'admin panel') + ) ); await Promise.all( @@ -75,24 +80,29 @@ export class AdminFeatureManagementResolver extends AvailableUserFeatureConfig { } @Mutation(() => Boolean) - async addWorkspaceFeature( - @Args('workspaceId') workspaceId: string, - @Args('feature', { type: () => Feature }) feature: WorkspaceFeatureName + async grantCommercialEntitlement( + @Args('targetType', { type: () => String }) + targetType: 'user' | 'workspace', + @Args('targetId', { type: () => String }) targetId: string, + @Args('plan', { type: () => String }) plan: string, + @Args('quantity', { type: () => Int, nullable: true }) quantity?: number ) { - await this.models.workspaceFeature.add( - workspaceId, - feature, - 'by administrator' - ); + await this.entitlement.upsertAdminGrant({ + targetType, + targetId, + plan, + quantity, + }); return true; } @Mutation(() => Boolean) - async removeWorkspaceFeature( - @Args('workspaceId') workspaceId: string, - @Args('feature', { type: () => Feature }) feature: WorkspaceFeatureName + async revokeCommercialEntitlement( + @Args('targetType', { type: () => String }) + targetType: 'user' | 'workspace', + @Args('targetId', { type: () => String }) targetId: string ) { - await this.models.workspaceFeature.remove(workspaceId, feature); + await this.entitlement.revokeAdminGrant(targetType, targetId); return true; } } diff --git a/packages/backend/server/src/core/features/types.ts b/packages/backend/server/src/core/features/types.ts index 9405109fd5..5b1cb0286c 100644 --- a/packages/backend/server/src/core/features/types.ts +++ b/packages/backend/server/src/core/features/types.ts @@ -5,24 +5,14 @@ import { Feature, UserFeatureName } from '../../models'; @Injectable() export class AvailableUserFeatureConfig { availableUserFeatures(): Set { - return new Set([ - Feature.Admin, - Feature.UnlimitedCopilot, - Feature.EarlyAccess, - Feature.AIEarlyAccess, - ]); + return new Set([Feature.Admin, Feature.EarlyAccess, Feature.AIEarlyAccess]); } configurableUserFeatures(): Set { return new Set( env.selfhosted - ? [Feature.Admin, Feature.UnlimitedCopilot] - : [ - Feature.EarlyAccess, - Feature.AIEarlyAccess, - Feature.Admin, - Feature.UnlimitedCopilot, - ] + ? [Feature.Admin] + : [Feature.EarlyAccess, Feature.AIEarlyAccess, Feature.Admin] ); } } diff --git a/packages/backend/server/src/core/notification/resolver.ts b/packages/backend/server/src/core/notification/resolver.ts index 41bdd6d000..36114c17a9 100644 --- a/packages/backend/server/src/core/notification/resolver.ts +++ b/packages/backend/server/src/core/notification/resolver.ts @@ -14,7 +14,7 @@ import { import { paginate, PaginationInput } from '../../base/graphql'; import { MentionNotificationCreateSchema } from '../../models'; import { CurrentUser } from '../auth/session'; -import { AccessController } from '../permission'; +import { PermissionAccess } from '../permission'; import { UserType } from '../user'; import { NotificationService } from './service'; import { @@ -28,7 +28,7 @@ import { export class UserNotificationResolver { constructor( private readonly service: NotificationService, - private readonly ac: AccessController + private readonly ac: PermissionAccess ) {} @ResolveField(() => PaginatedNotificationObjectType, { diff --git a/packages/backend/server/src/core/permission/__tests__/__snapshots__/actions.spec.ts.md b/packages/backend/server/src/core/permission/__tests__/__snapshots__/actions.spec.ts.md index 99948f14d0..a2b1b7516e 100644 --- a/packages/backend/server/src/core/permission/__tests__/__snapshots__/actions.spec.ts.md +++ b/packages/backend/server/src/core/permission/__tests__/__snapshots__/actions.spec.ts.md @@ -229,6 +229,7 @@ Generated by [AVA](https://avajs.dev). 'Doc.Comments.Delete': false, 'Doc.Comments.Read': false, 'Doc.Comments.Resolve': false, + 'Doc.Comments.Update': false, 'Doc.Copy': false, 'Doc.Delete': false, 'Doc.Duplicate': false, @@ -251,6 +252,7 @@ Generated by [AVA](https://avajs.dev). 'Doc.Comments.Delete': false, 'Doc.Comments.Read': true, 'Doc.Comments.Resolve': false, + 'Doc.Comments.Update': false, 'Doc.Copy': true, 'Doc.Delete': false, 'Doc.Duplicate': false, @@ -273,6 +275,7 @@ Generated by [AVA](https://avajs.dev). 'Doc.Comments.Delete': false, 'Doc.Comments.Read': true, 'Doc.Comments.Resolve': false, + 'Doc.Comments.Update': false, 'Doc.Copy': true, 'Doc.Delete': false, 'Doc.Duplicate': true, @@ -295,6 +298,7 @@ Generated by [AVA](https://avajs.dev). 'Doc.Comments.Delete': false, 'Doc.Comments.Read': true, 'Doc.Comments.Resolve': false, + 'Doc.Comments.Update': false, 'Doc.Copy': true, 'Doc.Delete': false, 'Doc.Duplicate': true, @@ -317,6 +321,7 @@ Generated by [AVA](https://avajs.dev). 'Doc.Comments.Delete': true, 'Doc.Comments.Read': true, 'Doc.Comments.Resolve': true, + 'Doc.Comments.Update': true, 'Doc.Copy': true, 'Doc.Delete': true, 'Doc.Duplicate': true, @@ -339,6 +344,7 @@ Generated by [AVA](https://avajs.dev). 'Doc.Comments.Delete': true, 'Doc.Comments.Read': true, 'Doc.Comments.Resolve': true, + 'Doc.Comments.Update': true, 'Doc.Copy': true, 'Doc.Delete': true, 'Doc.Duplicate': true, @@ -361,6 +367,7 @@ Generated by [AVA](https://avajs.dev). 'Doc.Comments.Delete': true, 'Doc.Comments.Read': true, 'Doc.Comments.Resolve': true, + 'Doc.Comments.Update': true, 'Doc.Copy': true, 'Doc.Delete': true, 'Doc.Duplicate': true, @@ -412,6 +419,7 @@ Generated by [AVA](https://avajs.dev). 'Doc.Comments.Delete': 'Editor', 'Doc.Comments.Read': 'External', 'Doc.Comments.Resolve': 'Editor', + 'Doc.Comments.Update': 'Editor', 'Doc.Copy': 'External', 'Doc.Delete': 'Editor', 'Doc.Duplicate': 'Reader', diff --git a/packages/backend/server/src/core/permission/__tests__/__snapshots__/actions.spec.ts.snap b/packages/backend/server/src/core/permission/__tests__/__snapshots__/actions.spec.ts.snap index f956445cbdb41e064ba50c8af43fb7c05e63fd45..dd3eb6108b183364aed1e10cc8d6bce3535b4e78 100644 GIT binary patch literal 1638 zcmV-s2ATOmRzVG5`19vI>Cs8b1k*7@Qo30!kR)T{ae|k zoAK-Eh#!jx00000000B+ThDLYL=^sJoJ|wjCRvj3LqVz~C~av8LP&@r4wP&UNI+12 z5P=ZFyE{oty`IH(nxx_Y0#wlh2Tn*FP;){^MR0)=;)(UNJrHFmuu4HFK4Fm;Aa_Vc~C(G0Y(5!F-KRIQ&vy15ix z{|vLL%!|lQXOo5akZK`CH8{m`V;Z^5<@pG|cpo9cDTwU+Xj zbt;={dev<)sY62a;}mxstIRzs;NCXEVlp3H{?Is=k8lg&@k~C(Erdu;87FevGD4UOn?k!HI-7kNUTN43aHowdjk3(N}lJXQ(JI`e9_@7vtuQ+FwUTMnH~(jDZYel+ZC^fE0@4MBCuWX6)O3 z(9D7UWVoG8932rxjm4v0ad_FEe9iWQoX|%2h0tLEIxa#NJUdWWP>iutyly*OG5u%) z=P@f_C0>yuM`C1&IZQS2!!c~yTe4jH3ezmXNi1@>}832ktIa*E@z3cdlas&gahMT$Agn)`1XNUzh*=K;v(n@yDBhmAYeB^jjensfZjGkUGnc zHnA&Z$M%;sYg12lH2m9dbQ(0C$#LZ!jBqMLn!Bi5%SaQDS>g}Hu#QuJn&gc5{?tuJ z659jah;#0;ZaPYi-PDaZ+2PX7hl;jHPdf)nI~@gcPdmG)ov8TZ(oWz(P>wWQddlf3 zXOEQQshZ!1OFSLN9Xl6~esbN-^j5_q0`^Wk(%A2=c%+`$hIpjAs9Q@{FFVQSJ$lON zDJO?=q-#6K%IP>W)F`Jef5fRLot|{|lXOJfevpnUf2!}i>!C<>??7m-6^i6o^3j8=fo)$D5>Ld7k*#jJWlX1N!DWn6 zHsmtK$=h?PJ=AvTG`pei*QsX9Y}+*kHUlBoKw*8dff-Fa%?dNyd8#R8T|Lb_+WuaK zuQR%RIV@~9{W50Kwg6S`)^!3^wy7I}qJ^yl`CLYVOuvbxN{z-ZvI6A>fLj2H1Sk>U zG6B9N!0!aOLx9s1c$WfKDexl&{-wa60X%B}?-;<0~U7bu0~yD)icl|uVnmwYvg zCQ#MYi*gcW6{8%Hu4XjWh;@yUQBvO&+Ak^Pt#fqO1$D)vw5g-=Da3D$XiFfC9UaAx k-u3xaEYwCyTCqk_0_C-mdJAMlQ|!Y30fcq8Ig>d60L{`2`v3p{ literal 1615 zcmV-V2C(@-RzVJ5 zLPA~JlVs_7$Lwy%48{Vgz%(r}-tzPw(__AF#ip(&JHZy!~ELf|{ zRin&HhRbbcEV#U4tZ?^+S2as4_|0_64S?l2X4PBbzW2qGo8_|Uo6|WQL+wNID3G4l zxx<)k7MZbZ2P?Vkd=*#uqh zG&9T04auHOCkyZa)qH~L_yn(1nBzCt&Zo0AhDdF@RE@m<7 zY&zG}inGn64hhkZ6Wq4VB6m%nyE_Pr@oaS2L*qg=!WP2gg=~y1gh)mk<6$zFpgR`2 zJ>+{n8($(;W)g(?y6+WwP|(G~MQh#D0^QYe#d4&fqNNc7SfojYh?eS^6smY|Xo*Cv zi$p?9?8>Oj(I6&mktk~zoKE41`Cr2#dcS^6>U|A~=zEbQ>3a>4xXB_>-18b5(eEOS z;P+h^6!*G@M)bKzvs0goJb}-{TIA*CT4a&=Mv1#FEBQv1xfRRvEbe$gQIg<}%8*Y2 ztKROR4ZHft==#=|n*hECa38=)0$e1(bpm`&fbR(K2LT2saD@Uc1vV-03kCkAz}Xxy zp94P30pD(Y|1Ad$limX9ohrV^N|35+ z&w<`}u$^@r9TG;3M5A7@dC?nx)$;s|&|3Hf&_Mw@DnjR7%U4)Xh_Dm9YS~;d{csHD zGSg?1yd*~sMaW5JGu6b8MX)J%(R8eJrdiVC5&8{tt)>_4zK+Gc6kO>(TyNTieB?=MUmPDa3VGk3< zU|6q=37H~(tJ2LD-E7g0Y>`dN@dapv*QRVCj@XB6p)p}kvV|t<71_d<-OxVVSt3>X zhc`=TLTmR~BAX(+Sz*c<{SEFEZ;op0Pq2P1>i1# zzW|I8;4%T;A;3ojxI=(n3GfdAPEp_@1+G)zCIvPqaF?d$n1ZOmNnRQc&#LqlauTWg zvo0i%>buqv0q1f10|jMY0IjacFF#VdtTU;<4p^?*R!P4=sq94LsDRW}&?oBH<)UqQ zOPaN*Cp!^b_G-NY&3X#K{hq65dy|F1l*iXEijHsa-X6Pz_sC2c^)xxf{;HsL) zher-gM+Uo>g9vR^4kBRh^U?MM35C*MMAC|@8a1+Mad|s zRPv3NWwKT*x(u(bU6eL8)iC+!?FDTyqp_o@tkF9%yB>uqM@cJIJ4&Fu>QQfjto(^w N_&<40ui_gx002Xs3p@Y- diff --git a/packages/backend/server/src/core/permission/__tests__/doc.spec.ts b/packages/backend/server/src/core/permission/__tests__/doc.spec.ts index 1b7a682361..76a0acc53e 100644 --- a/packages/backend/server/src/core/permission/__tests__/doc.spec.ts +++ b/packages/backend/server/src/core/permission/__tests__/doc.spec.ts @@ -10,14 +10,13 @@ import { WorkspaceMemberStatus, WorkspaceRole, } from '../../../models'; -import { DocAccessController } from '../doc'; -import { PermissionModule } from '../index'; +import { PermissionAccess, PermissionModule } from '../index'; import { WorkspacePolicyService } from '../policy'; import { DocRole, mapDocRoleToPermissions } from '../types'; let module: TestingModule; let models: Models; -let ac: DocAccessController; +let ac: PermissionAccess; let policy: WorkspacePolicyService; let user: User; let ws: Workspace; @@ -26,7 +25,7 @@ let underReviewUserId: string; test.before(async () => { module = await createTestingModule({ imports: [PermissionModule] }); models = module.get(Models); - ac = module.get(DocAccessController); + ac = module.get(PermissionAccess); policy = module.get(WorkspacePolicyService); }); @@ -40,6 +39,21 @@ test.after.always(async () => { await module.close(); }); +function doc(resource: { + workspaceId: string; + docId: string; + userId: string; + allowLocal?: boolean; +}) { + const checker = ac + .user(resource.userId) + .doc(resource.workspaceId, resource.docId); + if (resource.allowLocal) { + checker.allowLocal(); + } + return checker; +} + const roleCases: Array<{ title: string; setup?: () => Promise; @@ -90,7 +104,7 @@ const roleCases: Array<{ expectedRole: DocRole.Owner, }, { - title: 'should fallback to [External] if workspace is public', + title: 'should not grant private doc role if workspace is public', setup: async () => { await models.workspace.update(ws.id, { public: true, @@ -101,7 +115,7 @@ const roleCases: Array<{ docId: 'doc1', userId: 'random-user-id', }), - expectedRole: DocRole.External, + expectedRole: null, }, { title: 'should return null even if workspace has other public doc', @@ -131,9 +145,13 @@ const roleCases: Array<{ title: 'should return null if doc role is [None]', setup: async () => { await models.doc.setDefaultRole(ws.id, 'doc1', DocRole.None); + const u2 = await models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + underReviewUserId = u2.id; await models.workspaceUser.set( ws.id, - user.id, + underReviewUserId, WorkspaceRole.Collaborator, { status: WorkspaceMemberStatus.Accepted, @@ -143,7 +161,7 @@ const roleCases: Array<{ resource: () => ({ workspaceId: ws.id, docId: 'doc1', - userId: user.id, + userId: underReviewUserId, }), expectedRole: null, }, @@ -151,14 +169,6 @@ const roleCases: Array<{ title: 'should return [External] if doc role is [None] but doc is public', setup: async () => { await models.doc.setDefaultRole(ws.id, 'doc1', DocRole.None); - await models.workspaceUser.set( - ws.id, - user.id, - WorkspaceRole.Collaborator, - { - status: WorkspaceMemberStatus.Accepted, - } - ); await models.doc.publish(ws.id, 'doc1'); }, resource: () => ({ @@ -174,18 +184,18 @@ for (const roleCase of roleCases) { test(roleCase.title, async t => { await roleCase.setup?.(); const resource = roleCase.resource(); - const role = await ac.getRole(resource); + const role = (await doc(resource).permissions()).role; t.is(role, roleCase.expectedRole); }); } test('should return mapped permissions', async t => { - const { permissions } = await ac.role({ + const { permissions } = await doc({ workspaceId: ws.id, docId: 'doc1', userId: user.id, - }); + }).permissions(); t.deepEqual(permissions, mapDocRoleToPermissions(DocRole.Owner)); }); @@ -195,11 +205,11 @@ test('should deny publish permission when workspace sharing is disabled', async enableSharing: false, }); - const { permissions } = await ac.role({ + const { permissions } = await doc({ workspaceId: ws.id, docId: 'doc1', userId: user.id, - }); + }).permissions(); t.false(permissions['Doc.Publish']); t.true(permissions['Doc.Read']); @@ -211,24 +221,18 @@ test('should deny publish assert when workspace sharing is disabled', async t => }); await t.throwsAsync( - ac.assert( - { - workspaceId: ws.id, - docId: 'doc1', - userId: user.id, - }, - 'Doc.Publish' - ) + doc({ + workspaceId: ws.id, + docId: 'doc1', + userId: user.id, + }).assert('Doc.Publish') ); await t.notThrowsAsync( - ac.assert( - { - workspaceId: ws.id, - docId: 'doc1', - userId: user.id, - }, - 'Doc.Read' - ) + doc({ + workspaceId: ws.id, + docId: 'doc1', + userId: user.id, + }).assert('Doc.Read') ); }); @@ -239,34 +243,27 @@ test('should deny external read assert when sharing is disabled even if doc is p }); await t.throwsAsync( - ac.assert( - { - workspaceId: ws.id, - docId: 'doc1', - userId: 'random-user-id', - }, - 'Doc.Read' - ) + doc({ + workspaceId: ws.id, + docId: 'doc1', + userId: 'random-user-id', + }).assert('Doc.Read') ); }); test('should assert action', async t => { await t.notThrowsAsync( - ac.assert( - { - workspaceId: ws.id, - docId: 'doc1', - userId: user.id, - }, - 'Doc.Update' - ) + doc({ + workspaceId: ws.id, + docId: 'doc1', + userId: user.id, + }).assert('Doc.Update') ); const u2 = await models.user.create({ email: `${randomUUID()}@affine.pro` }); await t.throwsAsync( - ac.assert( - { workspaceId: ws.id, docId: 'doc1', userId: u2.id }, + doc({ workspaceId: ws.id, docId: 'doc1', userId: u2.id }).assert( 'Doc.Update' ) ); @@ -278,8 +275,7 @@ test('should assert action', async t => { await models.docUser.set(ws.id, 'doc1', u2.id, DocRole.Manager); await t.notThrowsAsync( - ac.assert( - { workspaceId: ws.id, docId: 'doc1', userId: u2.id }, + doc({ workspaceId: ws.id, docId: 'doc1', userId: u2.id }).assert( 'Doc.Delete' ) ); @@ -301,11 +297,11 @@ test('should apply readonly doc restrictions while keeping cleanup actions', asy } await policy.reconcileWorkspaceQuotaState(ws.id); - const { permissions } = await ac.role({ + const { permissions } = await doc({ workspaceId: ws.id, docId: 'doc1', userId: user.id, - }); + }).permissions(); t.false(permissions['Doc.Update']); t.false(permissions['Doc.Publish']); diff --git a/packages/backend/server/src/core/permission/__tests__/docs.spec.ts b/packages/backend/server/src/core/permission/__tests__/docs.spec.ts index d10bed8583..03eed70017 100644 --- a/packages/backend/server/src/core/permission/__tests__/docs.spec.ts +++ b/packages/backend/server/src/core/permission/__tests__/docs.spec.ts @@ -1,20 +1,84 @@ +import { randomUUID } from 'node:crypto'; + +import { Prisma, PrismaClient } from '@prisma/client'; import test from 'ava'; import { createModule } from '../../../__tests__/create-module'; import { Mockers } from '../../../__tests__/mocks'; +import { Models } from '../../../models'; import { AccessControllerBuilder } from '../builder'; +import { PermissionDiagnosticService } from '../diagnostic'; import { DocRole, PermissionModule, WorkspaceRole } from '../index'; +import { PermissionSqlPredicateBuilder } from '../sql-predicate'; +import type { DocAction } from '../types'; const module = await createModule({ imports: [PermissionModule], }); const builder = module.get(AccessControllerBuilder); +const models = module.get(Models); +const db = module.get(PrismaClient); +const diagnostic = module.get(PermissionDiagnosticService); +const sqlPredicate = module.get(PermissionSqlPredicateBuilder); test.after.always(async () => { await module.close(); }); +async function sqlReadableDocIds(input: { + workspaceId: string; + userId?: string; + action?: DocAction; + docIds: string[]; +}) { + const values = Prisma.join( + input.docIds.map((docId, index) => Prisma.sql`(${docId}, ${index})`) + ); + const predicate = sqlPredicate.docReadableByNewTablesSql({ + workspaceId: input.workspaceId, + userId: input.userId, + action: input.action ?? 'Doc.Read', + docIdColumn: Prisma.raw('c.doc_id'), + }); + const rows = await db.$queryRaw<{ docId: string }[]>` + WITH candidates(doc_id, ord) AS (VALUES ${values}) + SELECT c.doc_id AS "docId" + FROM candidates c + WHERE ${predicate} + ORDER BY c.ord ASC + `; + return rows.map(row => row.docId); +} + +async function resetProjection(workspaceId: string) { + await db.$executeRaw`DELETE FROM doc_grants WHERE workspace_id = ${workspaceId}`; + await db.$executeRaw`DELETE FROM doc_access_policies WHERE workspace_id = ${workspaceId}`; + await db.$executeRaw`DELETE FROM workspace_members WHERE workspace_id = ${workspaceId}`; + await db.$executeRaw` + INSERT INTO workspace_access_policies ( + workspace_id, + visibility, + sharing_enabled, + url_preview_enabled, + member_default_doc_role, + updated_at + ) + VALUES (${workspaceId}, 'private', true, false, 'none', now()) + ON CONFLICT (workspace_id) + DO UPDATE SET + visibility = EXCLUDED.visibility, + sharing_enabled = EXCLUDED.sharing_enabled, + url_preview_enabled = EXCLUDED.url_preview_enabled, + member_default_doc_role = EXCLUDED.member_default_doc_role, + updated_at = now() + `; + await models.workspaceRuntimeState.upsert(workspaceId, { + readonly: false, + readonlyReasons: [], + }); +} + test('should filter docs by Doc.Read', async t => { const owner = await module.create(Mockers.User); const workspace = await module.create(Mockers.Workspace, { @@ -79,11 +143,329 @@ test('should filter docs by Doc.Read', async t => { t.is(docs3.length, 0); }); +test('SQL doc read predicate matches Rust for projection default and public candidates', async t => { + const owner = await module.create(Mockers.User); + const member = await module.create(Mockers.User); + const workspace = await module.create(Mockers.Workspace, { + owner, + }); + await resetProjection(workspace.id); + await db.$executeRaw` + UPDATE workspace_access_policies + SET member_default_doc_role = 'reader' + WHERE workspace_id = ${workspace.id} + `; + await db.$executeRaw` + INSERT INTO workspace_members ( + workspace_id, + user_id, + role, + state, + source, + updated_at + ) + VALUES (${workspace.id}, ${member.id}, 'member', 'active', 'legacy', now()) + `; + await db.$executeRaw` + INSERT INTO doc_access_policies ( + workspace_id, + doc_id, + visibility, + public_role, + member_default_role, + updated_at + ) + VALUES + (${workspace.id}, 'member-default-none', 'private', NULL, 'none', now()), + (${workspace.id}, 'public-doc', 'public', 'external', NULL, now()) + `; + + const docIds = ['missing-policy', 'member-default-none', 'public-doc']; + const sqlReadable = await sqlReadableDocIds({ + workspaceId: workspace.id, + userId: member.id, + docIds, + }); + const shadow = await diagnostic.shadowSqlDocRead({ + workspaceId: workspace.id, + userId: member.id, + docs: docIds.map(docId => ({ docId })), + sqlReadableDocIds: sqlReadable, + }); + + t.deepEqual(sqlReadable, ['missing-policy', 'public-doc']); + t.true(shadow.matched); +}); + +test('SQL doc read predicate matches Rust for non-member grant and sharing disabled', async t => { + const owner = await module.create(Mockers.User); + const nonMember = await module.create(Mockers.User); + const workspace = await module.create(Mockers.Workspace, { + owner, + }); + await resetProjection(workspace.id); + await db.$executeRaw` + INSERT INTO doc_access_policies ( + workspace_id, + doc_id, + visibility, + public_role, + member_default_role, + updated_at + ) + VALUES + (${workspace.id}, 'public-doc', 'public', 'external', NULL, now()), + (${workspace.id}, 'private-doc', 'private', NULL, NULL, now()), + (${workspace.id}, 'explicit-grant', 'private', NULL, NULL, now()), + (${workspace.id}, 'explicit-owner-grant', 'private', NULL, NULL, now()) + `; + await db.$executeRaw` + INSERT INTO doc_grants ( + workspace_id, + doc_id, + principal_type, + principal_id, + role, + updated_at + ) + VALUES + ( + ${workspace.id}, + 'explicit-grant', + 'user', + ${nonMember.id}, + 'reader', + now() + ), + ( + ${workspace.id}, + 'explicit-owner-grant', + 'user', + ${nonMember.id}, + 'owner', + now() + ) + `; + + const docIds = [ + 'public-doc', + 'private-doc', + 'explicit-grant', + 'explicit-owner-grant', + ]; + const sharingEnabledReadable = await sqlReadableDocIds({ + workspaceId: workspace.id, + userId: nonMember.id, + docIds, + }); + const sharingEnabledShadow = await diagnostic.shadowSqlDocRead({ + workspaceId: workspace.id, + userId: nonMember.id, + docs: docIds.map(docId => ({ docId })), + sqlReadableDocIds: sharingEnabledReadable, + }); + const sharingEnabledUpdate = await sqlReadableDocIds({ + workspaceId: workspace.id, + userId: nonMember.id, + action: 'Doc.Update', + docIds, + }); + + await db.$executeRaw` + UPDATE workspace_access_policies + SET sharing_enabled = false + WHERE workspace_id = ${workspace.id} + `; + const sharingDisabledReadable = await sqlReadableDocIds({ + workspaceId: workspace.id, + userId: nonMember.id, + docIds, + }); + const sharingDisabledShadow = await diagnostic.shadowSqlDocRead({ + workspaceId: workspace.id, + userId: nonMember.id, + docs: docIds.map(docId => ({ docId })), + sqlReadableDocIds: sharingDisabledReadable, + }); + + t.deepEqual(sharingEnabledReadable, [ + 'public-doc', + 'explicit-grant', + 'explicit-owner-grant', + ]); + t.true(sharingEnabledShadow.matched); + t.deepEqual(sharingEnabledUpdate, ['explicit-owner-grant']); + t.deepEqual(sharingDisabledReadable, []); + t.true(sharingDisabledShadow.matched); +}); + +test('SQL doc predicate suppresses member default when explicit grant exists', async t => { + const owner = await module.create(Mockers.User); + const member = await module.create(Mockers.User); + const workspace = await module.create(Mockers.Workspace, { + owner, + }); + await resetProjection(workspace.id); + await db.$executeRaw` + UPDATE workspace_access_policies + SET member_default_doc_role = 'manager' + WHERE workspace_id = ${workspace.id} + `; + await db.$executeRaw` + INSERT INTO workspace_members ( + workspace_id, + user_id, + role, + state, + source, + updated_at + ) + VALUES (${workspace.id}, ${member.id}, 'member', 'active', 'legacy', now()) + `; + await db.$executeRaw` + INSERT INTO doc_access_policies ( + workspace_id, + doc_id, + visibility, + public_role, + member_default_role, + updated_at + ) + VALUES + (${workspace.id}, 'default-manager', 'private', NULL, NULL, now()), + (${workspace.id}, 'explicit-reader', 'private', NULL, NULL, now()) + `; + await db.$executeRaw` + INSERT INTO doc_grants ( + workspace_id, + doc_id, + principal_type, + principal_id, + role, + updated_at + ) + VALUES ( + ${workspace.id}, + 'explicit-reader', + 'user', + ${member.id}, + 'reader', + now() + ) + `; + + const docIds = ['default-manager', 'explicit-reader']; + const sqlUpdateAllowed = await sqlReadableDocIds({ + workspaceId: workspace.id, + userId: member.id, + action: 'Doc.Update', + docIds, + }); + + t.deepEqual(sqlUpdateAllowed, ['default-manager']); +}); + +test('legacy SQL doc predicate matches external row and explicit grant cap semantics', async t => { + const workspaceId = randomUUID(); + const memberId = randomUUID(); + const externalId = randomUUID(); + + async function fixtureLegacyDocIds(input: { + userId: string; + action: DocAction; + docIds: string[]; + }) { + const values = Prisma.join( + input.docIds.map((docId, index) => Prisma.sql`(${docId}, ${index})`) + ); + const predicate = sqlPredicate.docReadableByLegacyTablesSql({ + workspaceId, + userId: input.userId, + action: input.action, + docIdColumn: Prisma.raw('c.doc_id'), + }); + // Current triggers reject newly inserted legacy External workspace rows; + // CTEs let the same predicate run in Postgres against historical shapes. + const rows = await db.$queryRaw<{ docId: string }[]>` + WITH + workspaces(id, enable_sharing) AS ( + VALUES (${workspaceId}, true) + ), + workspace_pages(workspace_id, page_id, public, "defaultRole") AS ( + VALUES + (${workspaceId}, 'default-manager', false, ${DocRole.Manager}::smallint), + (${workspaceId}, 'explicit-reader', false, ${DocRole.Manager}::smallint), + (${workspaceId}, 'external-owner', false, ${DocRole.Manager}::smallint), + (${workspaceId}, 'dirty-external', false, ${DocRole.Manager}::smallint) + ), + workspace_user_permissions( + id, + workspace_id, + user_id, + status, + type + ) AS ( + VALUES + (${randomUUID()}, ${workspaceId}, ${memberId}, 'Accepted'::"WorkspaceMemberStatus", ${WorkspaceRole.Collaborator}::smallint), + (${randomUUID()}, ${workspaceId}, ${externalId}, 'Accepted'::"WorkspaceMemberStatus", ${WorkspaceRole.External}::smallint) + ), + workspace_page_user_permissions( + workspace_id, + page_id, + user_id, + type + ) AS ( + VALUES + (${workspaceId}, 'explicit-reader', ${memberId}, ${DocRole.Reader}::smallint), + (${workspaceId}, 'external-owner', ${externalId}, ${DocRole.Owner}::smallint), + (${workspaceId}, 'dirty-external', ${externalId}, ${DocRole.External}::smallint) + ), + candidates(doc_id, ord) AS (VALUES ${values}) + SELECT c.doc_id AS "docId" + FROM candidates c + WHERE ${predicate} + ORDER BY c.ord ASC + `; + return rows.map(row => row.docId); + } + + const memberUpdateAllowed = await fixtureLegacyDocIds({ + userId: memberId, + action: 'Doc.Update', + docIds: ['default-manager', 'explicit-reader'], + }); + const externalUpdateAllowed = await fixtureLegacyDocIds({ + userId: externalId, + action: 'Doc.Update', + docIds: ['external-owner', 'dirty-external'], + }); + const externalManageAllowed = await fixtureLegacyDocIds({ + userId: externalId, + action: 'Doc.Users.Manage', + docIds: ['external-owner', 'dirty-external'], + }); + const externalTransferAllowed = await fixtureLegacyDocIds({ + userId: externalId, + action: 'Doc.TransferOwner', + docIds: ['external-owner', 'dirty-external'], + }); + + t.deepEqual(memberUpdateAllowed, ['default-manager']); + t.deepEqual(externalUpdateAllowed, ['external-owner']); + t.deepEqual(externalManageAllowed, []); + t.deepEqual(externalTransferAllowed, []); +}); + test('should filter docs by Doc.Publish', async t => { const owner = await module.create(Mockers.User); const workspace = await module.create(Mockers.Workspace, { owner, }); + await models.workspace.update(workspace.id, { enableSharing: true }); + await models.workspaceRuntimeState.upsert(workspace.id, { + readonly: false, + readonlyReasons: [], + }); const docs1 = await builder .user(owner.id) diff --git a/packages/backend/server/src/core/permission/__tests__/policy.spec.ts b/packages/backend/server/src/core/permission/__tests__/policy.spec.ts index c6262bea32..c4ca002967 100644 --- a/packages/backend/server/src/core/permission/__tests__/policy.spec.ts +++ b/packages/backend/server/src/core/permission/__tests__/policy.spec.ts @@ -1,5 +1,6 @@ import { randomUUID } from 'node:crypto'; +import { PrismaClient } from '@prisma/client'; import ava, { TestFn } from 'ava'; import Sinon from 'sinon'; @@ -7,11 +8,6 @@ import { createTestingModule, type TestingModule, } from '../../../__tests__/utils'; -import { - DocActionDenied, - OwnerCanNotLeaveWorkspace, - SpaceAccessDenied, -} from '../../../base'; import { Models, User, @@ -19,25 +15,59 @@ import { WorkspaceMemberStatus, WorkspaceRole, } from '../../../models'; -import { QuotaService } from '../../quota/service'; import { QuotaServiceModule } from '../../quota/service.module'; +import { QuotaStateService } from '../../quota/state'; import { PermissionModule } from '../index'; import { WorkspacePolicyService } from '../policy'; interface Context { module: TestingModule; + db: PrismaClient; models: Models; policy: WorkspacePolicyService; } const test = ava as TestFn; -const READONLY_FEATURE = 'quota_exceeded_readonly_workspace_v1' as const; type WorkspaceQuotaSnapshot = Awaited< - ReturnType + ReturnType > & { - ownerQuota?: string; + readonlyReasons: string[]; }; + +const readonlyWorkspaceState = ( + workspaceId: string, + readonlyReasons: string[], + overrides: Partial = {} +) => + ({ + workspaceId, + plan: 'free', + sourceEntitlementId: null, + ownerUserId: owner.id, + usesOwnerQuota: true, + seatLimit: 3, + memberCount: 1, + overcapacityMemberCount: readonlyReasons.includes('member_overflow') + ? 1 + : 0, + blobLimit: BigInt(1), + storageQuota: BigInt(1), + usedStorageQuota: readonlyReasons.includes('storage_overflow') + ? BigInt(2) + : BigInt(0), + historyPeriodSeconds: 1, + readonly: readonlyReasons.length > 0, + readonlyReasons, + flags: {}, + known: true, + stale: false, + lastReconciledAt: new Date(), + staleAfter: new Date(), + createdAt: new Date(), + updatedAt: new Date(), + ...overrides, + }) satisfies WorkspaceQuotaSnapshot; async function addAcceptedMembers( models: Models, workspaceId: string, @@ -64,6 +94,7 @@ let workspace: Workspace; test.before(async t => { const module = await createTestingModule({ imports: [PermissionModule] }); t.context.module = module; + t.context.db = module.get(PrismaClient); t.context.models = module.get(Models); t.context.policy = module.get(WorkspacePolicyService); }); @@ -81,21 +112,23 @@ test.after.always(async t => { await t.context.module.close(); }); -test('should reuse quota service exported by quota service module', async t => { +test('should reuse quota state service exported by quota service module', async t => { const module = await createTestingModule( { imports: [PermissionModule, QuotaServiceModule] }, false ); try { - const quota = module.select(QuotaServiceModule).get(QuotaService, { - strict: true, - }); + const quotaState = module + .select(QuotaServiceModule) + .get(QuotaStateService, { + strict: true, + }); const policy = module.select(PermissionModule).get(WorkspacePolicyService, { strict: true, }); - t.is(Reflect.get(policy, 'quota'), quota); + t.is(Reflect.get(policy, 'quotaState'), quotaState); } finally { await module.close(); } @@ -108,12 +141,9 @@ test('should keep owned workspace writable when quota is within limit', async t t.false(state.isReadonly); t.deepEqual(state.readonlyReasons, []); - t.false( - await t.context.models.workspaceFeature.has(workspace.id, READONLY_FEATURE) - ); }); -test('should enter readonly mode when fallback owner member quota overflows', async t => { +test('should report readonly state when fallback owner member quota overflows', async t => { await addAcceptedMembers(t.context.models, workspace.id, 10); const state = await t.context.policy.reconcileWorkspaceQuotaState( @@ -124,91 +154,16 @@ test('should enter readonly mode when fallback owner member quota overflows', as t.true(state.canRecoverByRemovingMembers); t.false(state.canRecoverByDeletingBlobs); t.deepEqual(state.readonlyReasons, ['member_overflow']); - t.true( - await t.context.models.workspaceFeature.has(workspace.id, READONLY_FEATURE) - ); - await t.throwsAsync(t.context.policy.assertCanInviteMembers(workspace.id), { - instanceOf: SpaceAccessDenied, - }); -}); - -test('should deny blob uploads when user no longer has write access', async t => { - const external = await t.context.models.user.create({ - email: `${randomUUID()}@affine.pro`, - }); - await t.context.models.workspaceUser.set( - workspace.id, - external.id, - WorkspaceRole.External, - { status: WorkspaceMemberStatus.Accepted } - ); - - await t.throwsAsync( - t.context.policy.assertCanUploadBlob(external.id, workspace.id), - { instanceOf: SpaceAccessDenied } - ); -}); - -test('should deny publish through policy when workspace sharing is disabled', async t => { - await t.context.models.workspace.update(workspace.id, { - enableSharing: false, - }); - - await t.throwsAsync( - t.context.policy.assertCanPublishDoc(owner.id, workspace.id, 'doc1'), - { instanceOf: DocActionDenied } - ); - await t.notThrowsAsync( - t.context.policy.assertCanUnpublishDoc(owner.id, workspace.id, 'doc1') - ); -}); - -test('should allow managers to revoke invite links in readonly workspace', async t => { - await addAcceptedMembers(t.context.models, workspace.id, 10); - await t.context.policy.reconcileWorkspaceQuotaState(workspace.id); - - await t.notThrowsAsync( - t.context.policy.assertCanManageInviteLink(owner.id, workspace.id) - ); -}); - -test('should apply leave workspace policy by role', async t => { - const collaborator = await t.context.models.user.create({ - email: `${randomUUID()}@affine.pro`, - }); - await t.context.models.workspaceUser.set( - workspace.id, - collaborator.id, - WorkspaceRole.Collaborator, - { status: WorkspaceMemberStatus.Accepted } - ); - - await t.throwsAsync( - t.context.policy.assertCanLeaveWorkspace(owner.id, workspace.id), - { instanceOf: OwnerCanNotLeaveWorkspace } - ); - await t.notThrowsAsync( - t.context.policy.assertCanLeaveWorkspace(collaborator.id, workspace.id) - ); }); test('should enter readonly mode when fallback owner storage quota overflows', async t => { - const quota = Sinon.stub( - Reflect.get(t.context.policy, 'quota') as QuotaService, - 'getWorkspaceQuotaWithUsage' + const quotaState = Sinon.stub( + Reflect.get(t.context.policy, 'quotaState') as QuotaStateService, + 'reconcileWorkspaceQuotaState' + ); + quotaState.callsFake(async workspaceId => + readonlyWorkspaceState(workspaceId, ['storage_overflow']) ); - quota.resolves({ - name: 'Free', - blobLimit: 1, - storageQuota: 1, - usedStorageQuota: 2, - historyPeriod: 1, - memberLimit: 3, - memberCount: 1, - overcapacityMemberCount: 0, - usedSize: 2, - ownerQuota: owner.id, - } satisfies WorkspaceQuotaSnapshot); const state = await t.context.policy.reconcileWorkspaceQuotaState( workspace.id @@ -218,57 +173,26 @@ test('should enter readonly mode when fallback owner storage quota overflows', a t.false(state.canRecoverByRemovingMembers); t.true(state.canRecoverByDeletingBlobs); t.deepEqual(state.readonlyReasons, ['storage_overflow']); - t.true( - await t.context.models.workspaceFeature.has(workspace.id, READONLY_FEATURE) - ); }); -test('should leave readonly mode after workspace usage recovers', async t => { - const quota = Sinon.stub( - Reflect.get(t.context.policy, 'quota') as QuotaService, - 'getWorkspaceQuotaWithUsage' +test('should report recovered state after workspace usage recovers', async t => { + const quotaState = Sinon.stub( + Reflect.get(t.context.policy, 'quotaState') as QuotaStateService, + 'reconcileWorkspaceQuotaState' ); - quota.onFirstCall().resolves({ - name: 'Free', - blobLimit: 1, - storageQuota: 1, - usedStorageQuota: 2, - historyPeriod: 1, - memberLimit: 3, - memberCount: 1, - overcapacityMemberCount: 0, - usedSize: 2, - ownerQuota: owner.id, - } satisfies WorkspaceQuotaSnapshot); - quota.onSecondCall().resolves({ - name: 'Free', - blobLimit: 1, - storageQuota: 1, - usedStorageQuota: 0, - historyPeriod: 1, - memberLimit: 3, - memberCount: 1, - overcapacityMemberCount: 0, - usedSize: 0, - ownerQuota: owner.id, - } satisfies WorkspaceQuotaSnapshot); - quota.onThirdCall().resolves({ - name: 'Free', - blobLimit: 1, - storageQuota: 1, - usedStorageQuota: 0, - historyPeriod: 1, - memberLimit: 3, - memberCount: 1, - overcapacityMemberCount: 0, - usedSize: 0, - ownerQuota: owner.id, - } satisfies WorkspaceQuotaSnapshot); + quotaState + .onFirstCall() + .callsFake(async workspaceId => + readonlyWorkspaceState(workspaceId, ['storage_overflow']) + ); + quotaState + .onSecondCall() + .callsFake(async workspaceId => readonlyWorkspaceState(workspaceId, [])); + quotaState + .onThirdCall() + .callsFake(async workspaceId => readonlyWorkspaceState(workspaceId, [])); await t.context.policy.reconcileWorkspaceQuotaState(workspace.id); - t.true( - await t.context.models.workspaceFeature.has(workspace.id, READONLY_FEATURE) - ); const recovered = await t.context.policy.reconcileWorkspaceQuotaState( workspace.id @@ -276,10 +200,6 @@ test('should leave readonly mode after workspace usage recovers', async t => { t.false(recovered.isReadonly); t.deepEqual(recovered.readonlyReasons, []); - t.false( - await t.context.models.workspaceFeature.has(workspace.id, READONLY_FEATURE) - ); - await t.notThrowsAsync(t.context.policy.assertCanInviteMembers(workspace.id)); }); test('should roll back team cancellation cleanup when cleanup fails', async t => { @@ -289,11 +209,58 @@ test('should roll back team cancellation cleanup when cleanup fails', async t => const admin = await t.context.models.user.create({ email: `${randomUUID()}@affine.pro`, }); - await t.context.models.workspaceUser.set( - workspace.id, - pending.id, - WorkspaceRole.Collaborator - ); + await t.context.db.$transaction(async db => { + await db.$executeRaw` + SELECT set_config('affine.permission_projection.enabled', 'off', true) + `; + const pendingPermission = await db.workspaceUserRole.create({ + data: { + workspaceId: workspace.id, + userId: pending.id, + type: WorkspaceRole.Collaborator, + status: WorkspaceMemberStatus.Pending, + }, + }); + const [invitationShape] = await db.$queryRaw>` + SELECT EXISTS ( + SELECT 1 + FROM information_schema.columns + WHERE table_name = 'workspace_invitations' + AND column_name = 'requested_role' + ) AS "current" + `; + if (invitationShape?.current) { + await db.workspaceInvitation.create({ + data: { + workspaceId: workspace.id, + inviteeUserId: pending.id, + requestedRole: 'member', + status: 'pending', + kind: 'email', + legacyPermissionId: pendingPermission.id, + }, + }); + } else { + await db.$executeRaw` + INSERT INTO workspace_invitations ( + workspace_id, + invitee_user_id, + role, + state, + source, + updated_at + ) + VALUES ( + ${workspace.id}, + ${pending.id}, + ${'member'}, + ${'pending'}, + ${'email'}, + now() + ) + `; + } + }); await t.context.models.workspaceUser.set( workspace.id, admin.id, diff --git a/packages/backend/server/src/core/permission/__tests__/service.spec.ts b/packages/backend/server/src/core/permission/__tests__/service.spec.ts new file mode 100644 index 0000000000..95bc630b95 --- /dev/null +++ b/packages/backend/server/src/core/permission/__tests__/service.spec.ts @@ -0,0 +1,1313 @@ +import { WorkspaceMemberStatus } from '@prisma/client'; +import test from 'ava'; + +import { InternalServerError } from '../../../base'; +import { + DocRole, + WorkspaceRole, + type WorkspaceRuntimeState, +} from '../../../models'; +import { PermissionReadModel } from '../config'; +import { docLegacyBoundary } from '../context'; +import { PermissionContextLoader } from '../context-loader'; +import { + PERMISSION_SHADOW_MISMATCH_CATEGORIES, + PermissionDiagnosticService, +} from '../diagnostic'; +import { PermissionService } from '../service'; +import { PermissionSqlPredicateBuilder } from '../sql-predicate'; + +function createCls() { + const store = new Map(); + return { + get(key: string) { + return store.get(key); + }, + set(key: string, value: unknown) { + store.set(key, value); + }, + }; +} + +function createLoader( + runtimeState: WorkspaceRuntimeState = { + workspaceId: 'w1', + known: true, + stale: false, + readonly: false, + readonlyReasons: [], + updatedAt: new Date(), + lastReconciledAt: new Date(), + staleAfter: null, + }, + docGrantRole: DocRole | null = null, + queryRaw: (query: unknown) => Promise = async () => [] +) { + const calls = { + workspaceUser: 0, + workspace: 0, + runtime: 0, + docPolicies: 0, + docGrants: 0, + }; + const models = { + workspaceUser: { + get: async () => { + calls.workspaceUser += 1; + return { + type: WorkspaceRole.Owner, + status: WorkspaceMemberStatus.Accepted, + }; + }, + }, + workspace: { + get: async () => { + calls.workspace += 1; + return { + public: false, + enableSharing: true, + enableUrlPreview: true, + }; + }, + }, + workspaceRuntimeState: { + get: async () => { + calls.runtime += 1; + return runtimeState; + }, + }, + doc: { + findDefaultRoles: async (_workspaceId: string, docIds: string[]) => { + calls.docPolicies += 1; + return docIds.map(docId => ({ + workspace: docId === 'private' ? DocRole.None : DocRole.Manager, + external: docId === 'public' ? DocRole.External : null, + })); + }, + }, + docUser: { + findMany: async () => { + calls.docGrants += 1; + return docGrantRole === null + ? [] + : [ + { + docId: 'private', + type: docGrantRole, + }, + ]; + }, + }, + }; + return { + calls, + loader: new PermissionContextLoader( + models as never, + { $queryRaw: queryRaw } as never, + createCls() as never + ), + }; +} + +test('PermissionService maps native decisions to legacy role boundary', async t => { + const { loader } = createLoader(); + const service = new PermissionService(loader, undefined, undefined, { + permission: { + readModel: PermissionReadModel.Legacy, + fallbackLegacyLoader: false, + }, + } as never); + + const permissions = await service.docPermissions({ + userId: 'u1', + workspaceId: 'w1', + docId: 'private', + actions: ['Doc.TransferOwner'], + }); + + t.is(permissions.effectiveRole, 'owner'); + t.is(permissions.resourceOwnerRole, null); + t.is(permissions.legacyApiRole, DocRole.Owner); + t.true(permissions.decisions[0].allowed); +}); + +test('doc legacy boundary keeps resource owner and effective role separate', t => { + t.deepEqual( + docLegacyBoundary({ + docId: 'doc', + resourceOwnerRole: 'owner', + effectiveRole: 'manager', + decisions: [], + }), + { + resourceOwnerRole: 'owner', + effectiveRole: 'manager', + legacyApiRole: DocRole.Manager, + } + ); + t.deepEqual( + docLegacyBoundary({ + docId: 'public', + effectiveRole: 'external', + decisions: [], + }), + { + resourceOwnerRole: null, + effectiveRole: 'external', + legacyApiRole: DocRole.External, + } + ); +}); + +test('PermissionService uses new projection loader behind read flag', async t => { + const calls: string[] = []; + const service = new PermissionService( + { + load: async () => { + calls.push('old'); + return { version: 1 } as never; + }, + loadFromNewTables: async () => { + calls.push('new'); + return { version: 1 } as never; + }, + } as never, + undefined, + undefined, + { + permission: { + readModel: PermissionReadModel.Projection, + fallbackLegacyLoader: true, + }, + } as never + ); + service.evaluate = () => + ({ + version: 1, + workspace: { decisions: [] }, + docs: [ + { + docId: 'doc', + effectiveRole: 'reader', + decisions: [], + }, + ], + }) as never; + + await service.docPermissions({ + workspaceId: 'w1', + docId: 'doc', + actions: ['Doc.Read'], + }); + + t.deepEqual(calls, ['new']); +}); + +test('PermissionService falls back to old loader when new read fails', async t => { + const calls: string[] = []; + const service = new PermissionService( + { + load: async () => { + calls.push('old'); + return { version: 1 } as never; + }, + loadFromNewTables: async () => { + calls.push('new'); + throw new Error('projection unavailable'); + }, + } as never, + undefined, + undefined, + { + permission: { + readModel: PermissionReadModel.Projection, + fallbackLegacyLoader: true, + }, + } as never + ); + service.evaluate = () => + ({ + version: 1, + workspace: { decisions: [] }, + docs: [ + { + docId: 'doc', + effectiveRole: 'reader', + decisions: [], + }, + ], + }) as never; + + await service.docPermissions({ + workspaceId: 'w1', + docId: 'doc', + actions: ['Doc.Read'], + }); + + t.deepEqual(calls, ['new', 'old']); +}); + +test('PermissionService can disable old evaluator fallback', async t => { + const service = new PermissionService( + { + load: async () => ({ version: 1 }) as never, + loadFromNewTables: async () => { + throw new Error('projection unavailable'); + }, + } as never, + undefined, + undefined, + { + permission: { + readModel: PermissionReadModel.Projection, + fallbackLegacyLoader: false, + }, + } as never + ); + + await t.throwsAsync( + service.docPermissions({ + workspaceId: 'w1', + docId: 'doc', + actions: ['Doc.Read'], + }), + { message: 'projection unavailable' } + ); +}); + +test('PermissionService supports anonymous preview without doc read', async t => { + const service = new PermissionService( + createLoader().loader, + undefined, + undefined, + { + permission: { + readModel: PermissionReadModel.Legacy, + fallbackLegacyLoader: false, + }, + } as never + ); + + t.true( + await service.canPreviewDoc({ + workspaceId: 'w1', + docId: 'private', + }) + ); + t.false( + await service.canDoc({ + workspaceId: 'w1', + docId: 'private', + action: 'Doc.Read', + }) + ); +}); + +test('PermissionContextLoader passes missing runtime state as unknown and stale', async t => { + const { loader } = createLoader({ + workspaceId: 'w1', + known: false, + stale: true, + readonly: false, + readonlyReasons: [], + updatedAt: null, + lastReconciledAt: null, + staleAfter: null, + }); + + const input = await loader.load({ + userId: 'u1', + workspaceId: 'w1', + workspaceActions: ['Workspace.CreateDoc'], + }); + + t.is(input.runtime?.known, false); + t.is(input.runtime?.stale, true); +}); + +test('PermissionContextLoader falls back to quota runtime for missing legacy runtime state', async t => { + const { loader } = createLoader( + { + workspaceId: 'w1', + known: false, + stale: true, + readonly: false, + readonlyReasons: [], + updatedAt: null, + lastReconciledAt: null, + staleAfter: null, + }, + null, + async () => [ + { + known: true, + stale: false, + readonly: true, + readonlyReasons: ['storage_overflow'], + staleAfter: null, + }, + ] + ); + + const input = await loader.load({ + userId: 'u1', + workspaceId: 'w1', + workspaceActions: ['Workspace.CreateDoc'], + }); + + t.true(input.runtime?.known); + t.false(input.runtime?.stale); + t.true(input.runtime?.readonly); + t.is(input.runtime?.readonlyReason, 'storage_overflow'); +}); + +test('PermissionContextLoader drops dirty external and none explicit rows', async t => { + for (const role of [DocRole.None, DocRole.External]) { + const { loader } = createLoader(undefined, role); + const input = await loader.load({ + userId: 'u1', + workspaceId: 'w1', + docs: [{ docId: 'private', actions: ['Doc.Update'] }], + }); + + t.is(input.docs?.[0]?.explicitUserRole, undefined); + } +}); + +test('PermissionContextLoader loads shadow input from new projection tables', async t => { + const queryResults = [ + [{ role: 'admin', state: 'active' }], + [ + { + visibility: 'public', + sharingEnabled: true, + urlPreviewEnabled: true, + memberDefaultDocRole: 'reader', + }, + ], + [ + { + known: true, + stale: false, + readonly: true, + readonlyReasons: ['storage_overflow'], + staleAfter: null, + }, + ], + [ + { + docId: 'doc1', + visibility: 'public', + publicRole: 'external', + memberDefaultRole: null, + urlPreviewEnabled: false, + }, + ], + [{ docId: 'doc1', role: 'manager' }], + ]; + const { loader } = createLoader(); + const shadowLoader = new PermissionContextLoader( + (loader as never as { models: never }).models, + { + $queryRaw: async () => queryResults.shift() ?? [], + } as never, + createCls() as never + ); + + const input = await shadowLoader.loadFromNewTables({ + userId: 'u1', + workspaceId: 'w1', + workspaceActions: ['Workspace.Read'], + docs: [{ docId: 'doc1', actions: ['Doc.Read'] }], + }); + + t.is(input.workspace?.role, 'admin'); + t.true(input.runtime?.readonly); + t.is(input.runtime?.readonlyReason, 'storage_overflow'); + t.is(input.workspace?.public, true); + t.is(input.docs?.[0]?.explicitUserRole, 'manager'); + t.is(input.docs?.[0]?.memberDefaultRole, 'reader'); + t.is(input.docs?.[0]?.publicRole, 'external'); +}); + +test('PermissionContextLoader treats missing effective quota state as unknown and stale', async t => { + const queryResults = [[], [], [], [], []]; + const { loader } = createLoader(); + const shadowLoader = new PermissionContextLoader( + (loader as never as { models: never }).models, + { + $queryRaw: async () => queryResults.shift() ?? [], + workspace: { + findUnique: async () => ({ id: 'w1' }), + }, + } as never, + createCls() as never + ); + + const input = await shadowLoader.loadFromNewTables({ + userId: 'u1', + workspaceId: 'w1', + workspaceActions: ['Workspace.CreateDoc'], + }); + + t.false(input.runtime?.known); + t.true(input.runtime?.stale); +}); + +test('PermissionService refreshes cached runtime after write-action reconcile', async t => { + let runtimeQueries = 0; + const loader = new PermissionContextLoader( + { + workspaceUser: { + get: async () => null, + }, + workspace: { + get: async () => ({ + public: false, + enableSharing: true, + enableUrlPreview: false, + }), + }, + workspaceRuntimeState: { + get: async () => ({ + workspaceId: 'w1', + known: false, + stale: true, + readonly: false, + readonlyReasons: [], + }), + }, + doc: { + findDefaultRoles: async () => [], + }, + docUser: { + findMany: async () => [], + }, + } as never, + { + $queryRaw: async (strings: TemplateStringsArray) => { + const sql = strings.join(''); + if (sql.includes('workspace_members')) { + return [{ role: 'owner', state: 'active' }]; + } + if (sql.includes('workspace_access_policies')) { + return [ + { + visibility: 'private', + sharingEnabled: true, + urlPreviewEnabled: false, + memberDefaultDocRole: 'manager', + }, + ]; + } + if (sql.includes('effective_workspace_quota_states')) { + runtimeQueries += 1; + return runtimeQueries === 1 + ? [] + : [ + { + known: true, + stale: false, + readonly: false, + readonlyReasons: [], + staleAfter: null, + }, + ]; + } + return []; + }, + workspace: { + findUnique: async () => ({ id: 'w1' }), + }, + } as never, + createCls() as never + ); + const service = new PermissionService(loader, undefined, { + getWorkspaceState: async () => ({}), + } as never); + const runtimeKnownValues: Array = []; + service.evaluate = input => { + runtimeKnownValues.push(input.runtime?.known); + return { + version: 1, + workspace: { + decisions: [ + { + action: input.workspaceActions?.[0] ?? 'Workspace.Read', + allowed: input.runtime?.known !== false, + sources: [], + restrictions: [], + }, + ], + }, + docs: [], + }; + }; + + t.false( + await service.canWorkspace({ + userId: 'u1', + workspaceId: 'w1', + action: 'Workspace.Read', + }) + ); + t.true( + await service.canWorkspace({ + userId: 'u1', + workspaceId: 'w1', + action: 'Workspace.CreateDoc', + }) + ); + t.deepEqual(runtimeKnownValues, [false, true]); + t.is(runtimeQueries, 2); +}); + +test('PermissionService does not reconcile runtime for read-only workspace actions', async t => { + let runtimeQueries = 0; + let reconciles = 0; + const loader = new PermissionContextLoader( + { + workspaceUser: { + get: async () => null, + }, + workspace: { + get: async () => ({ + public: false, + enableSharing: true, + enableUrlPreview: false, + }), + }, + workspaceRuntimeState: { + get: async () => null, + }, + doc: { + findDefaultRoles: async () => [], + }, + docUser: { + findMany: async () => [], + }, + } as never, + { + $queryRaw: async (strings: TemplateStringsArray) => { + const sql = strings.join(''); + if (sql.includes('workspace_members')) { + return [{ role: 'owner', state: 'active' }]; + } + if (sql.includes('workspace_access_policies')) { + return [ + { + visibility: 'private', + sharingEnabled: true, + urlPreviewEnabled: false, + memberDefaultDocRole: 'manager', + }, + ]; + } + if (sql.includes('effective_workspace_quota_states')) { + runtimeQueries += 1; + return []; + } + return []; + }, + workspace: { + findUnique: async () => ({ id: 'w1' }), + }, + } as never, + createCls() as never + ); + const service = new PermissionService(loader, undefined, { + getWorkspaceState: async () => { + reconciles += 1; + return {}; + }, + } as never); + service.evaluate = input => ({ + version: 1, + workspace: { + decisions: [ + { + action: input.workspaceActions?.[0] ?? 'Workspace.Read', + allowed: true, + sources: [], + restrictions: [], + }, + ], + }, + docs: [], + }); + + t.true( + await service.canWorkspace({ + userId: 'u1', + workspaceId: 'w1', + action: 'Workspace.Settings.Read', + }) + ); + t.is(runtimeQueries, 1); + t.is(reconciles, 0); +}); + +test('PermissionService reconciles runtime for non-readonly write actions', async t => { + let runtimeQueries = 0; + let reconciles = 0; + const loader = new PermissionContextLoader( + { + workspaceUser: { + get: async () => null, + }, + workspace: { + get: async () => ({ + public: false, + enableSharing: true, + enableUrlPreview: false, + }), + }, + workspaceRuntimeState: { + get: async () => null, + }, + doc: { + findDefaultRoles: async () => [], + }, + docUser: { + findMany: async () => [], + }, + } as never, + { + $queryRaw: async (strings: TemplateStringsArray) => { + const sql = strings.join(''); + if (sql.includes('workspace_members')) { + return [{ role: 'owner', state: 'active' }]; + } + if (sql.includes('workspace_access_policies')) { + return [ + { + visibility: 'private', + sharingEnabled: true, + urlPreviewEnabled: false, + memberDefaultDocRole: 'manager', + }, + ]; + } + if (sql.includes('effective_workspace_quota_states')) { + runtimeQueries += 1; + return [ + { + known: true, + stale: false, + readonly: false, + readonlyReasons: [], + staleAfter: null, + }, + ]; + } + return []; + }, + workspace: { + findUnique: async () => ({ id: 'w1' }), + }, + } as never, + createCls() as never + ); + const service = new PermissionService(loader, undefined, { + getWorkspaceState: async () => { + reconciles += 1; + return {}; + }, + } as never); + service.evaluate = input => ({ + version: 1, + workspace: { + decisions: [ + { + action: input.workspaceActions?.[0] ?? 'Workspace.Read', + allowed: input.runtime?.known === true, + sources: [], + restrictions: [], + }, + ], + }, + docs: [], + }); + + t.true( + await service.canWorkspace({ + userId: 'u1', + workspaceId: 'w1', + action: 'Workspace.Delete', + }) + ); + t.is(runtimeQueries, 1); + t.is(reconciles, 1); +}); + +for (const projectionLocalCase of [ + { + name: 'does not treat missing projection as local workspace', + workspaceId: 'w1', + workspaceRow: { id: 'w1' }, + expectedLocal: false, + }, + { + name: 'allows local only when workspace is absent', + workspaceId: 'local', + workspaceRow: null, + expectedLocal: true, + }, +] as const) { + test(`PermissionContextLoader ${projectionLocalCase.name}`, async t => { + const queryResults = [[], [], [], [], []]; + const { loader } = createLoader(); + const shadowLoader = new PermissionContextLoader( + (loader as never as { models: never }).models, + { + $queryRaw: async () => queryResults.shift() ?? [], + workspace: { + findUnique: async () => projectionLocalCase.workspaceRow, + }, + } as never, + createCls() as never + ); + + const input = await shadowLoader.loadFromNewTables({ + userId: 'u1', + workspaceId: projectionLocalCase.workspaceId, + allowLocal: true, + }); + + t.is(input.workspace?.local, projectionLocalCase.expectedLocal); + }); +} + +for (const shadowDocCase of [ + { + name: 'reports projection mismatches', + expectedMismatchType: 'rust_rule', + buildEvaluate: () => { + let index = 0; + return () => + ({ + version: 1, + workspace: { + decisions: [], + }, + docs: [ + { + docId: 'doc1', + decisions: [ + { + action: 'Doc.Read', + allowed: index++ === 0, + sources: [], + restrictions: [], + }, + ], + }, + ], + }) as never; + }, + }, + { + name: 'classifies legacy API role mapping mismatches', + expectedMismatchType: 'legacy_api_role_mapping', + buildEvaluate: () => { + let index = 0; + return () => + ({ + version: 1, + workspace: { + decisions: [], + }, + docs: [ + { + docId: 'doc1', + effectiveRole: index++ === 0 ? 'owner' : 'manager', + resourceOwnerRole: null, + decisions: [], + }, + ], + }) as never; + }, + }, + { + name: 'accepts explicit expected delta categories', + expectedDeltaCategory: 'legacy_compat_delta', + expectedMismatchType: 'legacy_compat_delta', + buildEvaluate: () => { + let index = 0; + return () => + ({ + version: 1, + workspace: { + decisions: [], + }, + docs: [ + { + docId: 'doc1', + decisions: [ + { + action: 'Doc.Read', + allowed: index++ === 0, + sources: [], + restrictions: [], + }, + ], + }, + ], + }) as never; + }, + }, +] as const) { + test(`PermissionDiagnosticService ${shadowDocCase.name}`, async t => { + const loader = { + load: async () => ({ version: 1 }) as never, + loadFromNewTables: async () => ({ version: 1 }) as never, + } as never; + const permission = new PermissionService(loader); + permission.evaluate = shadowDocCase.buildEvaluate(); + const service = new PermissionDiagnosticService(loader, permission); + + const result = await service.shadowDocPermissions({ + workspaceId: 'w1', + docs: [{ docId: 'doc1', actions: ['Doc.Read'] }], + expectedDeltaCategory: shadowDocCase.expectedDeltaCategory, + }); + + t.false(result.matched); + t.is(result.mismatchType, shadowDocCase.expectedMismatchType); + }); +} + +test('PermissionDiagnosticService compares SQL predicate shadow doc read results', async t => { + const loader = { + loadFromNewTables: async () => ({ version: 1 }) as never, + } as never; + const permission = new PermissionService(loader); + permission.evaluate = () => + ({ + version: 1, + workspace: { decisions: [] }, + docs: ['ok', 'missing', 'extra'].map(docId => ({ + docId, + decisions: [ + { + action: 'Doc.Read', + allowed: docId !== 'extra', + sources: [], + restrictions: [], + }, + ], + })), + }) as never; + const service = new PermissionDiagnosticService(loader, permission); + + const result = await service.shadowSqlDocRead({ + userId: 'u1', + workspaceId: 'w1', + docs: [{ docId: 'ok' }, { docId: 'missing' }, { docId: 'extra' }], + sqlReadableDocIds: ['ok', 'extra'], + }); + + t.false(result.matched); + t.is(result.mismatchType, 'sql_predicate'); + t.deepEqual(result.missingInSql, ['missing']); + t.deepEqual(result.extraInSql, ['extra']); + t.true(result.predicate.sql.includes('doc_access_policies')); +}); + +test('PermissionDiagnosticService preview shadow flags preview/read coupling', async t => { + const loader = { + load: async () => ({ version: 1 }) as never, + loadFromNewTables: async () => ({ version: 1 }) as never, + } as never; + const permission = new PermissionService(loader); + const service = new PermissionDiagnosticService(loader, permission); + let index = 0; + permission.evaluate = () => + ({ + version: 1, + workspace: { + decisions: [], + }, + docs: [ + { + docId: 'doc1', + decisions: [ + { + action: 'Doc.Preview', + allowed: true, + sources: [], + restrictions: [], + }, + { + action: 'Doc.Read', + allowed: index++ === 1, + sources: [], + restrictions: [], + }, + ], + }, + ], + }) as never; + + const result = await service.shadowPreviewDoc({ + workspaceId: 'w1', + docId: 'doc1', + }); + + t.false(result.matched); + t.is(result.mismatchType, 'preview_read_mapping'); +}); + +test('PermissionDiagnosticService keeps public preview/read shadow as normal match', async t => { + const loader = { + load: async () => ({ version: 1 }) as never, + loadFromNewTables: async () => ({ version: 1 }) as never, + } as never; + const permission = new PermissionService(loader); + const service = new PermissionDiagnosticService(loader, permission); + permission.evaluate = () => + ({ + version: 1, + workspace: { + decisions: [], + }, + docs: [ + { + docId: 'doc1', + decisions: [ + { + action: 'Doc.Preview', + allowed: true, + sources: [], + restrictions: [], + }, + { + action: 'Doc.Read', + allowed: true, + sources: [], + restrictions: [], + }, + ], + }, + ], + }) as never; + + const result = await service.shadowPreviewDoc({ + workspaceId: 'w1', + docId: 'doc1', + }); + + t.true(result.matched); + t.is(result.mismatchType, null); +}); + +test('PermissionDiagnosticService exposes shadow mismatch categories', t => { + t.deepEqual(PERMISSION_SHADOW_MISMATCH_CATEGORIES, [ + 'legacy_compat_delta', + 'projection', + 'rust_rule', + 'loader', + 'sql_predicate', + 'legacy_api_role_mapping', + 'preview_read_mapping', + 'runtime_state', + 'projection_or_loader', + ]); +}); + +test('PermissionService maps native validation errors to internal errors', t => { + const service = new PermissionService(createLoader().loader); + + const error = t.throws(() => + service.evaluate({ + version: 2, + } as never) + ); + + t.true(error instanceof InternalServerError); +}); + +test('PermissionContextLoader memoizes permission context within a request', async t => { + const { loader, calls } = createLoader(); + + await loader.load({ + userId: 'u1', + workspaceId: 'w1', + workspaceActions: ['Workspace.Read'], + docs: [{ docId: 'public', actions: ['Doc.Read'] }], + }); + await loader.load({ + userId: 'u1', + workspaceId: 'w1', + workspaceActions: ['Workspace.Read'], + docs: [{ docId: 'public', actions: ['Doc.Read'] }], + }); + + t.deepEqual(calls, { + workspaceUser: 1, + workspace: 1, + runtime: 1, + docPolicies: 1, + docGrants: 1, + }); +}); + +test('PermissionSqlPredicateBuilder builds legacy-table predicate parameters', t => { + const predicate = + new PermissionSqlPredicateBuilder().docReadableByLegacyTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Users.Manage', + docIdColumn: 'docs.id', + }); + + t.true(predicate.sql.includes('workspace_page_user_permissions')); + t.deepEqual(predicate.params.slice(0, 3), ['u1', 'u1', 'w1']); + t.true( + predicate.params + .slice(3) + .filter(Array.isArray) + .every(value => value.every(Number.isInteger)) + ); +}); + +test('PermissionSqlPredicateBuilder rejects unsafe raw doc id columns', t => { + const builder = new PermissionSqlPredicateBuilder(); + + t.throws( + () => + builder.docReadableByLegacyTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Read', + docIdColumn: 'docs.id; DROP TABLE docs' as never, + }), + { message: 'Unsupported doc id column: docs.id; DROP TABLE docs' } + ); + t.throws( + () => + builder.docReadableByNewTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Read', + docIdColumn: 'docs.id; DROP TABLE docs' as never, + }), + { message: 'Unsupported doc id column: docs.id; DROP TABLE docs' } + ); +}); + +test('PermissionSqlPredicateBuilder drops dirty legacy external explicit grants', t => { + const predicate = + new PermissionSqlPredicateBuilder().docReadableByLegacyTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Read', + }); + + const roles = predicate.params[4] as DocRole[]; + t.false(roles.includes(DocRole.External)); + t.true(roles.includes(DocRole.Reader)); +}); + +test('PermissionSqlPredicateBuilder treats legacy workspace external rows as non-members', t => { + const predicate = + new PermissionSqlPredicateBuilder().docReadableByLegacyTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Users.Manage', + }); + + const activeMemberRoles = predicate.params[3] as WorkspaceRole[]; + t.false(activeMemberRoles.includes(WorkspaceRole.External)); +}); + +test('PermissionSqlPredicateBuilder caps non-member explicit grants below manager', t => { + const builder = new PermissionSqlPredicateBuilder(); + const update = builder.docReadableByNewTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Update', + }); + const transferOwner = builder.docReadableByNewTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.TransferOwner', + }); + + t.true((update.params[4] as string[]).includes('editor')); + t.true((update.params[4] as string[]).includes('manager')); + t.true((update.params[4] as string[]).includes('owner')); + t.deepEqual(transferOwner.params[3], ['owner']); + t.deepEqual(transferOwner.params[4], []); +}); + +test('PermissionSqlPredicateBuilder suppresses member default when explicit grant exists', t => { + const builder = new PermissionSqlPredicateBuilder(); + const legacy = builder.docReadableByLegacyTablesSql({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Update', + }); + const projection = builder.docReadableByNewTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Update', + }); + + t.true(legacy.sql.includes('p.user_id IS NULL')); + t.true(projection.sql.includes('dg.principal_id IS NULL')); +}); + +test('PermissionSqlPredicateBuilder accepts representative doc actions', t => { + const builder = new PermissionSqlPredicateBuilder(); + const actions = [ + 'Doc.Read', + 'Doc.Update', + 'Doc.Duplicate', + 'Doc.Users.Manage', + 'Doc.TransferOwner', + ] as const; + + for (const action of actions) { + const legacy = builder.docReadableByLegacyTables({ + workspaceId: 'w1', + userId: 'u1', + action: action as never, + }); + const projection = builder.docReadableByNewTables({ + workspaceId: 'w1', + userId: 'u1', + action: action as never, + }); + + t.true(legacy.sql.includes('workspace_page_user_permissions')); + t.true(projection.sql.includes('workspace_access_policies')); + } +}); + +test('PermissionService chooses SQL predicate from configured read model and fallback flag', t => { + const { loader } = createLoader(); + const legacy = new PermissionService(loader, undefined, undefined, { + permission: { + readModel: PermissionReadModel.Legacy, + fallbackLegacyLoader: true, + }, + } as never).docReadableSqlPredicate({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Read', + }); + const projection = new PermissionService(loader, undefined, undefined, { + permission: { + readModel: PermissionReadModel.Projection, + fallbackLegacyLoader: false, + }, + } as never).docReadableSqlPredicate({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Read', + }); + const projectionWithFallback = new PermissionService( + loader, + undefined, + undefined, + { + permission: { + readModel: PermissionReadModel.Projection, + fallbackLegacyLoader: true, + }, + } as never + ).docReadableSqlPredicate({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Read', + }); + const fallback = new PermissionService(loader, undefined, undefined, { + permission: { + readModel: PermissionReadModel.Projection, + fallbackLegacyLoader: true, + }, + } as never).fallbackDocReadableSqlPredicate({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Read', + }); + + t.true( + (legacy as unknown as { sql: string }).sql.includes( + 'workspace_user_permissions' + ) + ); + t.false( + (legacy as unknown as { sql: string }).sql.includes( + 'workspace_access_policies' + ) + ); + t.true( + (projection as unknown as { sql: string }).sql.includes( + 'workspace_access_policies' + ) + ); + t.true( + (projectionWithFallback as unknown as { sql: string }).sql.includes( + 'workspace_access_policies' + ) + ); + t.false( + (projectionWithFallback as unknown as { sql: string }).sql.includes( + 'workspace_user_permissions' + ) + ); + t.true( + (fallback as unknown as { sql: string }).sql.includes( + 'workspace_user_permissions' + ) + ); +}); + +test('PermissionSqlPredicateBuilder uses new projection tables for shadow read', t => { + const predicate = new PermissionSqlPredicateBuilder().docReadableByNewTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Read', + docIdColumn: 'docs.id', + }); + + t.true(predicate.sql.includes('FROM workspace_access_policies wap')); + t.true(predicate.sql.includes('LEFT JOIN doc_access_policies dap')); + t.true(predicate.sql.includes('workspace_members')); + t.true(predicate.sql.includes('doc_grants')); + t.true(predicate.sql.includes('dap.doc_id = docs.id')); + t.true( + predicate.sql.includes( + 'COALESCE(dap.member_default_role, wap.member_default_doc_role)' + ) + ); + t.deepEqual(predicate.params.slice(0, 3), ['u1', 'u1', 'w1']); +}); + +test('PermissionSqlPredicateBuilder emits new-table role parameters as strings', t => { + const builder = new PermissionSqlPredicateBuilder(); + const transferOwner = builder.docReadableByNewTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.TransferOwner', + }); + const manageUsers = builder.docReadableByNewTables({ + workspaceId: 'w1', + userId: 'u1', + action: 'Doc.Users.Manage', + }); + + for (const predicate of [transferOwner, manageUsers]) { + t.true( + predicate.params + .slice(3) + .every( + value => + Array.isArray(value) && + value.every(role => typeof role === 'string') + ) + ); + } +}); diff --git a/packages/backend/server/src/core/permission/__tests__/workspace.spec.ts b/packages/backend/server/src/core/permission/__tests__/workspace.spec.ts index d6d65bd949..d36b90bc6e 100644 --- a/packages/backend/server/src/core/permission/__tests__/workspace.spec.ts +++ b/packages/backend/server/src/core/permission/__tests__/workspace.spec.ts @@ -10,14 +10,13 @@ import { WorkspaceMemberStatus, WorkspaceRole, } from '../../../models'; -import { PermissionModule } from '../index'; +import { PermissionAccess, PermissionModule } from '../index'; import { WorkspacePolicyService } from '../policy'; import { mapWorkspaceRoleToPermissions } from '../types'; -import { WorkspaceAccessController } from '../workspace'; let module: TestingModule; let models: Models; -let ac: WorkspaceAccessController; +let ac: PermissionAccess; let policy: WorkspacePolicyService; let user: User; let ws: Workspace; @@ -26,7 +25,7 @@ let underReviewUserId: string; test.before(async () => { module = await createTestingModule({ imports: [PermissionModule] }); models = module.get(Models); - ac = module.get(WorkspaceAccessController); + ac = module.get(PermissionAccess); policy = module.get(WorkspacePolicyService); }); @@ -138,10 +137,34 @@ const roleCases: Array<{ }, ]; +async function getRole(resource: { + workspaceId: string; + userId: string; + allowLocal?: boolean; +}) { + const checker = ac.user(resource.userId).workspace(resource.workspaceId); + if (resource.allowLocal) { + checker.allowLocal(); + } + return (await checker.permissions()).role; +} + +function workspace(resource: { + workspaceId: string; + userId: string; + allowLocal?: boolean; +}) { + const checker = ac.user(resource.userId).workspace(resource.workspaceId); + if (resource.allowLocal) { + checker.allowLocal(); + } + return checker; +} + for (const roleCase of roleCases) { test(roleCase.title, async t => { await roleCase.setup?.(); - const role = await ac.getRole(roleCase.resource()); + const role = await getRole(roleCase.resource()); t.is(role, roleCase.expectedRole); }); @@ -150,10 +173,10 @@ for (const roleCase of roleCases) { test('should return mapped null permission even workspace has public docs', async t => { await models.doc.publish(ws.id, 'doc1'); - const { permissions } = await ac.role({ + const { permissions } = await workspace({ workspaceId: ws.id, userId: 'random-user-id', - }); + }).permissions(); t.deepEqual(permissions, mapWorkspaceRoleToPermissions(null)); }); @@ -162,13 +185,10 @@ test('should deny external read assert even workspace has public docs', async t await models.doc.publish(ws.id, 'doc1'); await t.throwsAsync( - ac.assert( - { - workspaceId: ws.id, - userId: 'random-user-id', - }, - 'Workspace.Read' - ) + workspace({ + workspaceId: ws.id, + userId: 'random-user-id', + }).assert('Workspace.Read') ); }); @@ -177,13 +197,10 @@ test('should deny external read assert when sharing disabled even if workspace h await models.workspace.update(ws.id, { enableSharing: false }); await t.throwsAsync( - ac.assert( - { - workspaceId: ws.id, - userId: 'random-user-id', - }, - 'Workspace.Read' - ) + workspace({ + workspaceId: ws.id, + userId: 'random-user-id', + }).assert('Workspace.Read') ); }); @@ -193,31 +210,27 @@ test('should reject external doc roles when sharing disabled', async t => { enableSharing: false, }); - const [docRole] = await ac.docRoles( - { - workspaceId: ws.id, - userId: 'random-user-id', - }, - ['doc1'] - ); + const docRole = await ac + .user('random-user-id') + .doc(ws.id, 'doc1') + .permissions(); t.is(docRole.role, null); t.false(docRole.permissions['Doc.Read']); }); test('should return mapped permissions', async t => { - const { permissions } = await ac.role({ + const { permissions } = await workspace({ workspaceId: ws.id, userId: user.id, - }); + }).permissions(); t.deepEqual(permissions, mapWorkspaceRoleToPermissions(WorkspaceRole.Owner)); }); test('should assert action', async t => { await t.notThrowsAsync( - ac.assert( - { workspaceId: ws.id, userId: user.id }, + workspace({ workspaceId: ws.id, userId: user.id }).assert( 'Workspace.TransferOwner' ) ); @@ -225,7 +238,7 @@ test('should assert action', async t => { const u2 = await models.user.create({ email: 'u2@affine.pro' }); await t.throwsAsync( - ac.assert({ workspaceId: ws.id, userId: u2.id }, 'Workspace.Sync') + workspace({ workspaceId: ws.id, userId: u2.id }).assert('Workspace.Sync') ); await models.workspaceUser.set(ws.id, u2.id, WorkspaceRole.Admin, { @@ -233,8 +246,7 @@ test('should assert action', async t => { }); await t.notThrowsAsync( - ac.assert( - { workspaceId: ws.id, userId: u2.id }, + workspace({ workspaceId: ws.id, userId: u2.id }).assert( 'Workspace.Settings.Update' ) ); @@ -256,10 +268,10 @@ test('should apply readonly workspace restrictions while keeping cleanup actions } await policy.reconcileWorkspaceQuotaState(ws.id); - const { permissions } = await ac.role({ + const { permissions } = await workspace({ workspaceId: ws.id, userId: user.id, - }); + }).permissions(); t.false(permissions['Workspace.CreateDoc']); t.false(permissions['Workspace.Settings.Update']); diff --git a/packages/backend/server/src/core/permission/builder.ts b/packages/backend/server/src/core/permission/builder.ts index b7ff4a9cd2..0116771068 100644 --- a/packages/backend/server/src/core/permission/builder.ts +++ b/packages/backend/server/src/core/permission/builder.ts @@ -1,26 +1,47 @@ import { Injectable } from '@nestjs/common'; import { DocID } from '../utils/doc'; -import { getAccessController } from './controller'; import { Resource } from './resource'; -import { DocAction, WorkspaceAction } from './types'; -import { WorkspaceAccessController } from './workspace'; +import { PermissionService } from './service'; +import { + DOC_ACTIONS, + DocAction, + DocRole, + WORKSPACE_ACTIONS, + WorkspaceAction, + WorkspaceRole, +} from './types'; + +function assertPerm(permission?: PermissionService) { + if (!permission) { + throw new Error('PermissionService is required for permission checks.'); + } + return permission; +} @Injectable() export class AccessControllerBuilder { + constructor(private readonly permission?: PermissionService) {} + user(userId: string) { - return new UserAccessControllerBuilder(userId); + return new UserAccessControllerBuilder(userId, this.permission); } } export class UserAccessControllerBuilder { - constructor(private readonly userId: string) {} + constructor( + private readonly userId: string, + private readonly permission?: PermissionService + ) {} workspace(workspaceId: string) { - return new WorkspaceAccessControllerBuilder({ - userId: this.userId, - workspaceId, - }); + return new WorkspaceAccessControllerBuilder( + { + userId: this.userId, + workspaceId, + }, + this.permission + ); } doc( @@ -45,16 +66,22 @@ export class UserAccessControllerBuilder { docId = docIdOrWorkspaceId.docId; } - return new DocAccessControllerBuilder({ - userId: this.userId, - workspaceId, - docId, - }); + return new DocAccessControllerBuilder( + { + userId: this.userId, + workspaceId, + docId, + }, + this.permission + ); } } class WorkspaceAccessControllerBuilder { - constructor(public readonly data: Resource<'ws'>) {} + constructor( + public readonly data: Resource<'ws'>, + private readonly permission?: PermissionService + ) {} allowLocal() { this.data.allowLocal = true; @@ -62,10 +89,13 @@ class WorkspaceAccessControllerBuilder { } doc(docId: string) { - return new DocAccessControllerBuilder({ - ...this.data, - docId, - }); + return new DocAccessControllerBuilder( + { + ...this.data, + docId, + }, + this.permission + ); } /** @@ -79,35 +109,61 @@ class WorkspaceAccessControllerBuilder { action: DocAction ): Promise { const docIds = items.map(item => item.docId); - const checker = getAccessController('ws') as WorkspaceAccessController; - const docRoles = await checker.docRoles(this.data, docIds); + const docRoles = await assertPerm(this.permission).batchDocPermissions({ + userId: this.data.userId, + workspaceId: this.data.workspaceId, + docs: docIds.map(docId => ({ + docId, + actions: [action], + })), + allowLocal: this.data.allowLocal, + }); const docRolesMap = new Map( docRoles.map((role, index) => [docIds[index], role]) ); return items.filter(item => { - return docRolesMap.get(item.docId)?.permissions[action]; + return docRolesMap + .get(item.docId) + ?.decisions.some( + decision => decision.action === action && decision.allowed + ); }); } async assert(action: WorkspaceAction) { - const checker = getAccessController('ws'); - await checker.assert(this.data, action); + await assertPerm(this.permission).assertWorkspace({ + ...this.data, + action, + }); } async can(action: WorkspaceAction) { - const checker = getAccessController('ws'); - return await checker.can(this.data, action); + return await assertPerm(this.permission).canWorkspace({ + ...this.data, + action, + }); } async permissions() { - const checker = getAccessController('ws'); - return await checker.role(this.data); + const result = await assertPerm(this.permission).workspacePermissions({ + ...this.data, + actions: [...WORKSPACE_ACTIONS], + }); + return { + role: result.legacyApiRole as WorkspaceRole | null, + permissions: Object.fromEntries( + result.decisions.map(decision => [decision.action, decision.allowed]) + ) as Record, + }; } } class DocAccessControllerBuilder { - constructor(public readonly data: Resource<'doc'>) {} + constructor( + public readonly data: Resource<'doc'>, + private readonly permission?: PermissionService + ) {} allowLocal() { this.data.allowLocal = true; @@ -115,17 +171,29 @@ class DocAccessControllerBuilder { } async assert(action: DocAction) { - const checker = getAccessController('doc'); - await checker.assert(this.data, action); + await assertPerm(this.permission).assertDoc({ + ...this.data, + action, + }); } async can(action: DocAction) { - const checker = getAccessController('doc'); - return await checker.can(this.data, action); + return await assertPerm(this.permission).canDoc({ + ...this.data, + action, + }); } async permissions() { - const checker = getAccessController('doc'); - return await checker.role(this.data); + const result = await assertPerm(this.permission).docPermissions({ + ...this.data, + actions: [...DOC_ACTIONS], + }); + return { + role: result.legacyApiRole as DocRole | null, + permissions: Object.fromEntries( + result.decisions.map(decision => [decision.action, decision.allowed]) + ) as Record, + }; } } diff --git a/packages/backend/server/src/core/permission/config.ts b/packages/backend/server/src/core/permission/config.ts new file mode 100644 index 0000000000..e0a605e542 --- /dev/null +++ b/packages/backend/server/src/core/permission/config.ts @@ -0,0 +1,31 @@ +import { z } from 'zod'; + +import { defineModuleConfig } from '../../base'; + +export enum PermissionReadModel { + Legacy = 'legacy', + Projection = 'projection', +} + +declare global { + interface AppConfigSchema { + permission: { + readModel: PermissionReadModel; + fallbackLegacyLoader: boolean; + }; + } +} + +defineModuleConfig('permission', { + readModel: { + desc: 'Permission data source for Rust evaluation', + default: PermissionReadModel.Projection, + shape: z.nativeEnum(PermissionReadModel), + env: ['AFFINE_PERMISSION_READ_MODEL', 'string'], + }, + fallbackLegacyLoader: { + desc: 'Fallback from projection loader to legacy loader when projection input loading fails', + default: false, + env: ['AFFINE_PERMISSION_FALLBACK_LEGACY_LOADER', 'boolean'], + }, +}); diff --git a/packages/backend/server/src/core/permission/context-loader.ts b/packages/backend/server/src/core/permission/context-loader.ts new file mode 100644 index 0000000000..0ff6a0d023 --- /dev/null +++ b/packages/backend/server/src/core/permission/context-loader.ts @@ -0,0 +1,463 @@ +import { Injectable } from '@nestjs/common'; +import { PrismaClient } from '@prisma/client'; +import { ClsService } from 'nestjs-cls'; + +import { DocRole, Models } from '../../models'; +import type { PermissionEvaluationInputV1 } from '../../native'; +import { + toNativeDocRole, + toNativeExplicitDocGrantRole, + toNativeMemberState, + toNativeWorkspaceRole, +} from './context'; +import type { DocAction, WorkspaceAction } from './types'; + +type PermissionRequestCache = { + workspaceMember: Map< + string, + Awaited> + >; + workspacePolicy: Map>>; + workspaceRuntime: Map< + string, + Awaited> + >; + workspaceQuotaRuntime: Map; + docPolicies: Map< + string, + Awaited> + >; + docGrants: Map>>; +}; + +type NewWorkspaceMemberRow = { + role: 'owner' | 'admin' | 'member'; + state: 'active' | 'suspended' | 'left'; +}; + +type NewWorkspacePolicyRow = { + visibility: 'private' | 'public'; + sharingEnabled: boolean; + urlPreviewEnabled: boolean; + memberDefaultDocRole: 'none' | 'reader' | 'commenter' | 'editor' | 'manager'; +}; + +type NewDocPolicyRow = { + docId: string; + visibility: 'private' | 'public'; + publicRole: 'external' | null; + memberDefaultRole: + | 'none' + | 'reader' + | 'commenter' + | 'editor' + | 'manager' + | null; + urlPreviewEnabled: boolean; +}; + +type NewDocGrantRow = { + docId: string; + role: 'owner' | 'manager' | 'editor' | 'commenter' | 'reader'; +}; + +type NewWorkspaceRuntimeState = { + known: boolean; + stale: boolean; + readonly: boolean; + readonlyReasons: string[]; + staleAfter: Date | null; +}; + +const CACHE_KEY = 'permission.context.cache'; + +function createPermissionRequestCache(): PermissionRequestCache { + return { + workspaceMember: new Map(), + workspacePolicy: new Map(), + workspaceRuntime: new Map(), + workspaceQuotaRuntime: new Map(), + docPolicies: new Map(), + docGrants: new Map(), + }; +} + +export type PermissionWorkspaceAction = WorkspaceAction | 'Workspace.Preview'; +export type PermissionDocAction = DocAction | 'Doc.Preview'; + +function cacheKey(parts: readonly unknown[]) { + return parts.join('\0'); +} + +@Injectable() +export class PermissionContextLoader { + constructor( + private readonly models: Models, + private readonly db: PrismaClient, + private readonly cls?: ClsService + ) {} + + async load(input: { + userId?: string; + workspaceId: string; + allowLocal?: boolean; + workspaceActions?: PermissionWorkspaceAction[]; + docs?: Array<{ docId: string; actions: PermissionDocAction[] }>; + }): Promise { + const docs = input.docs ?? []; + const [member, workspace, runtime, docPolicies, docGrants] = + await Promise.all([ + input.userId + ? this.workspaceMember(input.workspaceId, input.userId) + : Promise.resolve(null), + this.workspacePolicy(input.workspaceId), + this.workspaceRuntime(input.workspaceId), + this.docPolicies( + input.workspaceId, + docs.map(doc => doc.docId) + ), + input.userId + ? this.docGrants( + input.workspaceId, + docs.map(doc => doc.docId), + input.userId + ) + : Promise.resolve([]), + ]); + + const docGrantMap = new Map(docGrants.map(grant => [grant.docId, grant])); + const workspaceSharingEnabled = workspace?.enableSharing ?? true; + + return { + version: 1, + legacyCompatMode: true, + subject: { + userId: input.userId, + groupIds: [], + allowLocal: input.allowLocal, + }, + runtime: { + known: runtime.known, + stale: runtime.stale, + readonly: runtime.readonly, + readonlyReason: runtime.readonlyReasons[0], + sharingEnabled: workspaceSharingEnabled, + urlPreviewEnabled: workspace?.enableUrlPreview ?? false, + }, + workspace: { + role: toNativeWorkspaceRole(member?.type), + memberState: toNativeMemberState(member?.status), + public: workspace?.public ?? false, + sharingEnabled: workspaceSharingEnabled, + urlPreviewEnabled: workspace?.enableUrlPreview ?? false, + local: !workspace, + }, + workspaceActions: input.workspaceActions, + docs: docs.map((doc, index) => { + const policy = docPolicies[index]; + const grant = docGrantMap.get(doc.docId); + return { + docId: doc.docId, + actions: doc.actions, + explicitUserRole: toNativeExplicitDocGrantRole(grant?.type), + groupGrants: [], + groupGrantsEnabled: false, + memberDefaultRole: toNativeDocRole( + policy?.workspace ?? DocRole.Manager + ), + publicRole: policy?.external === null ? undefined : 'external', + visibility: policy?.external === null ? 'private' : 'public', + sharingEnabled: workspaceSharingEnabled, + previewEnabled: policy?.external !== null, + }; + }), + }; + } + + async loadFromNewTables(input: { + userId?: string; + workspaceId: string; + allowLocal?: boolean; + workspaceActions?: PermissionWorkspaceAction[]; + docs?: Array<{ docId: string; actions: PermissionDocAction[] }>; + }): Promise { + const docs = input.docs ?? []; + const docIds = docs.map(doc => doc.docId); + const [member, workspacePolicy, runtime, docPolicies, docGrants] = + await Promise.all([ + input.userId + ? this.newWorkspaceMember(input.workspaceId, input.userId) + : Promise.resolve(null), + this.newWorkspacePolicy(input.workspaceId), + this.newWorkspaceRuntime(input.workspaceId), + this.newDocPolicies(input.workspaceId, docIds), + input.userId + ? this.newDocGrants(input.workspaceId, docIds, input.userId) + : Promise.resolve([]), + ]); + const docPolicyMap = new Map( + docPolicies.map(policy => [policy.docId, policy]) + ); + const docGrantMap = new Map(docGrants.map(grant => [grant.docId, grant])); + const local = + !workspacePolicy && + !!input.allowLocal && + !(await this.workspaceExists(input.workspaceId)); + const sharingEnabled = workspacePolicy?.sharingEnabled ?? true; + const urlPreviewEnabled = workspacePolicy?.urlPreviewEnabled ?? false; + + return { + version: 1, + legacyCompatMode: true, + subject: { + userId: input.userId, + groupIds: [], + allowLocal: input.allowLocal, + }, + runtime: { + known: runtime.known, + stale: runtime.stale, + readonly: runtime.readonly, + readonlyReason: runtime.readonlyReasons[0], + sharingEnabled, + urlPreviewEnabled, + }, + workspace: { + role: member?.role, + memberState: member?.state === 'active' ? 'active' : undefined, + public: workspacePolicy?.visibility === 'public', + sharingEnabled, + urlPreviewEnabled, + local, + }, + workspaceActions: input.workspaceActions, + docs: docs.map(doc => { + const policy = docPolicyMap.get(doc.docId); + const grant = docGrantMap.get(doc.docId); + const visibility = policy?.visibility ?? 'private'; + const publicRole = policy?.publicRole ?? undefined; + return { + docId: doc.docId, + actions: doc.actions, + explicitUserRole: grant?.role, + groupGrants: [], + groupGrantsEnabled: false, + memberDefaultRole: + policy?.memberDefaultRole ?? + workspacePolicy?.memberDefaultDocRole ?? + 'manager', + publicRole: publicRole === 'external' ? 'external' : undefined, + visibility, + sharingEnabled, + previewEnabled: + visibility === 'public' || + policy?.urlPreviewEnabled || + urlPreviewEnabled, + }; + }), + }; + } + + private get cache(): PermissionRequestCache { + if (!this.cls) { + return createPermissionRequestCache(); + } + + if (typeof this.cls.isActive === 'function' && !this.cls.isActive()) { + return createPermissionRequestCache(); + } + + const existing = this.cls.get(CACHE_KEY) as + | PermissionRequestCache + | undefined; + if (existing) { + return existing; + } + + const created = createPermissionRequestCache(); + this.cls.set(CACHE_KEY, created); + return created; + } + + private memo( + map: Map | T>, + key: string, + load: () => Promise + ) { + const cached = map.get(key); + if (cached) { + return Promise.resolve(cached); + } + const promise = load(); + map.set(key, promise); + return promise; + } + + private workspaceMember(workspaceId: string, userId: string) { + return this.memo( + this.cache.workspaceMember, + cacheKey([workspaceId, userId]), + () => this.models.workspaceUser.get(workspaceId, userId) + ); + } + + private workspacePolicy(workspaceId: string) { + return this.memo(this.cache.workspacePolicy, workspaceId, () => + this.models.workspace.get(workspaceId) + ); + } + + private async workspaceRuntime(workspaceId: string) { + return this.memo(this.cache.workspaceRuntime, workspaceId, () => + this.models.workspaceRuntimeState.get(workspaceId).then(async state => { + if (state.known || !state.stale) { + return state; + } + + const quotaState = await this.newWorkspaceRuntime(workspaceId); + if (!quotaState.known) { + return state; + } + + return { + workspaceId, + known: quotaState.known, + stale: quotaState.stale, + readonly: quotaState.readonly, + readonlyReasons: quotaState.readonlyReasons, + updatedAt: null, + lastReconciledAt: null, + staleAfter: quotaState.staleAfter, + }; + }) + ); + } + + invalidateWorkspaceQuotaRuntime(workspaceId: string) { + this.cache.workspaceQuotaRuntime.delete(workspaceId); + } + + private newWorkspaceRuntime(workspaceId: string) { + return this.memo( + this.cache.workspaceQuotaRuntime, + workspaceId, + async () => { + const rows = await this.db.$queryRaw` + SELECT + known, + stale, + readonly, + readonly_reasons AS "readonlyReasons", + stale_after AS "staleAfter" + FROM effective_workspace_quota_states + WHERE workspace_id = ${workspaceId} + LIMIT 1 + `; + const state = rows[0]; + if (!state) { + return { + known: false, + stale: true, + readonly: false, + readonlyReasons: [], + staleAfter: null, + }; + } + + return { + ...state, + stale: + state.stale || + (state.staleAfter !== null && state.staleAfter <= new Date()), + }; + } + ); + } + + private docPolicies(workspaceId: string, docIds: string[]) { + const uniqueDocIds = [...new Set(docIds)]; + return this.memo( + this.cache.docPolicies, + cacheKey([workspaceId, ...uniqueDocIds]), + () => this.models.doc.findDefaultRoles(workspaceId, uniqueDocIds) + ); + } + + private docGrants(workspaceId: string, docIds: string[], userId: string) { + const uniqueDocIds = [...new Set(docIds)]; + return this.memo( + this.cache.docGrants, + cacheKey([workspaceId, userId, ...uniqueDocIds]), + () => this.models.docUser.findMany(workspaceId, uniqueDocIds, userId) + ); + } + + private async newWorkspaceMember(workspaceId: string, userId: string) { + const rows = await this.db.$queryRaw` + SELECT role, state + FROM workspace_members + WHERE workspace_id = ${workspaceId} + AND user_id = ${userId} + AND state = 'active' + LIMIT 1 + `; + return rows[0] ?? null; + } + + private async newWorkspacePolicy(workspaceId: string) { + const rows = await this.db.$queryRaw` + SELECT + visibility, + sharing_enabled AS "sharingEnabled", + url_preview_enabled AS "urlPreviewEnabled", + member_default_doc_role AS "memberDefaultDocRole" + FROM workspace_access_policies + WHERE workspace_id = ${workspaceId} + LIMIT 1 + `; + return rows[0] ?? null; + } + + async workspaceExists(workspaceId: string) { + const workspace = await this.db.workspace.findUnique({ + where: { id: workspaceId }, + select: { id: true }, + }); + return !!workspace; + } + + private async newDocPolicies(workspaceId: string, docIds: string[]) { + if (docIds.length === 0) { + return []; + } + return await this.db.$queryRaw` + SELECT + doc_id AS "docId", + visibility, + public_role AS "publicRole", + member_default_role AS "memberDefaultRole", + url_preview_enabled AS "urlPreviewEnabled" + FROM doc_access_policies + WHERE workspace_id = ${workspaceId} + AND doc_id = ANY(${[...new Set(docIds)]}) + `; + } + + private async newDocGrants( + workspaceId: string, + docIds: string[], + userId: string + ) { + if (docIds.length === 0) { + return []; + } + return await this.db.$queryRaw` + SELECT doc_id AS "docId", role + FROM doc_grants + WHERE workspace_id = ${workspaceId} + AND principal_type = 'user' + AND principal_id = ${userId} + AND doc_id = ANY(${[...new Set(docIds)]}) + `; + } +} diff --git a/packages/backend/server/src/core/permission/context.ts b/packages/backend/server/src/core/permission/context.ts new file mode 100644 index 0000000000..c61192da01 --- /dev/null +++ b/packages/backend/server/src/core/permission/context.ts @@ -0,0 +1,125 @@ +import { WorkspaceMemberStatus } from '@prisma/client'; + +import type { + PermissionDocRole, + PermissionEvaluationInputV1, + PermissionEvaluationOutputV1, + PermissionWorkspaceRole, +} from '../../native'; +import { DocRole, WorkspaceRole } from './types'; + +export type PermissionRuntimeState = NonNullable< + PermissionEvaluationInputV1['runtime'] +>; + +export type PermissionWorkspaceContext = NonNullable< + PermissionEvaluationInputV1['workspace'] +>; + +export type PermissionDocContext = NonNullable< + NonNullable[number] +>; + +export type PermissionLegacyRoleBoundary = { + resourceOwnerRole: PermissionDocRole | PermissionWorkspaceRole | null; + effectiveRole: PermissionDocRole | PermissionWorkspaceRole | null; + legacyApiRole: DocRole | WorkspaceRole | null; +}; + +const WORKSPACE_ROLE_TO_NATIVE = new Map< + WorkspaceRole, + PermissionWorkspaceRole +>([ + [WorkspaceRole.External, 'external'], + [WorkspaceRole.Collaborator, 'member'], + [WorkspaceRole.Admin, 'admin'], + [WorkspaceRole.Owner, 'owner'], +]); + +const DOC_ROLE_TO_NATIVE = new Map([ + [DocRole.None, 'none'], + [DocRole.External, 'external'], + [DocRole.Reader, 'reader'], + [DocRole.Commenter, 'commenter'], + [DocRole.Editor, 'editor'], + [DocRole.Manager, 'manager'], + [DocRole.Owner, 'owner'], +]); + +const NATIVE_WORKSPACE_ROLE_TO_LEGACY = new Map< + PermissionWorkspaceRole, + WorkspaceRole +>([ + ['external', WorkspaceRole.External], + ['member', WorkspaceRole.Collaborator], + ['admin', WorkspaceRole.Admin], + ['owner', WorkspaceRole.Owner], +]); + +const NATIVE_DOC_ROLE_TO_LEGACY = new Map([ + ['none', DocRole.None], + ['external', DocRole.External], + ['reader', DocRole.Reader], + ['commenter', DocRole.Commenter], + ['editor', DocRole.Editor], + ['manager', DocRole.Manager], + ['owner', DocRole.Owner], +]); + +export function toNativeWorkspaceRole(role: WorkspaceRole | null | undefined) { + return role == null ? undefined : WORKSPACE_ROLE_TO_NATIVE.get(role); +} + +export function toNativeDocRole(role: DocRole | null | undefined) { + return role == null ? undefined : DOC_ROLE_TO_NATIVE.get(role); +} + +export function toNativeExplicitDocGrantRole(role: DocRole | null | undefined) { + if (role === DocRole.None || role === DocRole.External) { + return undefined; + } + return toNativeDocRole(role); +} + +export function toNativeMemberState(status?: WorkspaceMemberStatus | null) { + switch (status) { + case WorkspaceMemberStatus.Accepted: + return 'active'; + case WorkspaceMemberStatus.UnderReview: + return 'waiting_review'; + case WorkspaceMemberStatus.AllocatingSeat: + case WorkspaceMemberStatus.NeedMoreSeat: + case WorkspaceMemberStatus.NeedMoreSeatAndReview: + return 'waiting_seat'; + case WorkspaceMemberStatus.Pending: + return 'pending'; + default: + return undefined; + } +} + +export function workspaceLegacyBoundary( + workspace: PermissionEvaluationOutputV1['workspace'] +): PermissionLegacyRoleBoundary { + const effectiveRole = workspace.effectiveRole ?? null; + return { + resourceOwnerRole: workspace.resourceOwnerRole ?? null, + effectiveRole, + legacyApiRole: effectiveRole + ? (NATIVE_WORKSPACE_ROLE_TO_LEGACY.get(effectiveRole) ?? null) + : null, + }; +} + +export function docLegacyBoundary( + doc: PermissionEvaluationOutputV1['docs'][number] +): PermissionLegacyRoleBoundary { + const effectiveRole = doc.effectiveRole ?? null; + return { + resourceOwnerRole: doc.resourceOwnerRole ?? null, + effectiveRole, + legacyApiRole: effectiveRole + ? (NATIVE_DOC_ROLE_TO_LEGACY.get(effectiveRole) ?? null) + : null, + }; +} diff --git a/packages/backend/server/src/core/permission/controller.ts b/packages/backend/server/src/core/permission/controller.ts deleted file mode 100644 index 21b258d993..0000000000 --- a/packages/backend/server/src/core/permission/controller.ts +++ /dev/null @@ -1,53 +0,0 @@ -import { Logger, OnModuleInit } from '@nestjs/common'; - -import type { - Resource, - ResourceAction, - ResourceRole, - ResourceType, -} from './resource'; - -const ACTION_CHECKER_PROVIDERS = new Map>(); - -function registerAccessController( - type: Type, - provider: AccessController -) { - ACTION_CHECKER_PROVIDERS.set(type, provider); -} - -export function getAccessController( - type: Type -): AccessController { - const provider = ACTION_CHECKER_PROVIDERS.get(type); - if (!provider) { - throw new Error(`No action checker provider for type ${type}`); - } - return provider; -} - -export abstract class AccessController< - Type extends ResourceType, -> implements OnModuleInit { - protected abstract readonly type: Type; - protected logger = new Logger(AccessController.name); - - onModuleInit() { - registerAccessController(this.type, this); - } - - abstract assert( - resource: Resource, - action: ResourceAction - ): Promise; - - abstract can( - resource: Resource, - action: ResourceAction - ): Promise; - - abstract role(resource: Resource): Promise<{ - role: ResourceRole | null; - permissions: Record, boolean>; - }>; -} diff --git a/packages/backend/server/src/core/permission/diagnostic.ts b/packages/backend/server/src/core/permission/diagnostic.ts new file mode 100644 index 0000000000..06b81038fe --- /dev/null +++ b/packages/backend/server/src/core/permission/diagnostic.ts @@ -0,0 +1,326 @@ +import { Inject, Injectable, Optional } from '@nestjs/common'; + +import { metrics } from '../../base'; +import type { PermissionEvaluationOutputV1 } from '../../native'; +import { docLegacyBoundary, workspaceLegacyBoundary } from './context'; +import { + PermissionContextLoader, + type PermissionDocAction, + type PermissionWorkspaceAction, +} from './context-loader'; +import { PermissionService } from './service'; +import { PermissionSqlPredicateBuilder } from './sql-predicate'; + +export const PERMISSION_SHADOW_MISMATCH_CATEGORIES = [ + 'legacy_compat_delta', + 'projection', + 'rust_rule', + 'loader', + 'sql_predicate', + 'legacy_api_role_mapping', + 'preview_read_mapping', + 'runtime_state', + 'projection_or_loader', +] as const; + +type PermissionShadowMismatchCategory = + (typeof PERMISSION_SHADOW_MISMATCH_CATEGORIES)[number]; + +@Injectable() +export class PermissionDiagnosticService { + constructor( + private readonly loader: PermissionContextLoader, + private readonly permission: PermissionService, + @Optional() + @Inject(PermissionSqlPredicateBuilder) + private readonly sqlPredicate = new PermissionSqlPredicateBuilder() + ) {} + + async shadowDocPermissions(input: { + userId?: string; + workspaceId: string; + docs: Array<{ docId: string; actions: PermissionDocAction[] }>; + allowLocal?: boolean; + expectedDeltaCategory?: PermissionShadowMismatchCategory; + }) { + const [legacyOutput, newOutput] = await Promise.all([ + this.loader.load(input).then(input => this.permission.evaluate(input)), + this.loader + .loadFromNewTables(input) + .then(input => this.permission.evaluate(input)), + ]); + + const legacy = legacyOutput.docs.map(doc => ({ + docId: doc.docId, + ...docLegacyBoundary(doc), + decisions: doc.decisions, + })); + const current = newOutput.docs.map(doc => ({ + docId: doc.docId, + ...docLegacyBoundary(doc), + decisions: doc.decisions, + })); + const matched = JSON.stringify(legacy) === JSON.stringify(current); + const mismatchType = matched + ? null + : (input.expectedDeltaCategory ?? + this.classifyDocShadowMismatch(legacy, current)); + this.recordShadowMismatch('doc', mismatchType); + + return { + matched, + legacy, + current, + mismatchType, + }; + } + + async shadowWorkspacePermissions(input: { + userId?: string; + workspaceId: string; + actions: PermissionWorkspaceAction[]; + allowLocal?: boolean; + expectedDeltaCategory?: PermissionShadowMismatchCategory; + }) { + const legacyInput = { + userId: input.userId, + workspaceId: input.workspaceId, + workspaceActions: input.actions, + allowLocal: input.allowLocal, + }; + const [legacyOutput, newOutput] = await Promise.all([ + this.loader + .load(legacyInput) + .then(input => this.permission.evaluate(input)), + this.loader + .loadFromNewTables(legacyInput) + .then(input => this.permission.evaluate(input)), + ]); + + const legacy = { + ...workspaceLegacyBoundary(legacyOutput.workspace), + decisions: legacyOutput.workspace.decisions, + }; + const current = { + ...workspaceLegacyBoundary(newOutput.workspace), + decisions: newOutput.workspace.decisions, + }; + const matched = JSON.stringify(legacy) === JSON.stringify(current); + const mismatchType = matched + ? null + : (input.expectedDeltaCategory ?? + this.classifyShadowMismatch(legacyOutput, newOutput)); + this.recordShadowMismatch('workspace', mismatchType); + + return { + matched, + legacy, + current, + mismatchType, + }; + } + + async shadowSqlDocRead(input: { + userId: string; + workspaceId: string; + docs: Array<{ docId: string }>; + sqlReadableDocIds: string[]; + allowLocal?: boolean; + expectedDeltaCategory?: PermissionShadowMismatchCategory; + }) { + const rustOutput = this.permission.evaluate( + await this.loader.loadFromNewTables({ + userId: input.userId, + workspaceId: input.workspaceId, + docs: input.docs.map(doc => ({ + docId: doc.docId, + actions: ['Doc.Read'], + })), + allowLocal: input.allowLocal, + }) + ); + const rustReadable = new Set( + rustOutput.docs + .filter(doc => doc.decisions[0]?.allowed) + .map(doc => doc.docId) + ); + const sqlReadable = new Set(input.sqlReadableDocIds); + const missingInSql = [...rustReadable].filter(id => !sqlReadable.has(id)); + const extraInSql = [...sqlReadable].filter(id => !rustReadable.has(id)); + const mismatchType = + missingInSql.length || extraInSql.length + ? (input.expectedDeltaCategory ?? 'sql_predicate') + : null; + this.recordShadowMismatch('sql_predicate', mismatchType); + + return { + matched: mismatchType === null, + predicate: this.sqlPredicate.docReadableByNewTables({ + workspaceId: input.workspaceId, + userId: input.userId, + action: 'Doc.Read', + }), + rustReadableDocIds: [...rustReadable], + sqlReadableDocIds: [...sqlReadable], + missingInSql, + extraInSql, + mismatchType, + }; + } + + async shadowPreviewDoc(input: { + userId?: string; + workspaceId: string; + docId: string; + allowLocal?: boolean; + }) { + const result = await this.shadowDocPermissions({ + ...input, + docs: [{ docId: input.docId, actions: ['Doc.Preview', 'Doc.Read'] }], + }); + const legacy = result.legacy[0]; + const current = result.current[0]; + const legacyPreviewAllowed = legacy?.decisions.find( + decision => decision.action === 'Doc.Preview' + )?.allowed; + const legacyReadAllowed = legacy?.decisions.find( + decision => decision.action === 'Doc.Read' + )?.allowed; + const previewAllowed = current?.decisions.find( + decision => decision.action === 'Doc.Preview' + )?.allowed; + const readAllowed = current?.decisions.find( + decision => decision.action === 'Doc.Read' + )?.allowed; + const mismatchType = + legacyPreviewAllowed !== previewAllowed || + (previewAllowed && readAllowed && !legacyReadAllowed) + ? 'preview_read_mapping' + : result.mismatchType; + this.recordShadowMismatch('preview', mismatchType); + + return { + ...result, + matched: result.matched && mismatchType === null, + mismatchType, + }; + } + + async shadowPreviewWorkspace(input: { + userId?: string; + workspaceId: string; + allowLocal?: boolean; + }) { + const result = await this.shadowWorkspacePermissions({ + ...input, + actions: ['Workspace.Preview', 'Workspace.Read'], + }); + const legacyPreviewAllowed = result.legacy.decisions.find( + decision => decision.action === 'Workspace.Preview' + )?.allowed; + const legacyReadAllowed = result.legacy.decisions.find( + decision => decision.action === 'Workspace.Read' + )?.allowed; + const previewAllowed = result.current.decisions.find( + decision => decision.action === 'Workspace.Preview' + )?.allowed; + const readAllowed = result.current.decisions.find( + decision => decision.action === 'Workspace.Read' + )?.allowed; + const mismatchType = + legacyPreviewAllowed !== previewAllowed || + (previewAllowed && readAllowed && !legacyReadAllowed) + ? 'preview_read_mapping' + : result.mismatchType; + this.recordShadowMismatch('preview', mismatchType); + + return { + ...result, + matched: result.matched && mismatchType === null, + mismatchType, + }; + } + + private classifyShadowMismatch( + legacyOutput: PermissionEvaluationOutputV1, + newOutput: PermissionEvaluationOutputV1 + ) { + if (JSON.stringify(legacyOutput) === JSON.stringify(newOutput)) { + return null; + } + const legacyRestrictions = + JSON.stringify(legacyOutput).includes('runtime_'); + const newRestrictions = JSON.stringify(newOutput).includes('runtime_'); + if (legacyRestrictions || newRestrictions) { + return 'runtime_state'; + } + if (legacyOutput.docs.length !== newOutput.docs.length) { + return 'loader'; + } + if (JSON.stringify(legacyOutput.docs) !== JSON.stringify(newOutput.docs)) { + return 'rust_rule'; + } + return 'projection'; + } + + private classifyDocShadowMismatch( + legacy: Array< + ReturnType & { decisions: unknown } + >, + current: Array< + ReturnType & { decisions: unknown } + > + ) { + if (JSON.stringify(legacy) === JSON.stringify(current)) { + return null; + } + + const legacyApi = legacy.map(doc => ({ + effectiveRole: doc.effectiveRole, + legacyApiRole: doc.legacyApiRole, + resourceOwnerRole: doc.resourceOwnerRole, + })); + const currentApi = current.map(doc => ({ + effectiveRole: doc.effectiveRole, + legacyApiRole: doc.legacyApiRole, + resourceOwnerRole: doc.resourceOwnerRole, + })); + if (JSON.stringify(legacyApi) !== JSON.stringify(currentApi)) { + return 'legacy_api_role_mapping'; + } + + if ( + JSON.stringify(legacy).includes('runtime_') || + JSON.stringify(current).includes('runtime_') + ) { + return 'runtime_state'; + } + + if (legacy.length !== current.length) { + return 'loader'; + } + + const legacyDecisions = legacy.map(doc => doc.decisions); + const currentDecisions = current.map(doc => doc.decisions); + if (JSON.stringify(legacyDecisions) !== JSON.stringify(currentDecisions)) { + return 'rust_rule'; + } + + return 'projection'; + } + + private recordShadowMismatch( + scope: string, + category: PermissionShadowMismatchCategory | null + ) { + if (!category) { + return; + } + + metrics.permission + .counter('shadow_mismatches', { + description: 'Permission shadow-read mismatch count', + }) + .add(1, { scope, category }); + } +} diff --git a/packages/backend/server/src/core/permission/doc.ts b/packages/backend/server/src/core/permission/doc.ts deleted file mode 100644 index 011a1bb0be..0000000000 --- a/packages/backend/server/src/core/permission/doc.ts +++ /dev/null @@ -1,75 +0,0 @@ -import { Injectable } from '@nestjs/common'; - -import { DocActionDenied } from '../../base'; -import { AccessController, getAccessController } from './controller'; -import { WorkspacePolicyService } from './policy'; -import type { Resource } from './resource'; -import { - DocAction, - docActionRequiredRole, - DocRole, - mapDocRoleToPermissions, -} from './types'; -import { WorkspaceAccessController } from './workspace'; - -@Injectable() -export class DocAccessController extends AccessController<'doc'> { - protected readonly type = 'doc'; - constructor(private readonly policy: WorkspacePolicyService) { - super(); - } - - async role(resource: Resource<'doc'>) { - const role = await this.getRole(resource); - const permissions = await this.policy.applyDocPermissions( - resource.workspaceId, - mapDocRoleToPermissions(role) - ); - const sharingAllowed = await this.policy.canPublishDoc( - resource.workspaceId - ); - if (!sharingAllowed) { - permissions['Doc.Publish'] = false; - } - - return { role, permissions }; - } - - async can(resource: Resource<'doc'>, action: DocAction) { - const { permissions, role } = await this.role(resource); - const allow = permissions[action] || false; - - if (!allow) { - this.logger.debug('Doc access check failed', { - action, - resource, - role, - requiredRole: docActionRequiredRole(action), - }); - } - - return allow; - } - - async assert(resource: Resource<'doc'>, action: DocAction) { - const allow = await this.can(resource, action); - - if (!allow) { - throw new DocActionDenied({ - docId: resource.docId, - spaceId: resource.workspaceId, - action, - }); - } - } - - async getRole(payload: Resource<'doc'>): Promise { - const workspaceController = getAccessController( - 'ws' - ) as WorkspaceAccessController; - const docRoles = await workspaceController.getDocRoles(payload, [ - payload.docId, - ]); - return docRoles[0]; - } -} diff --git a/packages/backend/server/src/core/permission/index.ts b/packages/backend/server/src/core/permission/index.ts index 321f43e9a2..6a013d473d 100644 --- a/packages/backend/server/src/core/permission/index.ts +++ b/packages/backend/server/src/core/permission/index.ts @@ -1,27 +1,50 @@ +import './config'; + import { Module } from '@nestjs/common'; import { QuotaServiceModule } from '../quota/service.module'; import { AccessControllerBuilder } from './builder'; -import { DocAccessController } from './doc'; +import { PermissionContextLoader } from './context-loader'; +import { PermissionDiagnosticService } from './diagnostic'; import { EventsListener } from './event'; import { WorkspacePolicyService } from './policy'; -import { WorkspaceAccessController } from './workspace'; +import { PermissionProjectionChecker } from './projection-checker'; +import { PermissionService } from './service'; +import { PermissionSqlPredicateBuilder } from './sql-predicate'; @Module({ imports: [QuotaServiceModule], providers: [ - WorkspaceAccessController, - DocAccessController, AccessControllerBuilder, EventsListener, WorkspacePolicyService, + PermissionProjectionChecker, + PermissionSqlPredicateBuilder, + PermissionContextLoader, + PermissionDiagnosticService, + PermissionService, + ], + exports: [ + AccessControllerBuilder, + WorkspacePolicyService, + PermissionProjectionChecker, + PermissionSqlPredicateBuilder, + PermissionDiagnosticService, + PermissionService, ], - exports: [AccessControllerBuilder, WorkspacePolicyService], }) export class PermissionModule {} -export { AccessControllerBuilder as AccessController } from './builder'; +export { AccessControllerBuilder as PermissionAccess } from './builder'; +export { PermissionContextLoader } from './context-loader'; +export { + PERMISSION_SHADOW_MISMATCH_CATEGORIES, + PermissionDiagnosticService, +} from './diagnostic'; export { WorkspacePolicyService } from './policy'; +export { PermissionProjectionChecker } from './projection-checker'; +export { PermissionService } from './service'; +export { PermissionSqlPredicateBuilder } from './sql-predicate'; export { DOC_ACTIONS, type DocAction, diff --git a/packages/backend/server/src/core/permission/policy.ts b/packages/backend/server/src/core/permission/policy.ts index b2b05c302d..22be5a6563 100644 --- a/packages/backend/server/src/core/permission/policy.ts +++ b/packages/backend/server/src/core/permission/policy.ts @@ -1,29 +1,15 @@ import { Injectable } from '@nestjs/common'; import { Transactional } from '@nestjs-cls/transactional'; -import { - DocActionDenied, - OnEvent, - OwnerCanNotLeaveWorkspace, - SpaceAccessDenied, -} from '../../base'; +import { OnEvent } from '../../base'; import { Models, WorkspaceRole } from '../../models'; -import { QuotaService } from '../quota/service'; -import { getAccessController } from './controller'; -import type { Resource } from './resource'; -import { - type DocAction, - type DocActionPermissions, - mapWorkspaceRoleToPermissions, - type WorkspaceAction, - type WorkspaceActionPermissions, -} from './types'; +import { QuotaStateService } from '../quota/state'; export type WorkspaceReadonlyReason = 'member_overflow' | 'storage_overflow'; type WorkspaceQuotaSnapshot = Awaited< - ReturnType + ReturnType > & { - ownerQuota?: string; + readonlyReasons: WorkspaceReadonlyReason[]; }; export type WorkspaceState = { @@ -35,35 +21,6 @@ export type WorkspaceState = { usesFallbackOwnerQuota: boolean; }; -const READONLY_WORKSPACE_ACTIONS: WorkspaceAction[] = [ - 'Workspace.CreateDoc', - 'Workspace.Settings.Update', - 'Workspace.Properties.Create', - 'Workspace.Properties.Update', - 'Workspace.Properties.Delete', - 'Workspace.Blobs.Write', -]; - -const READONLY_DOC_ACTIONS: DocAction[] = [ - 'Doc.Update', - 'Doc.Duplicate', - 'Doc.Publish', - 'Doc.Comments.Create', - 'Doc.Comments.Update', - 'Doc.Comments.Resolve', -]; - -const READONLY_WORKSPACE_FEATURE = - 'quota_exceeded_readonly_workspace_v1' as const; - -type WorkspaceRoleChecker = { - getRole(resource: Resource<'ws'>): Promise; - docRoles( - resource: Resource<'ws'>, - docIds: string[] - ): Promise }>>; -}; - declare global { interface Events { 'workspace.blobs.updated': { @@ -76,39 +33,23 @@ declare global { export class WorkspacePolicyService { constructor( private readonly models: Models, - private readonly quota: QuotaService + private readonly quotaState: QuotaStateService ) {} async getWorkspaceState(workspaceId: string): Promise { - const [isTeamWorkspace, isUnlimitedWorkspace, quota] = await Promise.all([ - this.models.workspace.isTeamWorkspace(workspaceId), - this.models.workspaceFeature.has(workspaceId, 'unlimited_workspace'), - this.quota.getWorkspaceQuotaWithUsage(workspaceId), - ]); + const quota = + await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); const quotaSnapshot = quota as WorkspaceQuotaSnapshot; - const readonlyReasons: WorkspaceReadonlyReason[] = []; - const usesFallbackOwnerQuota = - !!quotaSnapshot.ownerQuota && !isUnlimitedWorkspace; - - if (usesFallbackOwnerQuota && quotaSnapshot.overcapacityMemberCount > 0) { - readonlyReasons.push('member_overflow'); - } - - if ( - usesFallbackOwnerQuota && - quotaSnapshot.usedStorageQuota > quotaSnapshot.storageQuota - ) { - readonlyReasons.push('storage_overflow'); - } + const readonlyReasons = quotaSnapshot.readonlyReasons; return { - isTeamWorkspace, + isTeamWorkspace: ['team', 'selfhost_team'].includes(quotaSnapshot.plan), isReadonly: readonlyReasons.length > 0, readonlyReasons, canRecoverByRemovingMembers: readonlyReasons.includes('member_overflow'), canRecoverByDeletingBlobs: readonlyReasons.includes('storage_overflow'), - usesFallbackOwnerQuota, + usesFallbackOwnerQuota: quotaSnapshot.usesOwnerQuota, }; } @@ -126,286 +67,19 @@ export class WorkspacePolicyService { } async reconcileWorkspaceQuotaState(workspaceId: string) { - const [state, isReadonlyFeatureEnabled] = await Promise.all([ - this.getWorkspaceState(workspaceId), - this.models.workspaceFeature.has(workspaceId, READONLY_WORKSPACE_FEATURE), - ]); - - if (state.isReadonly && !isReadonlyFeatureEnabled) { - await this.models.workspaceFeature.add( - workspaceId, - READONLY_WORKSPACE_FEATURE, - `workspace recovery mode: ${state.readonlyReasons.join(',')}` - ); - } else if (!state.isReadonly && isReadonlyFeatureEnabled) { - await this.models.workspaceFeature.remove( - workspaceId, - READONLY_WORKSPACE_FEATURE - ); - } - - return state; + return await this.getWorkspaceState(workspaceId); } - async isWorkspaceReadonly(workspaceId: string) { - const hasReadonlyFeature = await this.models.workspaceFeature.has( - workspaceId, - READONLY_WORKSPACE_FEATURE - ); - - if (!hasReadonlyFeature) { - return false; - } - - const state = await this.getWorkspaceState(workspaceId); - if (!state.isReadonly) { - await this.models.workspaceFeature.remove( - workspaceId, - READONLY_WORKSPACE_FEATURE - ); - return false; - } - - return true; - } - - async isSharingEnabled(workspaceId: string) { - return await this.models.workspace.allowSharing(workspaceId); - } - - async canReadWorkspaceByPublicFlag(workspaceId: string) { - const workspace = await this.models.workspace.get(workspaceId); - return !!workspace?.public && (workspace.enableSharing ?? true); - } - - async canReadWorkspaceBySharedDocs(workspaceId: string) { - const [sharingEnabled, hasPublicDocs] = await Promise.all([ - this.isSharingEnabled(workspaceId), - this.models.doc.hasPublic(workspaceId), - ]); - - return sharingEnabled && hasPublicDocs; - } - - async canReadSharedDoc(workspaceId: string, docId: string) { - const [sharingEnabled, isPublicDoc] = await Promise.all([ - this.isSharingEnabled(workspaceId), - this.models.doc.isPublic(workspaceId, docId), - ]); - - return sharingEnabled && isPublicDoc; - } - - async canPreviewDoc(workspaceId: string, docId: string) { - const [sharingEnabled, canReadSharedDoc, allowUrlPreview] = - await Promise.all([ - this.isSharingEnabled(workspaceId), - this.canReadSharedDoc(workspaceId, docId), - this.models.workspace.allowUrlPreview(workspaceId), - ]); - - return sharingEnabled && (canReadSharedDoc || allowUrlPreview); - } - - async canPreviewWorkspace(workspaceId: string) { - const [sharingEnabled, allowUrlPreview] = await Promise.all([ - this.isSharingEnabled(workspaceId), - this.models.workspace.allowUrlPreview(workspaceId), - ]); - - return sharingEnabled && allowUrlPreview; - } - - async canPublishDoc(workspaceId: string) { - return await this.isSharingEnabled(workspaceId); - } - - async applyWorkspacePermissions( - workspaceId: string, - permissions: WorkspaceActionPermissions - ) { - if (!(await this.isWorkspaceReadonly(workspaceId))) { - return permissions; - } - - const next = { ...permissions }; - READONLY_WORKSPACE_ACTIONS.forEach(action => { - next[action] = false; - }); - return next; - } - - async applyDocPermissions( - workspaceId: string, - permissions: DocActionPermissions - ) { - if (!(await this.isWorkspaceReadonly(workspaceId))) { - return permissions; - } - - const next = { ...permissions }; - READONLY_DOC_ACTIONS.forEach(action => { - next[action] = false; - }); - return next; - } - - async assertWorkspaceActionAllowed( - workspaceId: string, - action: WorkspaceAction - ) { - if ( - READONLY_WORKSPACE_ACTIONS.includes(action) && - (await this.isWorkspaceReadonly(workspaceId)) - ) { - throw new SpaceAccessDenied({ spaceId: workspaceId }); - } - } - - async assertDocActionAllowed( - workspaceId: string, - docId: string, - action: DocAction - ) { - if ( - READONLY_DOC_ACTIONS.includes(action) && - (await this.isWorkspaceReadonly(workspaceId)) - ) { - throw new DocActionDenied({ - action, - docId, - spaceId: workspaceId, - }); - } - } - - async assertWorkspaceRoleAction( - userId: string, - workspaceId: string, - action: WorkspaceAction - ) { - const checker = getAccessController( - 'ws' - ) as unknown as WorkspaceRoleChecker; - const role = await checker.getRole({ userId, workspaceId }); - const permissions = mapWorkspaceRoleToPermissions(role); - - if (!permissions[action]) { - throw new SpaceAccessDenied({ spaceId: workspaceId }); - } - } - - async assertDocRoleAction( - userId: string, - workspaceId: string, - docId: string, - action: DocAction - ) { - const checker = getAccessController( - 'ws' - ) as unknown as WorkspaceRoleChecker; - const [role] = await checker.docRoles({ userId, workspaceId }, [docId]); - - if (!role?.permissions[action]) { - throw new DocActionDenied({ - action, - docId, - spaceId: workspaceId, - }); - } - } - - async assertCanUploadBlob(userId: string, workspaceId: string) { - await this.assertWorkspaceRoleAction( - userId, - workspaceId, - 'Workspace.Blobs.Write' - ); - await this.assertWorkspaceActionAllowed( - workspaceId, - 'Workspace.Blobs.Write' - ); - } - - async assertCanDeleteBlob(userId: string, workspaceId: string) { - await this.assertWorkspaceRoleAction( - userId, - workspaceId, - 'Workspace.Blobs.Write' - ); - } - - async assertCanInviteMembers(workspaceId: string) { - if (await this.isWorkspaceReadonly(workspaceId)) { - throw new SpaceAccessDenied({ spaceId: workspaceId }); - } - } - - async assertCanRevokeMember( - userId: string, - workspaceId: string, - role: WorkspaceRole - ) { - await this.assertWorkspaceRoleAction( - userId, - workspaceId, - role === WorkspaceRole.Admin - ? 'Workspace.Administrators.Manage' - : 'Workspace.Users.Manage' - ); - } - - @Transactional() async handleTeamPlanCanceled(workspaceId: string) { - await this.models.workspaceUser.deleteNonAccepted(workspaceId); - await this.models.workspaceUser.demoteAcceptedAdmins(workspaceId); - await this.models.workspaceFeature.remove(workspaceId, 'team_plan_v1'); + await this.cleanupTeamPlanCanceled(workspaceId); return await this.reconcileWorkspaceQuotaState(workspaceId); } - async assertCanUnpublishDoc( - userId: string, - workspaceId: string, - docId: string - ) { - await this.assertDocRoleAction(userId, workspaceId, docId, 'Doc.Publish'); - } - - async assertCanPublishDoc( - userId: string, - workspaceId: string, - docId: string - ) { - await this.assertDocRoleAction(userId, workspaceId, docId, 'Doc.Publish'); - await this.assertDocActionAllowed(workspaceId, docId, 'Doc.Publish'); - - if (!(await this.canPublishDoc(workspaceId))) { - throw new DocActionDenied({ - action: 'Doc.Publish', - docId, - spaceId: workspaceId, - }); - } - } - - async assertCanManageInviteLink(userId: string, workspaceId: string) { - await this.assertWorkspaceRoleAction( - userId, - workspaceId, - 'Workspace.Users.Manage' - ); - } - - async assertCanLeaveWorkspace(userId: string, workspaceId: string) { - const role = await this.models.workspaceUser.getActive(workspaceId, userId); - - if (!role) { - throw new SpaceAccessDenied({ spaceId: workspaceId }); - } - - if (role.type === WorkspaceRole.Owner) { - throw new OwnerCanNotLeaveWorkspace(); - } + @Transactional() + private async cleanupTeamPlanCanceled(workspaceId: string) { + await this.models.workspaceUser.deleteNonAccepted(workspaceId); + await this.models.workspaceUser.demoteAcceptedAdmins(workspaceId); + await this.models.workspaceFeature.remove(workspaceId, 'team_plan_v1'); } @OnEvent('workspace.members.updated') diff --git a/packages/backend/server/src/core/permission/projection-checker.ts b/packages/backend/server/src/core/permission/projection-checker.ts new file mode 100644 index 0000000000..ccb5ef928e --- /dev/null +++ b/packages/backend/server/src/core/permission/projection-checker.ts @@ -0,0 +1,166 @@ +import { Injectable, Optional } from '@nestjs/common'; +import { PrismaClient } from '@prisma/client'; + +import { Models } from '../../models'; +import { + PermissionContextLoader, + type PermissionDocAction, + type PermissionWorkspaceAction, +} from './context-loader'; +import { PermissionService } from './service'; + +type ProjectionDecisionSample = { + category: string; + workspaceId: string; + docId: string | null; + userId: string | null; + workspaceActions: string[] | null; + docActions: string[] | null; +}; + +@Injectable() +export class PermissionProjectionChecker { + constructor( + private readonly db: PrismaClient, + private readonly models: Models, + @Optional() + private readonly loader?: PermissionContextLoader, + @Optional() + private readonly permission?: PermissionService + ) {} + + async checkLegacyProjection() { + const report = + await this.models.permissionProjection.checkLegacyProjection(); + return { + ...report, + oldNewDecisionMismatch: await this.checkOldNewLoaderDecisionMismatch(), + }; + } + + private async checkOldNewLoaderDecisionMismatch() { + const { loader, permission } = this; + if (!loader || !permission) { + return 0; + } + + const samples = await this.db.$queryRaw` + ( + SELECT + 'active_member_doc' AS category, + old_member.workspace_id AS "workspaceId", + old_doc.page_id AS "docId", + old_member.user_id AS "userId", + NULL::text[] AS "workspaceActions", + ARRAY['Doc.Read', 'Doc.Preview']::text[] AS "docActions" + FROM workspace_user_permissions old_member + INNER JOIN workspace_pages old_doc + ON old_doc.workspace_id = old_member.workspace_id + WHERE old_member.status = 'Accepted'::"WorkspaceMemberStatus" + AND affine_permission_legacy_workspace_role(old_member.type) IS NOT NULL + AND affine_permission_legacy_default_doc_role(old_doc."defaultRole") IS NOT NULL + ORDER BY md5(old_member.workspace_id || ':' || old_doc.page_id || ':' || old_member.user_id) + LIMIT 80 + ) + UNION ALL + ( + SELECT + 'workspace_invitation' AS category, + old_member.workspace_id AS "workspaceId", + NULL::text AS "docId", + old_member.user_id AS "userId", + ARRAY['Workspace.Read']::text[] AS "workspaceActions", + NULL::text[] AS "docActions" + FROM workspace_user_permissions old_member + WHERE old_member.status <> 'Accepted'::"WorkspaceMemberStatus" + AND affine_permission_workspace_invitation_state(old_member.status) IS NOT NULL + AND affine_permission_legacy_workspace_role(old_member.type) IS NOT NULL + ORDER BY md5(old_member.workspace_id || ':' || old_member.user_id) + LIMIT 40 + ) + UNION ALL + ( + SELECT + 'public_doc_anonymous' AS category, + old_doc.workspace_id AS "workspaceId", + old_doc.page_id AS "docId", + NULL::text AS "userId", + NULL::text[] AS "workspaceActions", + ARRAY['Doc.Read', 'Doc.Preview']::text[] AS "docActions" + FROM workspace_pages old_doc + WHERE old_doc.public + AND affine_permission_legacy_default_doc_role(old_doc."defaultRole") IS NOT NULL + ORDER BY md5(old_doc.workspace_id || ':' || old_doc.page_id) + LIMIT 40 + ) + UNION ALL + ( + SELECT + 'workspace_url_preview_private_doc' AS category, + old_doc.workspace_id AS "workspaceId", + old_doc.page_id AS "docId", + NULL::text AS "userId", + NULL::text[] AS "workspaceActions", + ARRAY['Doc.Preview', 'Doc.Read']::text[] AS "docActions" + FROM workspace_pages old_doc + INNER JOIN workspaces old_workspace + ON old_workspace.id = old_doc.workspace_id + WHERE old_workspace.enable_sharing + AND old_workspace.enable_url_preview + AND NOT old_doc.public + AND affine_permission_legacy_default_doc_role(old_doc."defaultRole") IS NOT NULL + ORDER BY md5(old_doc.workspace_id || ':' || old_doc.page_id) + LIMIT 40 + ) + UNION ALL + ( + SELECT + 'explicit_doc_grant' AS category, + old_grant.workspace_id AS "workspaceId", + old_grant.page_id AS "docId", + old_grant.user_id AS "userId", + NULL::text[] AS "workspaceActions", + ARRAY['Doc.Read', 'Doc.Update', 'Doc.Users.Manage', 'Doc.TransferOwner']::text[] AS "docActions" + FROM workspace_page_user_permissions old_grant + WHERE affine_permission_legacy_doc_role(old_grant.type) IS NOT NULL + ORDER BY md5(old_grant.workspace_id || ':' || old_grant.page_id || ':' || old_grant.user_id) + LIMIT 80 + ) + `; + + let mismatches = 0; + for (const sample of samples) { + const input = { + userId: sample.userId ?? undefined, + workspaceId: sample.workspaceId, + workspaceActions: sample.workspaceActions as + | PermissionWorkspaceAction[] + | undefined, + docs: + sample.docId && sample.docActions + ? [ + { + docId: sample.docId, + actions: sample.docActions as PermissionDocAction[], + }, + ] + : undefined, + }; + const [legacy, projection] = await Promise.all([ + loader.load(input).then(input => permission.evaluate(input)), + loader + .loadFromNewTables(input) + .then(input => permission.evaluate(input)), + ]); + if ( + JSON.stringify(legacy.workspace) !== + JSON.stringify(projection.workspace) || + JSON.stringify(legacy.docs) !== JSON.stringify(projection.docs) + ) { + mismatches += 1; + } + } + + return mismatches; + } +} diff --git a/packages/backend/server/src/core/permission/service.ts b/packages/backend/server/src/core/permission/service.ts new file mode 100644 index 0000000000..04d5de3ffa --- /dev/null +++ b/packages/backend/server/src/core/permission/service.ts @@ -0,0 +1,317 @@ +import { Inject, Injectable, Optional } from '@nestjs/common'; +import { Prisma } from '@prisma/client'; + +import { + Config, + DocActionDenied, + InternalServerError, + metrics, + SpaceAccessDenied, +} from '../../base'; +import { + evaluatePermissionV1, + type PermissionEvaluationInputV1, + type PermissionEvaluationOutputV1, +} from '../../native'; +import { PermissionReadModel } from './config'; +import { docLegacyBoundary, workspaceLegacyBoundary } from './context'; +import { + PermissionContextLoader, + type PermissionDocAction, + type PermissionWorkspaceAction, +} from './context-loader'; +import { WorkspacePolicyService } from './policy'; +import { PermissionSqlPredicateBuilder } from './sql-predicate'; +import type { DocAction } from './types'; + +const RUNTIME_RESTRICTED_WORKSPACE_ACTIONS = new Set( + [ + 'Workspace.Sync', + 'Workspace.CreateDoc', + 'Workspace.Delete', + 'Workspace.TransferOwner', + 'Workspace.Users.Manage', + 'Workspace.Administrators.Manage', + 'Workspace.Settings.Update', + 'Workspace.Properties.Create', + 'Workspace.Properties.Update', + 'Workspace.Properties.Delete', + 'Workspace.Blobs.Write', + 'Workspace.Payment.Manage', + ] +); + +const RUNTIME_RESTRICTED_DOC_ACTIONS = new Set([ + 'Doc.Duplicate', + 'Doc.Trash', + 'Doc.Restore', + 'Doc.Delete', + 'Doc.Update', + 'Doc.Publish', + 'Doc.TransferOwner', + 'Doc.Properties.Update', + 'Doc.Users.Manage', + 'Doc.Comments.Create', + 'Doc.Comments.Update', + 'Doc.Comments.Delete', + 'Doc.Comments.Resolve', +]); + +@Injectable() +export class PermissionService { + constructor( + private readonly loader: PermissionContextLoader, + @Optional() + @Inject(PermissionSqlPredicateBuilder) + private readonly sqlPredicate = new PermissionSqlPredicateBuilder(), + @Optional() + private readonly workspacePolicy?: WorkspacePolicyService, + @Optional() + private readonly config?: Config + ) {} + + readModel() { + return this.config?.permission.readModel ?? PermissionReadModel.Projection; + } + + docReadableSqlPredicate(input: { + userId: string; + workspaceId: string; + action: DocAction; + docIdColumn?: Prisma.Sql; + }) { + if (this.readModel() === PermissionReadModel.Projection) { + return this.sqlPredicate.docReadableByNewTablesSql(input); + } + + return this.sqlPredicate.docReadableByLegacyTablesSql(input); + } + + fallbackDocReadableSqlPredicate(input: { + userId: string; + workspaceId: string; + action: DocAction; + docIdColumn?: Prisma.Sql; + }) { + if ( + this.readModel() === PermissionReadModel.Projection && + (this.config?.permission.fallbackLegacyLoader ?? false) + ) { + return this.sqlPredicate.docReadableByLegacyTablesSql(input); + } + + return null; + } + + evaluate(input: PermissionEvaluationInputV1) { + try { + return evaluatePermissionV1(input); + } catch (error) { + throw new InternalServerError( + error instanceof Error ? error.message : undefined + ); + } + } + + async workspacePermissions(input: { + userId?: string; + workspaceId: string; + actions: PermissionWorkspaceAction[]; + allowLocal?: boolean; + }) { + const output = await this.evaluateLoaded({ + userId: input.userId, + workspaceId: input.workspaceId, + workspaceActions: input.actions, + allowLocal: input.allowLocal, + }); + return { + ...workspaceLegacyBoundary(output.workspace), + decisions: output.workspace.decisions, + }; + } + + async canWorkspace(input: { + userId?: string; + workspaceId: string; + action: PermissionWorkspaceAction; + allowLocal?: boolean; + }) { + const output = await this.workspacePermissions({ + ...input, + actions: [input.action], + }); + return output.decisions[0]?.allowed ?? false; + } + + async assertWorkspace(input: { + userId?: string; + workspaceId: string; + action: PermissionWorkspaceAction; + allowLocal?: boolean; + }) { + if (!(await this.canWorkspace(input))) { + throw new SpaceAccessDenied({ spaceId: input.workspaceId }); + } + } + + async docPermissions(input: { + userId?: string; + workspaceId: string; + docId: string; + actions: PermissionDocAction[]; + allowLocal?: boolean; + }) { + const output = await this.evaluateLoaded({ + userId: input.userId, + workspaceId: input.workspaceId, + docs: [{ docId: input.docId, actions: input.actions }], + allowLocal: input.allowLocal, + }); + const doc = output.docs[0]; + return { + ...docLegacyBoundary(doc), + decisions: doc.decisions, + }; + } + + async canDoc(input: { + userId?: string; + workspaceId: string; + docId: string; + action: PermissionDocAction; + allowLocal?: boolean; + }) { + const output = await this.docPermissions({ + ...input, + actions: [input.action], + }); + return output.decisions[0]?.allowed ?? false; + } + + async assertDoc(input: { + userId?: string; + workspaceId: string; + docId: string; + action: PermissionDocAction; + allowLocal?: boolean; + }) { + if (!(await this.canDoc(input))) { + throw new DocActionDenied({ + action: input.action, + docId: input.docId, + spaceId: input.workspaceId, + }); + } + } + + async filterReadableDocs(input: { + userId?: string; + workspaceId: string; + docs: T[]; + allowLocal?: boolean; + }) { + const decisions = await this.batchDocPermissions({ + ...input, + docs: input.docs.map(doc => ({ + docId: doc.docId, + actions: ['Doc.Read'], + })), + }); + const readableDocIds = new Set( + decisions.filter(doc => doc.decisions[0]?.allowed).map(doc => doc.docId) + ); + return input.docs.filter(doc => readableDocIds.has(doc.docId)); + } + + async batchDocPermissions(input: { + userId?: string; + workspaceId: string; + docs: Array<{ docId: string; actions: PermissionDocAction[] }>; + allowLocal?: boolean; + }) { + const output = await this.evaluateLoaded(input); + return output.docs.map(doc => ({ + docId: doc.docId, + ...docLegacyBoundary(doc), + decisions: doc.decisions, + })); + } + + async canPreviewWorkspace(input: { + userId?: string; + workspaceId: string; + allowLocal?: boolean; + }) { + return await this.canWorkspace({ + ...input, + action: 'Workspace.Preview', + }); + } + + async canPreviewDoc(input: { + userId?: string; + workspaceId: string; + docId: string; + allowLocal?: boolean; + }) { + return await this.canDoc({ + ...input, + action: 'Doc.Preview', + }); + } + + private async evaluateLoaded( + input: Parameters[0] + ) { + if (this.readModel() === PermissionReadModel.Projection) { + try { + if ( + this.needsFreshRuntimeState(input) && + (await this.loader.workspaceExists(input.workspaceId)) + ) { + await this.workspacePolicy?.getWorkspaceState(input.workspaceId); + this.loader.invalidateWorkspaceQuotaRuntime(input.workspaceId); + } + return this.evaluate(await this.loader.loadFromNewTables(input)); + } catch (error) { + if ( + input.allowLocal && + error instanceof Error && + error.message === 'Workspace owner not found' + ) { + const loaded = await this.loader.loadFromNewTables(input); + if (loaded.workspace?.local) { + return this.evaluate(loaded); + } + } + if (!(this.config?.permission.fallbackLegacyLoader ?? false)) { + throw error; + } + metrics.permission + .counter('projection_loader_fallbacks', { + description: 'Permission projection loader fallback count', + }) + .add(1); + } + } + + return this.evaluate(await this.loader.load(input)); + } + + private needsFreshRuntimeState( + input: Parameters[0] + ) { + return ( + input.workspaceActions?.some(action => + RUNTIME_RESTRICTED_WORKSPACE_ACTIONS.has(action) + ) || + input.docs?.some(doc => + doc.actions.some(action => RUNTIME_RESTRICTED_DOC_ACTIONS.has(action)) + ) || + false + ); + } +} + +export type PermissionServiceEvaluationOutput = PermissionEvaluationOutputV1; diff --git a/packages/backend/server/src/core/permission/sql-predicate.ts b/packages/backend/server/src/core/permission/sql-predicate.ts new file mode 100644 index 0000000000..166b3066ba --- /dev/null +++ b/packages/backend/server/src/core/permission/sql-predicate.ts @@ -0,0 +1,306 @@ +import { Injectable } from '@nestjs/common'; +import { Prisma } from '@prisma/client'; + +import { permissionActionRoleMatrixV1 } from '../../native'; +import { type DocAction, DocRole, WorkspaceRole } from './types'; + +export type PermissionSqlPredicate = { + sql: string; + params: unknown[]; +}; + +type RawDocIdColumn = 'doc_id' | 'docs.id'; + +@Injectable() +export class PermissionSqlPredicateBuilder { + private readonly matrix = permissionActionRoleMatrixV1() as { + doc?: { roles?: Record }; + workspace?: { roles?: Record }; + }; + + private readonly legacyDocRoleValues = new Map([ + ['external', DocRole.External], + ['reader', DocRole.Reader], + ['commenter', DocRole.Commenter], + ['editor', DocRole.Editor], + ['manager', DocRole.Manager], + ['owner', DocRole.Owner], + ]); + + private readonly legacyWorkspaceRoleValues = new Map([ + ['member', WorkspaceRole.Collaborator], + ['admin', WorkspaceRole.Admin], + ['owner', WorkspaceRole.Owner], + ]); + + private docRolesForAction(action: DocAction) { + return Object.entries(this.matrix.doc?.roles ?? {}) + .filter(([, actions]) => actions.includes(action)) + .map(([role]) => role) + .filter(role => role !== 'none'); + } + + private inheritedWorkspaceRolesForDocAction(action: DocAction) { + const docRoles = new Set(this.docRolesForAction(action)); + return [ + docRoles.has('owner') ? 'owner' : null, + docRoles.has('manager') ? 'admin' : null, + ].filter((role): role is string => role !== null); + } + + private nonMemberDocGrantRolesForAction(action: DocAction) { + const roles = new Set(this.docRolesForAction(action)); + roles.delete('external'); + roles.delete('manager'); + roles.delete('owner'); + if (roles.has('editor')) { + roles.add('manager'); + roles.add('owner'); + } + return [...roles]; + } + + private legacyNonMemberDocGrantRolesForAction(action: DocAction) { + return this.nonMemberDocGrantRolesForAction(action) + .map(role => this.legacyDocRoleValues.get(role)) + .filter(role => role !== undefined); + } + + private rawDocIdColumn(column: RawDocIdColumn = 'doc_id') { + switch (column) { + case 'doc_id': + case 'docs.id': + return column; + default: + throw new Error(`Unsupported doc id column: ${column}`); + } + } + + docReadableByLegacyTables(input: { + workspaceId: string; + userId: string; + action: DocAction; + docIdColumn?: RawDocIdColumn; + }): PermissionSqlPredicate { + const roles = this.docRolesForAction(input.action) + .map(role => this.legacyDocRoleValues.get(role)) + .filter(role => role !== undefined); + const grantRoles = roles.filter(role => role !== DocRole.External); + const nonMemberGrantRoles = this.legacyNonMemberDocGrantRolesForAction( + input.action + ); + const legacyActiveMemberRoles = [ + WorkspaceRole.Collaborator, + WorkspaceRole.Admin, + WorkspaceRole.Owner, + ]; + const inheritedWorkspaceRoles = this.inheritedWorkspaceRolesForDocAction( + input.action + ) + .map(role => this.legacyWorkspaceRoleValues.get(role)) + .filter(role => role !== undefined); + const docIdColumn = this.rawDocIdColumn(input.docIdColumn); + + return { + sql: [ + `EXISTS (SELECT 1 FROM workspaces w`, + `LEFT JOIN workspace_pages wp ON wp.workspace_id = w.id`, + `AND wp.page_id = ${docIdColumn}`, + `LEFT JOIN workspace_user_permissions wup ON wup.workspace_id = w.id`, + `AND wup.user_id = ? AND wup.status = 'Accepted'`, + `LEFT JOIN workspace_page_user_permissions p ON p.workspace_id = w.id`, + `AND p.user_id = ? AND p.page_id = ${docIdColumn}`, + `WHERE w.id = ? AND (`, + `(wup.type = ANY(?::smallint[]) AND p.type = ANY(?::smallint[]))`, + `OR ((wup.id IS NULL OR wup.type <> ALL(?::smallint[])) AND w.enable_sharing AND p.type = ANY(?::smallint[]))`, + `OR wup.type = ANY(?::smallint[])`, + `OR (wup.type = ANY(?::smallint[]) AND (p.user_id IS NULL OR p.type IN (?, ?))`, + `AND COALESCE(wp."defaultRole", 30) = ANY(?::smallint[]))`, + `OR (w.enable_sharing AND wp.public AND ? = ANY(?::smallint[]))`, + `))`, + ].join(' '), + params: [ + input.userId, + input.userId, + input.workspaceId, + legacyActiveMemberRoles, + grantRoles, + legacyActiveMemberRoles, + nonMemberGrantRoles, + inheritedWorkspaceRoles, + legacyActiveMemberRoles, + DocRole.None, + DocRole.External, + grantRoles, + DocRole.External, + roles, + ], + }; + } + + docReadableByLegacyTablesSql(input: { + workspaceId: string; + userId: string; + action: DocAction; + docIdColumn?: Prisma.Sql; + }): Prisma.Sql { + const docRoles = this.docRolesForAction(input.action); + const legacyDocRoles = docRoles + .map(role => this.legacyDocRoleValues.get(role)) + .filter(role => role !== undefined); + const legacyGrantRoles = legacyDocRoles.filter( + role => role !== DocRole.External + ); + const legacyNonMemberGrantRoles = + this.legacyNonMemberDocGrantRolesForAction(input.action); + const inheritedWorkspaceRoles = this.inheritedWorkspaceRolesForDocAction( + input.action + ) + .map(role => this.legacyWorkspaceRoleValues.get(role)) + .filter(role => role !== undefined); + const legacyActiveMemberRoles = [ + WorkspaceRole.Collaborator, + WorkspaceRole.Admin, + WorkspaceRole.Owner, + ]; + const docIdColumn = input.docIdColumn ?? Prisma.raw('doc_id'); + + return Prisma.sql` + EXISTS ( + SELECT 1 + FROM workspaces w + LEFT JOIN workspace_pages wp + ON wp.workspace_id = w.id + AND wp.page_id = ${docIdColumn} + LEFT JOIN workspace_user_permissions wup + ON wup.workspace_id = w.id + AND wup.user_id = ${input.userId} + AND wup.status = 'Accepted'::"WorkspaceMemberStatus" + LEFT JOIN workspace_page_user_permissions p + ON p.workspace_id = w.id + AND p.page_id = ${docIdColumn} + AND p.user_id = ${input.userId} + WHERE w.id = ${input.workspaceId} + AND ( + ( + wup.type = ANY(${Prisma.sql`${legacyActiveMemberRoles}::smallint[]`}) + AND p.type = ANY(${Prisma.sql`${legacyGrantRoles}::smallint[]`}) + ) + OR ( + ( + wup.id IS NULL + OR wup.type <> ALL(${Prisma.sql`${legacyActiveMemberRoles}::smallint[]`}) + ) + AND w.enable_sharing + AND p.type = ANY(${Prisma.sql`${legacyNonMemberGrantRoles}::smallint[]`}) + ) + OR wup.type = ANY(${Prisma.sql`${inheritedWorkspaceRoles}::smallint[]`}) + OR ( + wup.type = ANY(${Prisma.sql`${legacyActiveMemberRoles}::smallint[]`}) + AND ( + p.user_id IS NULL + OR p.type IN (${DocRole.None}, ${DocRole.External}) + ) + AND COALESCE(wp."defaultRole", 30) = ANY(${Prisma.sql`${legacyGrantRoles}::smallint[]`}) + ) + OR ( + w.enable_sharing + AND wp.public + AND 0 = ANY(${Prisma.sql`${legacyDocRoles}::smallint[]`}) + ) + ) + ) + `; + } + + docReadableByNewTables(input: { + workspaceId: string; + userId?: string; + action: DocAction; + docIdColumn?: RawDocIdColumn; + }): PermissionSqlPredicate { + const docRoles = this.docRolesForAction(input.action); + const inheritedWorkspaceRoles = this.inheritedWorkspaceRolesForDocAction( + input.action + ); + const grantRoles = docRoles.filter(role => role !== 'external'); + const nonMemberGrantRoles = this.nonMemberDocGrantRolesForAction( + input.action + ); + const docIdColumn = this.rawDocIdColumn(input.docIdColumn); + + return { + sql: [ + `EXISTS (SELECT 1 FROM workspace_access_policies wap`, + `LEFT JOIN doc_access_policies dap ON dap.workspace_id = wap.workspace_id`, + `AND dap.doc_id = ${docIdColumn}`, + `LEFT JOIN workspace_members wm ON wm.workspace_id = wap.workspace_id`, + `AND wm.user_id = ? AND wm.state = 'active'`, + `LEFT JOIN doc_grants dg ON dg.workspace_id = wap.workspace_id`, + `AND dg.doc_id = ${docIdColumn} AND dg.principal_type = 'user' AND dg.principal_id = ?`, + `WHERE wap.workspace_id = ?`, + `AND (`, + `(wm.id IS NOT NULL AND dg.role = ANY(?::text[]))`, + `OR (wm.id IS NULL AND wap.sharing_enabled AND dg.role = ANY(?::text[]))`, + `OR wm.role = ANY(?::text[])`, + `OR (wm.id IS NOT NULL AND dg.principal_id IS NULL AND COALESCE(dap.member_default_role, wap.member_default_doc_role) = ANY(?::text[]))`, + `OR (wap.sharing_enabled AND dap.visibility = 'public' AND dap.public_role = ANY(?::text[]))`, + `))`, + ].join(' '), + params: [ + input.userId, + input.userId, + input.workspaceId, + grantRoles, + nonMemberGrantRoles, + inheritedWorkspaceRoles, + grantRoles, + docRoles, + ], + }; + } + + docReadableByNewTablesSql(input: { + workspaceId: string; + userId?: string; + action: DocAction; + docIdColumn?: Prisma.Sql; + }): Prisma.Sql { + const docRoles = this.docRolesForAction(input.action); + const grantRoles = docRoles.filter(role => role !== 'external'); + const nonMemberGrantRoles = this.nonMemberDocGrantRolesForAction( + input.action + ); + const inheritedWorkspaceRoles = this.inheritedWorkspaceRolesForDocAction( + input.action + ); + const docIdColumn = input.docIdColumn ?? Prisma.raw('doc_id'); + + return Prisma.sql` + EXISTS ( + SELECT 1 + FROM workspace_access_policies wap + LEFT JOIN doc_access_policies dap + ON dap.workspace_id = wap.workspace_id + AND dap.doc_id = ${docIdColumn} + LEFT JOIN workspace_members wm + ON wm.workspace_id = wap.workspace_id + AND wm.user_id = ${input.userId} + AND wm.state = 'active' + LEFT JOIN doc_grants dg + ON dg.workspace_id = wap.workspace_id + AND dg.doc_id = ${docIdColumn} + AND dg.principal_type = 'user' + AND dg.principal_id = ${input.userId} + WHERE wap.workspace_id = ${input.workspaceId} + AND ( + (wm.id IS NOT NULL AND dg.role = ANY(${Prisma.sql`${grantRoles}::text[]`})) + OR (wm.id IS NULL AND wap.sharing_enabled AND dg.role = ANY(${Prisma.sql`${nonMemberGrantRoles}::text[]`})) + OR wm.role = ANY(${Prisma.sql`${inheritedWorkspaceRoles}::text[]`}) + OR (wm.id IS NOT NULL AND dg.principal_id IS NULL AND COALESCE(dap.member_default_role, wap.member_default_doc_role) = ANY(${Prisma.sql`${grantRoles}::text[]`})) + OR (wap.sharing_enabled AND dap.visibility = 'public' AND dap.public_role = ANY(${Prisma.sql`${docRoles}::text[]`})) + ) + ) + `; + } +} diff --git a/packages/backend/server/src/core/permission/types.ts b/packages/backend/server/src/core/permission/types.ts index 18608b03a4..1065ef7c46 100644 --- a/packages/backend/server/src/core/permission/types.ts +++ b/packages/backend/server/src/core/permission/types.ts @@ -145,6 +145,7 @@ export const RoleActionsMap = { Action.Doc.Delete, Action.Doc.Properties.Update, Action.Doc.Update, + Action.Doc.Comments.Update, Action.Doc.Comments.Resolve, Action.Doc.Comments.Delete, ]; diff --git a/packages/backend/server/src/core/permission/workspace.ts b/packages/backend/server/src/core/permission/workspace.ts deleted file mode 100644 index b89c59293c..0000000000 --- a/packages/backend/server/src/core/permission/workspace.ts +++ /dev/null @@ -1,210 +0,0 @@ -import { Injectable } from '@nestjs/common'; - -import { SpaceAccessDenied } from '../../base'; -import { DocRole, Models } from '../../models'; -import { AccessController } from './controller'; -import { WorkspacePolicyService } from './policy'; -import type { Resource } from './resource'; -import { - fixupDocRole, - mapDocRoleToPermissions, - mapWorkspaceRoleToPermissions, - WorkspaceAction, - workspaceActionRequiredRole, - WorkspaceRole, -} from './types'; - -@Injectable() -export class WorkspaceAccessController extends AccessController<'ws'> { - protected readonly type = 'ws'; - - constructor( - private readonly models: Models, - private readonly policy: WorkspacePolicyService - ) { - super(); - } - - async role(resource: Resource<'ws'>) { - const role = await this.getRole(resource); - - return { - role, - permissions: await this.policy.applyWorkspacePermissions( - resource.workspaceId, - mapWorkspaceRoleToPermissions(role) - ), - }; - } - - async can(resource: Resource<'ws'>, action: WorkspaceAction) { - const { permissions, role } = await this.role(resource); - const allow = permissions[action] || false; - - if (!allow) { - this.logger.debug('Workspace access check failed', { - action, - resource, - role, - requiredRole: workspaceActionRequiredRole(action), - }); - } - - return allow; - } - - async assert(resource: Resource<'ws'>, action: WorkspaceAction) { - const allow = await this.can(resource, action); - - if (!allow) { - throw new SpaceAccessDenied({ spaceId: resource.workspaceId }); - } - } - - async getRole(payload: Resource<'ws'>) { - const userRole = await this.models.workspaceUser.getActive( - payload.workspaceId, - payload.userId - ); - - let role = userRole?.type as WorkspaceRole | null; - - if (!role) { - role = await this.defaultWorkspaceRole(payload); - } - - return role; - } - - async docRoles(payload: Resource<'ws'>, docIds: string[]) { - const docRoles = await this.getDocRoles(payload, docIds); - return docRoles.map(role => ({ - role, - permissions: mapDocRoleToPermissions(role), - })); - } - - async getDocRoles(payload: Resource<'ws'>, docIds: string[]) { - const docRoles: (DocRole | null)[] = []; - - if (docIds.length === 0) { - return docRoles; - } - - const workspaceRole = await this.getRole(payload); - const sharingAllowed = await this.policy.isSharingEnabled( - payload.workspaceId - ); - if ( - !sharingAllowed && - (workspaceRole === null || workspaceRole === WorkspaceRole.External) - ) { - return docIds.map(() => null); - } - - const userRoles = await this.models.docUser.findMany( - payload.workspaceId, - docIds, - payload.userId - ); - const userRolesMap = new Map(userRoles.map(role => [role.docId, role])); - - const noUserRoleDocIds = docIds.filter(docId => { - const userRole = userRolesMap.get(docId); - return (userRole?.type ?? null) === null; - }); - const defaultDocRoles = - noUserRoleDocIds.length > 0 - ? await this.getDocDefaultRoles( - payload, - noUserRoleDocIds, - workspaceRole - ) - : []; - const defaultDocRolesMap = new Map( - defaultDocRoles.map((role, index) => [noUserRoleDocIds[index], role]) - ); - - for (const docId of docIds) { - const userRole = userRolesMap.get(docId); - - let docRole: DocRole | null = userRole?.type ?? null; - - // fallback logic - if (docRole === null) { - docRole = defaultDocRolesMap.get(docId) ?? null; - } - - // we need to fixup doc role to make sure it's not miss set - // for example: workspace owner will have doc owner role - // workspace external will not have role higher than editor - const role = fixupDocRole(workspaceRole, docRole); - - // never return [None] - docRoles.push(role === DocRole.None ? null : role); - } - - return docRoles; - } - - private async getDocDefaultRoles( - payload: Resource<'ws'>, - docIds: string[], - workspaceRole: WorkspaceRole | null - ) { - const fallbackDocRoles: (DocRole | null)[] = []; - - if (docIds.length === 0) { - return fallbackDocRoles; - } - - const defaultDocRoles = await this.models.doc.findDefaultRoles( - payload.workspaceId, - docIds - ); - - for (const defaultDocRole of defaultDocRoles) { - let docRole: DocRole | null; - // if user is in workspace but doc role is not set, fallback to default doc role - if (workspaceRole !== null && workspaceRole !== WorkspaceRole.External) { - docRole = - defaultDocRole.external !== null - ? // edgecase: when doc role set to [None] for workspace member, but doc is public, we should fallback to external role - Math.max(defaultDocRole.workspace, defaultDocRole.external) - : defaultDocRole.workspace; - } else { - // else fallback to external doc role - docRole = defaultDocRole.external; - } - - fallbackDocRoles.push(docRole); - } - - return fallbackDocRoles; - } - - private async defaultWorkspaceRole(payload: Resource<'ws'>) { - const ws = await this.models.workspace.get(payload.workspaceId); - - // NOTE(@forehalo): - // we allow user to use online service with local workspace - // so we always return owner role for local workspace - // copilot session for local workspace is an example - if (!ws) { - if (payload.allowLocal) { - return WorkspaceRole.Owner; - } - - return null; - } - - if (ws.public) { - const sharingAllowed = await this.policy.canReadWorkspaceByPublicFlag( - ws.id - ); - return sharingAllowed ? WorkspaceRole.External : null; - } - - return null; - } -} diff --git a/packages/backend/server/src/core/quota/__tests__/state.spec.ts b/packages/backend/server/src/core/quota/__tests__/state.spec.ts new file mode 100644 index 0000000000..277a60c958 --- /dev/null +++ b/packages/backend/server/src/core/quota/__tests__/state.spec.ts @@ -0,0 +1,495 @@ +import { randomUUID } from 'node:crypto'; + +import { PrismaClient } from '@prisma/client'; +import ava, { ExecutionContext, TestFn } from 'ava'; + +import { + createTestingModule, + type TestingModule, +} from '../../../__tests__/utils'; +import { EventBus } from '../../../base'; +import { + Models, + Workspace, + WorkspaceMemberStatus, + WorkspaceRole, +} from '../../../models'; +import { + SubscriptionPlan, + SubscriptionRecurring, +} from '../../../plugins/payment/types'; +import { EntitlementModule, EntitlementService } from '../../entitlement'; +import { QuotaService } from '../service'; +import { QuotaServiceModule } from '../service.module'; +import { QuotaStateService } from '../state'; + +interface Context { + module: TestingModule; + db: PrismaClient; + models: Models; + entitlement: EntitlementService; + quota: QuotaService; + state: QuotaStateService; +} + +const test = ava.serial as TestFn; +const ONE_GB = 1024 * 1024 * 1024; +const ONE_DAY_SECONDS = 24 * 60 * 60; +type CaseState = { + userId?: string; + workspaceId?: string; +}; + +test.before(async t => { + const module = await createTestingModule({ + imports: [EntitlementModule, QuotaServiceModule], + }); + t.context.module = module; + t.context.db = module.get(PrismaClient); + t.context.models = module.get(Models); + t.context.entitlement = module.get(EntitlementService); + t.context.quota = module.get(QuotaService); + t.context.state = module.get(QuotaStateService); +}); + +test('quota service ignores dirty legacy commercial features', async t => { + const { owner, workspace } = await createWorkspace(t); + await t.context.models.userFeature.add( + owner.id, + 'pro_plan_v1', + 'dirty legacy feature' + ); + await t.context.models.userFeature.add( + owner.id, + 'unlimited_copilot', + 'dirty legacy feature' + ); + await t.context.models.workspaceFeature.add( + workspace.id, + 'team_plan_v1', + 'dirty legacy feature', + { + memberLimit: 100, + } + ); + + const userQuota = await t.context.quota.getUserQuota(owner.id); + const workspaceSeats = await t.context.quota.getWorkspaceSeatQuota( + workspace.id + ); + + t.is(userQuota.name, 'Free'); + t.is(userQuota.copilotActionLimit, 10); + t.is(workspaceSeats.memberLimit, 3); +}); + +test('workspace quota state ignores dirty legacy readonly feature', async t => { + const { workspace } = await createWorkspace(t); + await t.context.models.workspaceFeature.add( + workspace.id, + 'quota_exceeded_readonly_workspace_v1', + 'dirty legacy feature' + ); + + const state = await t.context.state.reconcileWorkspaceQuotaState( + workspace.id + ); + + t.false(state.readonly); + t.deepEqual(state.readonlyReasons, []); +}); + +test('workspace quota state ignores dirty legacy permission rows', async t => { + const { workspace } = await createWorkspace(t); + const member = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + await t.context.models.workspaceUser.set( + workspace.id, + member.id, + WorkspaceRole.Collaborator, + { + status: WorkspaceMemberStatus.Accepted, + } + ); + await t.context.db.$transaction(async tx => { + await tx.$executeRaw` + SELECT set_config('affine.permission_projection.enabled', 'off', true) + `; + await tx.workspaceMember.deleteMany({ + where: { + workspaceId: workspace.id, + userId: member.id, + }, + }); + }); + + const state = await t.context.state.reconcileWorkspaceQuotaState( + workspace.id + ); + + t.is(state.memberCount, 1); +}); + +test('quota service exposes history period in seconds', async t => { + const { owner, workspace } = await createWorkspace(t); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: owner.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Monthly, + status: 'active', + }); + + const userState = await t.context.state.reconcileUserQuotaState(owner.id); + const workspaceState = await t.context.state.reconcileWorkspaceQuotaState( + workspace.id + ); + const workspaceQuota = await t.context.quota.getWorkspaceQuota(workspace.id); + + t.is(userState.historyPeriodSeconds, 30 * ONE_DAY_SECONDS); + t.is(workspaceState.historyPeriodSeconds, 30 * ONE_DAY_SECONDS); + t.is(workspaceQuota.historyPeriod, 30 * ONE_DAY_SECONDS); + t.is( + t.context.quota.formatWorkspaceQuota({ + ...workspaceQuota, + usedStorageQuota: 0, + memberCount: 1, + overcapacityMemberCount: 0, + usedSize: 0, + }).historyPeriod, + '30 days' + ); +}); + +test('quota state reconcile does not publish unchanged snapshots', async t => { + const user = await t.context.models.user.create({ + email: 'quota-event-owner@affine.pro', + }); + await t.context.db.effectiveUserQuotaState.deleteMany({ + where: { userId: user.id }, + }); + const event = t.context.module.get(EventBus); + let changes = 0; + event.on('user.quota_state.changed', ({ userId }) => { + if (userId === user.id) { + changes += 1; + } + }); + + await t.context.state.reconcileUserQuotaState(user.id); + await t.context.state.reconcileUserQuotaState(user.id); + + t.is(changes, 1); +}); + +test('workspace quota state requires owner from new permission table', async t => { + const { owner, workspace } = await createWorkspace(t); + await t.context.db.$transaction(async tx => { + await tx.$executeRaw` + SELECT set_config('affine.permission_projection.enabled', 'off', true) + `; + await tx.workspaceMember.deleteMany({ + where: { + workspaceId: workspace.id, + userId: owner.id, + }, + }); + }); + + await t.throwsAsync( + t.context.state.reconcileWorkspaceQuotaState(workspace.id), + { message: 'Workspace owner not found' } + ); +}); + +test('user quota state aggregates owned storage from new permission table only', async t => { + const { owner, workspace } = await createWorkspace(t); + await addBlob(t, workspace, 'blob', ONE_GB); + + const first = await t.context.state.reconcileUserQuotaState(owner.id); + await t.context.db.$transaction(async tx => { + await tx.$executeRaw` + SELECT set_config('affine.permission_projection.enabled', 'off', true) + `; + await tx.workspaceMember.deleteMany({ + where: { + workspaceId: workspace.id, + userId: owner.id, + }, + }); + }); + const second = await t.context.state.reconcileUserQuotaState(owner.id); + + t.is(first.usedStorageQuota, BigInt(ONE_GB)); + t.is(second.usedStorageQuota, 0n); +}); + +test('user quota state keeps ai capability alongside pro entitlement', async t => { + const { owner } = await createWorkspace(t); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: owner.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Yearly, + status: 'active', + }); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: owner.id, + plan: SubscriptionPlan.AI, + recurring: SubscriptionRecurring.Monthly, + status: 'active', + }); + + const state = await t.context.state.reconcileUserQuotaState(owner.id); + const quota = await t.context.quota.getUserQuota(owner.id); + + t.is(state.plan, 'pro'); + t.deepEqual(state.flags, { unlimitedCopilot: true }); + t.is(quota.copilotActionLimit, undefined); +}); + +test('ai entitlement is a capability overlay on free quota', async t => { + const { owner } = await createWorkspace(t); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: owner.id, + plan: SubscriptionPlan.AI, + recurring: SubscriptionRecurring.Monthly, + status: 'active', + }); + + const state = await t.context.state.reconcileUserQuotaState(owner.id); + const quota = await t.context.quota.getUserQuota(owner.id); + + t.is(state.plan, 'free'); + t.deepEqual(state.flags, { unlimitedCopilot: true }); + t.is(quota.name, 'Free'); + t.is(quota.copilotActionLimit, undefined); +}); + +test('workspace team status ignores dirty legacy feature', async t => { + const { workspace } = await createWorkspace(t); + await t.context.models.workspaceFeature.add( + workspace.id, + 'team_plan_v1', + 'dirty legacy feature', + { + memberLimit: 100, + } + ); + + t.false(await t.context.models.workspace.isTeamWorkspace(workspace.id)); + + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: workspace.id, + plan: SubscriptionPlan.Team, + recurring: SubscriptionRecurring.Yearly, + status: 'active', + quantity: 5, + }); + + t.true(await t.context.models.workspace.isTeamWorkspace(workspace.id)); +}); + +test('selfhosted builtin free has cloud pro quota rights', async t => { + const previousDeploymentType = globalThis.env.DEPLOYMENT_TYPE; + // @ts-expect-error test mutates env singleton for deployment-specific quota semantics + globalThis.env.DEPLOYMENT_TYPE = 'selfhosted'; + try { + const { owner, workspace } = await createWorkspace(t); + + const userState = await t.context.state.reconcileUserQuotaState(owner.id); + const userQuota = await t.context.quota.getUserQuota(owner.id); + const workspaceState = await t.context.state.reconcileWorkspaceQuotaState( + workspace.id + ); + const workspaceQuota = await t.context.quota.getWorkspaceQuota( + workspace.id + ); + + t.is(userState.plan, 'selfhost_free'); + t.is(userState.storageQuota, BigInt(100 * ONE_GB)); + t.is(userQuota.name, 'Pro'); + t.is(userQuota.memberLimit, 10); + t.is(workspaceState.plan, 'selfhost_free'); + t.is(workspaceQuota.name, 'Pro'); + t.is(workspaceQuota.memberLimit, 10); + } finally { + // @ts-expect-error restore mutable test env singleton + globalThis.env.DEPLOYMENT_TYPE = previousDeploymentType; + } +}); + +test.beforeEach(async t => { + await t.context.module.initTestingDB(); +}); + +test.after.always(async t => { + await t.context.module.close(); +}); + +test('reconciles quota states from entitlements and business tables', async t => { + const cases = [ + { + name: 'owner fallback uses user entitlement and owner storage usage', + setup: async () => { + const { owner, workspace } = await createWorkspace(t); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: owner.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Monthly, + status: 'active', + }); + await addBlob(t, workspace, 'blob', ONE_GB); + + return { userId: owner.id, workspaceId: workspace.id }; + }, + assert: async ({ userId, workspaceId }: CaseState) => { + const user = await t.context.state.reconcileUserQuotaState(userId!); + const workspace = await t.context.state.reconcileWorkspaceQuotaState( + workspaceId! + ); + + t.is(user.plan, 'pro'); + t.is(user.usedStorageQuota, BigInt(ONE_GB)); + t.true(workspace.usesOwnerQuota); + t.is(workspace.plan, 'pro'); + t.is( + (await t.context.quota.getWorkspaceQuota(workspaceId!)).name, + 'Pro' + ); + t.is(workspace.storageQuota, BigInt(100 * ONE_GB)); + t.is(workspace.usedStorageQuota, BigInt(ONE_GB)); + }, + }, + { + name: 'team entitlement owns workspace quota', + setup: async () => { + const { workspace } = await createWorkspace(t); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: workspace.id, + plan: SubscriptionPlan.Team, + recurring: SubscriptionRecurring.Yearly, + status: 'active', + quantity: 5, + }); + + return { workspaceId: workspace.id }; + }, + assert: async ({ workspaceId }: CaseState) => { + const workspace = await t.context.state.reconcileWorkspaceQuotaState( + workspaceId! + ); + + t.false(workspace.usesOwnerQuota); + t.is(workspace.seatLimit, 5); + t.is(workspace.storageQuota, BigInt(200 * ONE_GB)); + }, + }, + { + name: 'overcapacity members set readonly state', + setup: async () => { + const { workspace } = await createWorkspace(t); + await addAcceptedMembers(t, workspace.id, 4); + + return { workspaceId: workspace.id }; + }, + assert: async ({ workspaceId }: CaseState) => { + const workspace = await t.context.state.reconcileWorkspaceQuotaState( + workspaceId! + ); + + t.true(workspace.readonly); + t.deepEqual(workspace.readonlyReasons, ['member_overflow']); + t.is(workspace.overcapacityMemberCount, 2); + }, + }, + { + name: 'storage overflow sets readonly state', + setup: async () => { + const { workspace } = await createWorkspace(t); + for (let index = 0; index < 11; index++) { + await addBlob(t, workspace, `blob-${index}`, ONE_GB); + } + + return { workspaceId: workspace.id }; + }, + assert: async ({ workspaceId }: CaseState) => { + const workspace = await t.context.state.reconcileWorkspaceQuotaState( + workspaceId! + ); + + t.true(workspace.readonly); + t.deepEqual(workspace.readonlyReasons, ['storage_overflow']); + t.is(workspace.usedStorageQuota, BigInt(11 * ONE_GB)); + }, + }, + { + name: 'expired entitlement falls back to free state', + setup: async () => { + const { owner } = await createWorkspace(t); + await t.context.entitlement.upsertFromCloudSubscription({ + targetId: owner.id, + plan: SubscriptionPlan.Pro, + recurring: SubscriptionRecurring.Monthly, + status: 'canceled', + }); + + return { userId: owner.id }; + }, + assert: async ({ userId }: CaseState) => { + const user = await t.context.state.reconcileUserQuotaState(userId!); + + t.is(user.plan, 'free'); + t.is(user.sourceEntitlementId, null); + }, + }, + ]; + + for (const item of cases) { + await t.context.module.initTestingDB(); + const state = await item.setup(); + await item.assert(state); + } +}); + +async function createWorkspace(t: ExecutionContext) { + const owner = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + const workspace = await t.context.models.workspace.create(owner.id); + + return { owner, workspace }; +} + +async function addAcceptedMembers( + t: ExecutionContext, + workspaceId: string, + count: number +) { + for (let index = 0; index < count; index++) { + const member = await t.context.models.user.create({ + email: `${randomUUID()}@affine.pro`, + }); + await t.context.models.workspaceUser.set( + workspaceId, + member.id, + WorkspaceRole.Collaborator, + { + status: WorkspaceMemberStatus.Accepted, + } + ); + } +} + +async function addBlob( + t: ExecutionContext, + workspace: Workspace, + key: string, + size: number +) { + await t.context.models.blob.upsert({ + workspaceId: workspace.id, + key, + mime: 'application/octet-stream', + size, + }); +} diff --git a/packages/backend/server/src/core/quota/index.ts b/packages/backend/server/src/core/quota/index.ts index 3d9dca86ea..ca40d046b4 100644 --- a/packages/backend/server/src/core/quota/index.ts +++ b/packages/backend/server/src/core/quota/index.ts @@ -19,4 +19,5 @@ export class QuotaModule {} export { QuotaService }; export { QuotaServiceModule }; +export { QuotaStateService } from './state'; export { WorkspaceQuotaHumanReadableType, WorkspaceQuotaType } from './types'; diff --git a/packages/backend/server/src/core/quota/realtime.ts b/packages/backend/server/src/core/quota/realtime.ts new file mode 100644 index 0000000000..ff15d85180 --- /dev/null +++ b/packages/backend/server/src/core/quota/realtime.ts @@ -0,0 +1,147 @@ +import { Injectable, OnModuleInit, Optional } from '@nestjs/common'; +import { z } from 'zod'; + +import { OnEvent, SpaceAccessDenied } from '../../base'; +import { Models } from '../../models'; +import { registerRealtimeLiveQuery } from '../realtime/provider'; +import { RealtimePublisher } from '../realtime/publisher'; +import { RealtimeRegistry } from '../realtime/registry'; +import { + realtimeUserQuotaStateRoom, + realtimeWorkspaceQuotaStateRoom, +} from '../realtime/rooms'; +import { QuotaStateService } from './state'; + +type UserQuotaStateSnapshot = import('@affine/realtime').UserQuotaStateSnapshot; +type WorkspaceQuotaStateSnapshot = + import('@affine/realtime').WorkspaceQuotaStateSnapshot; + +declare module '@affine/realtime' { + interface RealtimeRequestMap { + 'user.quota-state.get': { + input: Record; + output: { state: UserQuotaStateSnapshot }; + }; + 'workspace.quota-state.get': { + input: { workspaceId: string }; + output: { state: WorkspaceQuotaStateSnapshot }; + }; + } + + interface RealtimeTopicMap { + 'user.quota-state.changed': { + input: Record; + event: { changed: true }; + }; + 'workspace.quota-state.changed': { + input: { workspaceId: string }; + event: { changed: true }; + }; + } +} + +@Injectable() +export class QuotaStateRealtimeProvider implements OnModuleInit { + constructor( + private readonly models: Models, + private readonly quotaState: QuotaStateService, + @Optional() private readonly registry?: RealtimeRegistry, + @Optional() private readonly publisher?: RealtimePublisher + ) {} + + onModuleInit() { + const { registry } = this; + if (!registry) return; + + const workspaceInput = z.object({ workspaceId: z.string() }); + + registerRealtimeLiveQuery(registry, { + request: { + name: 'user.quota-state.get', + input: z.object({}), + handle: async user => ({ + state: this.serializeState( + await this.quotaState.reconcileUserQuotaState(user.id) + ) as unknown as UserQuotaStateSnapshot, + }), + }, + topic: { + name: 'user.quota-state.changed', + input: z.object({}), + authorize: async () => {}, + room: user => { + if (!user) { + throw new Error('Authenticated user is required'); + } + return realtimeUserQuotaStateRoom(user.id); + }, + }, + }); + + registerRealtimeLiveQuery(registry, { + request: { + name: 'workspace.quota-state.get', + input: workspaceInput, + handle: async (user, payload) => { + await this.assertWorkspace(user.id, payload.workspaceId); + return { + state: this.serializeState( + await this.quotaState.reconcileWorkspaceQuotaState( + payload.workspaceId + ) + ) as unknown as WorkspaceQuotaStateSnapshot, + }; + }, + }, + topic: { + name: 'workspace.quota-state.changed', + input: workspaceInput, + authorize: async (user, payload) => { + await this.assertWorkspace(user.id, payload.workspaceId); + }, + room: (_user, payload) => + realtimeWorkspaceQuotaStateRoom(payload.workspaceId), + }, + }); + } + + @OnEvent('user.quota_state.changed', { suppressError: true }) + async onUserQuotaStateChanged({ + userId, + }: Events['user.quota_state.changed']) { + this.publisher?.publish( + 'user.quota-state.changed', + {}, + { changed: true }, + { room: realtimeUserQuotaStateRoom(userId) } + ); + } + + @OnEvent('workspace.quota_state.changed', { suppressError: true }) + async onWorkspaceQuotaStateChanged({ + workspaceId, + }: Events['workspace.quota_state.changed']) { + this.publisher?.publish( + 'workspace.quota-state.changed', + { workspaceId }, + { changed: true }, + { room: realtimeWorkspaceQuotaStateRoom(workspaceId) } + ); + } + + private async assertWorkspace(userId: string, workspaceId: string) { + const role = await this.models.workspaceUser.getActive(workspaceId, userId); + if (!role) { + throw new SpaceAccessDenied({ spaceId: workspaceId }); + } + } + + private serializeState>(state: T) { + return Object.fromEntries( + Object.entries(state).map(([key, value]) => [ + key, + typeof value === 'bigint' ? Number(value) : value, + ]) + ); + } +} diff --git a/packages/backend/server/src/core/quota/service.module.ts b/packages/backend/server/src/core/quota/service.module.ts index 411830a39f..3de31656b5 100644 --- a/packages/backend/server/src/core/quota/service.module.ts +++ b/packages/backend/server/src/core/quota/service.module.ts @@ -1,11 +1,14 @@ import { Module } from '@nestjs/common'; +import { EntitlementModule } from '../entitlement'; import { StorageModule } from '../storage'; +import { QuotaStateRealtimeProvider } from './realtime'; import { QuotaService } from './service'; +import { QuotaStateService } from './state'; @Module({ - imports: [StorageModule], - providers: [QuotaService], - exports: [QuotaService], + imports: [StorageModule, EntitlementModule], + providers: [QuotaService, QuotaStateService, QuotaStateRealtimeProvider], + exports: [QuotaService, QuotaStateService], }) export class QuotaServiceModule {} diff --git a/packages/backend/server/src/core/quota/service.ts b/packages/backend/server/src/core/quota/service.ts index eff27ddc41..974dae1dcb 100644 --- a/packages/backend/server/src/core/quota/service.ts +++ b/packages/backend/server/src/core/quota/service.ts @@ -1,13 +1,11 @@ import { Injectable, Logger } from '@nestjs/common'; -import { InternalServerError, MemberQuotaExceeded, OnEvent } from '../../base'; +import { MemberQuotaExceeded, OnEvent } from '../../base'; import { - Models, type UserQuota, WorkspaceQuota as BaseWorkspaceQuota, - WorkspaceRole, } from '../../models'; -import { WorkspaceBlobStorage } from '../storage'; +import { QuotaStateService } from './state'; import { UserQuotaHumanReadableType, UserQuotaType, @@ -29,10 +27,7 @@ export type WorkspaceQuotaWithUsage = Omit< export class QuotaService { protected logger = new Logger(QuotaService.name); - constructor( - private readonly models: Models, - private readonly storage: WorkspaceBlobStorage - ) {} + constructor(private readonly quotaState: QuotaStateService) {} @OnEvent('user.postCreated') async onUserCreated({ id }: Events['user.postCreated']) { @@ -40,121 +35,48 @@ export class QuotaService { } async getUserQuota(userId: string): Promise { - let quota = await this.models.userFeature.getQuota(userId); + const state = await this.quotaState.reconcileUserQuotaState(userId); - // not possible, but just in case, we do a little fix for user to avoid system dump - if (!quota) { - await this.setupUserBaseQuota(userId); - quota = await this.models.userFeature.getQuota(userId); - } - - const unlimitedCopilot = await this.models.userFeature.has( - userId, - 'unlimited_copilot' - ); - - if (!quota) { - throw new InternalServerError( - 'User quota not found and can not be created.' - ); - } - - return { - ...quota.configs, - copilotActionLimit: unlimitedCopilot - ? undefined - : quota.configs.copilotActionLimit, - } as UserQuotaWithUsage; + return this.userQuotaFromState(state); } async getUserQuotaWithUsage(userId: string): Promise { - const quota = await this.getUserQuota(userId); - const usedStorageQuota = await this.getUserStorageUsage(userId); + const state = await this.quotaState.reconcileUserQuotaState(userId); + const quota = this.userQuotaFromState(state); - return { ...quota, usedStorageQuota }; + return { ...quota, usedStorageQuota: Number(state.usedStorageQuota) }; } async getUserStorageUsage(userId: string) { - const workspaces = await this.models.workspaceUser.getUserActiveRoles( - userId, - { - role: WorkspaceRole.Owner, - } - ); - - const ids = workspaces.map(w => w.workspaceId); - - const workspacesWithQuota = - await this.models.workspaceFeature.batchHasQuota(ids); - - const sizes = await Promise.allSettled( - ids - .filter(w => !workspacesWithQuota.includes(w)) - .map(workspace => this.storage.totalSize(workspace)) - ); - - return sizes.reduce((total, size) => { - if (size.status === 'fulfilled') { - // ensure that size is within the safe range of gql - const totalSize = total + size.value; - if (Number.isSafeInteger(totalSize)) { - return totalSize; - } else { - this.logger.error(`Workspace size is invalid: ${size.value}`); - } - } else { - this.logger.error(`Failed to get workspace size`, size.reason); - } - return total; - }, 0); + const state = await this.quotaState.reconcileUserQuotaState(userId); + return Number(state.usedStorageQuota); } async getWorkspaceStorageUsage(workspaceId: string) { - const totalSize = await this.storage.totalSize(workspaceId); - // ensure that size is within the safe range of gql - if (Number.isSafeInteger(totalSize)) { - return totalSize; - } else { - this.logger.error(`Workspace size is invalid: ${totalSize}`); - } - - return 0; + const state = + await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); + return Number(state.usedStorageQuota); } async getWorkspaceQuota(workspaceId: string): Promise { - const quota = await this.models.workspaceFeature.getQuota(workspaceId); - - if (!quota) { - // get and convert to workspace quota from owner's quota - const owner = await this.models.workspaceUser.getOwner(workspaceId); - const ownerQuota = await this.getUserQuota(owner.id); - - return { - ...ownerQuota, - ownerQuota: owner.id, - }; - } - - return quota.configs; + const state = + await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); + return this.workspaceQuotaFromState(state); } async getWorkspaceQuotaWithUsage( workspaceId: string ): Promise { - const quota = await this.getWorkspaceQuota(workspaceId); - const usedStorageQuota = quota.ownerQuota - ? await this.getUserStorageUsage(quota.ownerQuota) - : await this.getWorkspaceStorageUsage(workspaceId); - const memberCount = - await this.models.workspaceUser.chargedCount(workspaceId); - const overcapacityMemberCount = memberCount - quota.memberLimit; + const state = + await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); + const quota = this.workspaceQuotaFromState(state); return { ...quota, - usedStorageQuota, - memberCount, - overcapacityMemberCount, - usedSize: usedStorageQuota, + usedStorageQuota: Number(state.usedStorageQuota), + memberCount: state.memberCount, + overcapacityMemberCount: state.overcapacityMemberCount, + usedSize: Number(state.usedStorageQuota), }; } @@ -175,13 +97,12 @@ export class QuotaService { } async getWorkspaceSeatQuota(workspaceId: string) { - const quota = await this.getWorkspaceQuota(workspaceId); - const memberCount = - await this.models.workspaceUser.chargedCount(workspaceId); + const state = + await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); return { - memberCount, - memberLimit: quota.memberLimit, + memberCount: state.memberCount, + memberLimit: state.seatLimit, }; } @@ -215,42 +136,27 @@ export class QuotaService { } async getUserQuotaCalculator(userId: string) { - const quota = await this.getUserQuota(userId); - const usedSize = await this.getUserStorageUsage(userId); + const quota = await this.getUserQuotaWithUsage(userId); return this.generateQuotaCalculator( quota.storageQuota, quota.blobLimit, - usedSize + quota.usedStorageQuota ); } async getWorkspaceQuotaCalculator(workspaceId: string) { - const quota = await this.getWorkspaceQuota(workspaceId); - const unlimited = await this.models.workspaceFeature.has( - workspaceId, - 'unlimited_workspace' - ); - - // quota check will be disabled for unlimited workspace - // we save a complicated db read for used size - if (unlimited) { - return this.generateQuotaCalculator(0, quota.blobLimit, 0, true); - } - - const usedSize = quota.ownerQuota - ? await this.getUserStorageUsage(quota.ownerQuota) - : await this.getWorkspaceStorageUsage(workspaceId); + const quota = await this.getWorkspaceQuotaWithUsage(workspaceId); return this.generateQuotaCalculator( quota.storageQuota, quota.blobLimit, - usedSize + quota.usedStorageQuota ); } private async setupUserBaseQuota(userId: string) { - await this.models.userFeature.add(userId, 'free_plan_v1', 'sign up'); + await this.quotaState.reconcileUserQuotaState(userId); } private generateQuotaCalculator( @@ -278,4 +184,60 @@ export class QuotaService { }; return checkExceeded; } + + private userQuotaFromState( + state: Awaited> + ): UserQuota { + const flags = state.flags as { unlimitedCopilot?: boolean }; + return { + name: this.planName(state.plan), + blobLimit: Number(state.blobLimit), + storageQuota: Number(state.storageQuota), + historyPeriod: state.historyPeriodSeconds, + memberLimit: this.userMemberLimit(state.plan), + copilotActionLimit: flags.unlimitedCopilot + ? undefined + : (state.copilotActionLimit ?? undefined), + }; + } + + private workspaceQuotaFromState( + state: Awaited< + ReturnType + > + ): WorkspaceQuota { + return { + name: this.planName(state.plan), + blobLimit: Number(state.blobLimit), + storageQuota: Number(state.storageQuota), + historyPeriod: state.historyPeriodSeconds, + memberLimit: state.seatLimit, + ownerQuota: state.usesOwnerQuota + ? (state.ownerUserId ?? undefined) + : undefined, + }; + } + + private userMemberLimit(plan: string) { + return plan === 'pro' || plan === 'lifetime_pro' || plan === 'selfhost_free' + ? 10 + : 3; + } + + private planName(plan: string) { + switch (plan) { + case 'pro': + case 'selfhost_free': + return 'Pro'; + case 'lifetime_pro': + return 'Lifetime Pro'; + case 'ai': + return 'AI'; + case 'team': + case 'selfhost_team': + return 'Team'; + default: + return 'Free'; + } + } } diff --git a/packages/backend/server/src/core/quota/state.ts b/packages/backend/server/src/core/quota/state.ts new file mode 100644 index 0000000000..435e492a78 --- /dev/null +++ b/packages/backend/server/src/core/quota/state.ts @@ -0,0 +1,413 @@ +import { Injectable } from '@nestjs/common'; +import { PrismaClient } from '@prisma/client'; + +import { EventBus, OnEvent } from '../../base'; +import { EntitlementService } from '../entitlement'; + +type Quota = Awaited< + ReturnType +>['quota']; + +const STATE_TTL = 1000 * 60 * 10; + +declare global { + interface Events { + 'user.quota_state.changed': { + userId: string; + }; + 'workspace.quota_state.changed': { + workspaceId: string; + }; + } +} + +@Injectable() +export class QuotaStateService { + constructor( + private readonly db: PrismaClient, + private readonly entitlement: EntitlementService, + private readonly event: EventBus + ) {} + + async reconcileUserQuotaState(userId: string) { + const [previous, entitlement, entitlements, resolved, usedStorageQuota] = + await Promise.all([ + this.db.effectiveUserQuotaState.findUnique({ where: { userId } }), + this.entitlement.getBestEntitlement('user', userId), + this.entitlement.getActiveEntitlements('user', userId), + this.entitlement.resolveUserEntitlement(userId), + this.getOwnerStorageUsage(userId), + ]); + const flags = { + ...resolved.flags, + unlimitedCopilot: entitlements.some( + entitlement => entitlement.plan === 'ai' + ), + }; + const now = new Date(); + + const state = await this.db.effectiveUserQuotaState.upsert({ + where: { userId }, + update: { + plan: resolved.plan, + sourceEntitlementId: entitlement?.id ?? null, + ...this.quotaData(resolved.quota), + usedStorageQuota, + flags, + known: true, + stale: false, + lastReconciledAt: now, + staleAfter: this.staleAfter(now), + }, + create: { + userId, + plan: resolved.plan, + sourceEntitlementId: entitlement?.id ?? null, + ...this.quotaData(resolved.quota), + usedStorageQuota, + flags, + known: true, + stale: false, + lastReconciledAt: now, + staleAfter: this.staleAfter(now), + }, + }); + if (this.userQuotaStateChanged(previous, state)) { + await this.event.emitAsync('user.quota_state.changed', { userId }); + } + return state; + } + + async reconcileWorkspaceQuotaState(workspaceId: string) { + const owner = await this.getWorkspaceOwner(workspaceId); + const [ + previous, + entitlement, + resolved, + memberCount, + workspaceStorageUsage, + ] = await Promise.all([ + this.db.effectiveWorkspaceQuotaState.findUnique({ + where: { workspaceId }, + }), + this.entitlement.getBestEntitlement('workspace', workspaceId), + this.entitlement.resolveWorkspaceEntitlement(workspaceId), + this.getChargedMemberCount(workspaceId), + this.getWorkspaceStorageUsage(workspaceId), + ]); + const usesOwnerQuota = !this.hasStandaloneWorkspaceQuota(resolved.plan); + const [ownerState, ownerEntitlement] = usesOwnerQuota + ? await Promise.all([ + this.reconcileUserQuotaState(owner.id), + this.entitlement.resolveUserEntitlement(owner.id), + ]) + : [null, null]; + const quota = ownerEntitlement?.quota ?? resolved.quota; + const plan = ownerEntitlement?.plan ?? resolved.plan; + const usedStorageQuota = ownerState + ? ownerState.usedStorageQuota + : workspaceStorageUsage; + const storageQuota = BigInt(quota.storageQuota); + const seatLimit = quota.seatLimit ?? 0; + const overcapacityMemberCount = Math.max(memberCount - seatLimit, 0); + const readonlyReasons = [ + overcapacityMemberCount > 0 ? 'member_overflow' : null, + usedStorageQuota > storageQuota ? 'storage_overflow' : null, + ].filter((reason): reason is string => !!reason); + const now = new Date(); + + const state = await this.db.effectiveWorkspaceQuotaState.upsert({ + where: { workspaceId }, + update: { + plan, + sourceEntitlementId: entitlement?.id ?? null, + ownerUserId: owner.id, + usesOwnerQuota, + seatLimit, + memberCount, + overcapacityMemberCount, + ...this.workspaceQuotaData(quota), + usedStorageQuota, + readonly: readonlyReasons.length > 0, + readonlyReasons, + flags: resolved.flags, + known: true, + stale: false, + lastReconciledAt: now, + staleAfter: this.staleAfter(now), + }, + create: { + workspaceId, + plan, + sourceEntitlementId: entitlement?.id ?? null, + ownerUserId: owner.id, + usesOwnerQuota, + seatLimit, + memberCount, + overcapacityMemberCount, + ...this.workspaceQuotaData(quota), + usedStorageQuota, + readonly: readonlyReasons.length > 0, + readonlyReasons, + flags: resolved.flags, + known: true, + stale: false, + lastReconciledAt: now, + staleAfter: this.staleAfter(now), + }, + }); + if (this.workspaceQuotaStateChanged(previous, state)) { + await this.event.emitAsync('workspace.quota_state.changed', { + workspaceId, + }); + } + return state; + } + + async reconcileAllEntitlementProjection() { + const [users, workspaces] = await Promise.all([ + this.db.user.findMany({ select: { id: true } }), + this.db.workspace.findMany({ select: { id: true } }), + ]); + + await this.reconcileMany([ + ...users.map(user => () => this.reconcileUserQuotaState(user.id)), + ...workspaces.map( + workspace => () => this.reconcileWorkspaceQuotaState(workspace.id) + ), + ]); + } + + @OnEvent('entitlement.changed') + async onEntitlementChanged({ + targetType, + targetId, + }: Events['entitlement.changed']) { + if (targetType === 'user') { + await this.reconcileUserQuotaState(targetId); + await this.reconcileOwnedWorkspaces(targetId); + } else if (targetType === 'workspace') { + await this.reconcileWorkspaceQuotaState(targetId); + } + } + + @OnEvent('workspace.members.updated') + async onWorkspaceMembersUpdated({ + workspaceId, + }: Events['workspace.members.updated']) { + await this.reconcileWorkspaceQuotaState(workspaceId); + } + + @OnEvent('workspace.owner.changed') + async onWorkspaceOwnerChanged({ + workspaceId, + from, + to, + }: Events['workspace.owner.changed']) { + await this.reconcileWorkspaceQuotaState(workspaceId); + await Promise.all([ + this.reconcileUserQuotaState(from), + this.reconcileUserQuotaState(to), + ]); + } + + @OnEvent('workspace.blobs.updated') + async onWorkspaceBlobsUpdated({ + workspaceId, + }: Events['workspace.blobs.updated']) { + const owner = await this.getWorkspaceOwner(workspaceId); + await Promise.all([ + this.reconcileWorkspaceQuotaState(workspaceId), + this.reconcileUserQuotaState(owner.id), + ]); + } + + private async reconcileOwnedWorkspaces(userId: string) { + const workspaces = await this.getOwnedWorkspaceIds(userId); + + await this.reconcileMany( + workspaces.map( + workspaceId => () => this.reconcileWorkspaceQuotaState(workspaceId) + ) + ); + } + + private async getOwnerStorageUsage(userId: string) { + const workspaces = await this.getOwnedWorkspaceIds(userId); + const usages = await this.mapMany(workspaces, async workspaceId => { + const entitlement = + await this.entitlement.resolveWorkspaceEntitlement(workspaceId); + + return this.hasStandaloneWorkspaceQuota(entitlement.plan) + ? 0n + : this.getWorkspaceStorageUsage(workspaceId); + }); + + return usages.reduce((total, usage) => total + usage, 0n); + } + + private async getWorkspaceOwner(workspaceId: string) { + const owner = await this.db.workspaceMember.findFirst({ + where: { + workspaceId, + role: 'owner', + state: 'active', + }, + select: { + user: { + select: { + id: true, + }, + }, + }, + }); + if (!owner) { + throw new Error('Workspace owner not found'); + } + return owner.user; + } + + private async getChargedMemberCount(workspaceId: string) { + const [members, invitations] = await Promise.all([ + this.db.workspaceMember.count({ + where: { workspaceId, state: 'active' }, + }), + this.db.workspaceInvitation.count({ + where: { + workspaceId, + status: { + not: 'waiting_review', + }, + }, + }), + ]); + return members + invitations; + } + + private async getOwnedWorkspaceIds(userId: string) { + const workspaces = await this.db.workspaceMember.findMany({ + where: { + userId, + role: 'owner', + state: 'active', + }, + select: { + workspaceId: true, + }, + }); + return workspaces.map(workspace => workspace.workspaceId); + } + + private async getWorkspaceStorageUsage(workspaceId: string) { + const sum = await this.db.blob.aggregate({ + where: { + workspaceId, + deletedAt: null, + }, + _sum: { + size: true, + }, + }); + + return BigInt(sum._sum.size ?? 0); + } + + private hasStandaloneWorkspaceQuota(plan: string) { + return plan === 'team' || plan === 'selfhost_team'; + } + + private quotaData(quota: Quota) { + return { + blobLimit: BigInt(quota.blobLimit), + storageQuota: BigInt(quota.storageQuota), + historyPeriodSeconds: quota.historyPeriod, + copilotActionLimit: quota.copilotActionLimit ?? null, + }; + } + + private workspaceQuotaData(quota: Quota) { + return { + blobLimit: BigInt(quota.blobLimit), + storageQuota: BigInt(quota.storageQuota), + historyPeriodSeconds: quota.historyPeriod, + }; + } + + private async reconcileMany(tasks: Array<() => Promise>) { + await this.mapMany(tasks, task => task()); + } + + private async mapMany(items: T[], mapper: (item: T) => Promise) { + const batchSize = 16; + const results: U[] = []; + for (let index = 0; index < items.length; index += batchSize) { + results.push( + ...(await Promise.all( + items.slice(index, index + batchSize).map(item => mapper(item)) + )) + ); + } + return results; + } + + private userQuotaStateChanged( + previous: Awaited< + ReturnType + >, + current: Awaited< + ReturnType + > + ) { + if (!previous) { + return true; + } + return ( + previous.plan !== current.plan || + previous.sourceEntitlementId !== current.sourceEntitlementId || + previous.blobLimit !== current.blobLimit || + previous.storageQuota !== current.storageQuota || + previous.usedStorageQuota !== current.usedStorageQuota || + previous.historyPeriodSeconds !== current.historyPeriodSeconds || + previous.copilotActionLimit !== current.copilotActionLimit || + previous.known !== current.known || + previous.stale !== current.stale || + JSON.stringify(previous.flags) !== JSON.stringify(current.flags) + ); + } + + private workspaceQuotaStateChanged( + previous: Awaited< + ReturnType + >, + current: Awaited< + ReturnType + > + ) { + if (!previous) { + return true; + } + return ( + previous.plan !== current.plan || + previous.sourceEntitlementId !== current.sourceEntitlementId || + previous.ownerUserId !== current.ownerUserId || + previous.usesOwnerQuota !== current.usesOwnerQuota || + previous.seatLimit !== current.seatLimit || + previous.memberCount !== current.memberCount || + previous.overcapacityMemberCount !== current.overcapacityMemberCount || + previous.blobLimit !== current.blobLimit || + previous.storageQuota !== current.storageQuota || + previous.usedStorageQuota !== current.usedStorageQuota || + previous.historyPeriodSeconds !== current.historyPeriodSeconds || + previous.readonly !== current.readonly || + previous.known !== current.known || + previous.stale !== current.stale || + previous.readonlyReasons.join(',') !== + current.readonlyReasons.join(',') || + JSON.stringify(previous.flags) !== JSON.stringify(current.flags) + ); + } + + private staleAfter(now: Date) { + return new Date(now.getTime() + STATE_TTL); + } +} diff --git a/packages/backend/server/src/core/quota/utils.ts b/packages/backend/server/src/core/quota/utils.ts index b4688799d1..edf00ad357 100644 --- a/packages/backend/server/src/core/quota/utils.ts +++ b/packages/backend/server/src/core/quota/utils.ts @@ -1,4 +1,4 @@ -import { OneDay, OneKB } from '../../base'; +import { OneKB } from '../../base'; export const ByteUnit = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB']; @@ -14,6 +14,6 @@ export function formatSize(bytes: number, decimals: number = 2): string { ); } -export function formatDate(ms: number): string { - return `${(ms / OneDay).toFixed(0)} days`; +export function formatDate(seconds: number): string { + return `${(seconds / (24 * 60 * 60)).toFixed(0)} days`; } diff --git a/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts b/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts index aff9cc4f0d..529f3ac09c 100644 --- a/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts +++ b/packages/backend/server/src/core/realtime/__tests__/registry.spec.ts @@ -1,4 +1,7 @@ -import { getRealtimeInputKey } from '@affine/realtime'; +import { + getRealtimeInputKey, + type WorkspaceQuotaStateSnapshot, +} from '@affine/realtime'; import test from 'ava'; import { z } from 'zod'; @@ -7,13 +10,15 @@ import { CopilotTranscriptRealtimeProvider } from '../../../plugins/copilot/tran import type { CurrentUser } from '../../auth'; import { CommentRealtimeProvider } from '../../comment/realtime'; import { NotificationRealtimeProvider } from '../../notification/realtime'; -import type { AccessController } from '../../permission'; +import type { PermissionAccess } from '../../permission'; +import { QuotaStateRealtimeProvider } from '../../quota/realtime'; import { RealtimeGateway } from '../gateway'; import { realtimeCommentRoom, realtimeNotificationRoom, realtimeTranscriptTaskRoom, realtimeWorkspaceEmbeddingProgressRoom, + realtimeWorkspaceQuotaStateRoom, registerRealtimeLiveQuery, } from '../index'; import { RealtimePublisher } from '../publisher'; @@ -214,6 +219,103 @@ test('realtime providers expose runtime injection metadata for registry dependen CopilotTranscriptRealtimeProvider ).includes(RealtimeRegistry) ); + t.true( + Reflect.getMetadata( + 'design:paramtypes', + QuotaStateRealtimeProvider + ).includes(RealtimeRegistry) + ); +}); + +test('quota realtime provider exposes effective quota state snapshots', async t => { + const registry = new RealtimeRegistry(); + const provider = new QuotaStateRealtimeProvider( + { + workspaceUser: { + getActive: async () => ({ role: 'admin' }), + }, + } as never, + { + reconcileUserQuotaState: async () => ({ + userId: 'u1', + plan: 'pro', + sourceEntitlementId: null, + blobLimit: 1n, + storageQuota: 2n, + usedStorageQuota: 3n, + historyPeriodSeconds: 4, + copilotActionLimit: null, + flags: {}, + known: true, + stale: false, + lastReconciledAt: null, + staleAfter: null, + createdAt: new Date(0), + updatedAt: new Date(0), + }), + reconcileWorkspaceQuotaState: async () => ({ + workspaceId: 'space', + plan: 'team', + sourceEntitlementId: null, + ownerUserId: 'u1', + usesOwnerQuota: false, + seatLimit: 5, + memberCount: 4, + overcapacityMemberCount: 0, + blobLimit: 6n, + storageQuota: 7n, + usedStorageQuota: 8n, + historyPeriodSeconds: 9, + readonly: false, + readonlyReasons: [], + flags: {}, + known: true, + stale: false, + lastReconciledAt: null, + staleAfter: null, + createdAt: new Date(0), + updatedAt: new Date(0), + }), + } as never, + registry + ); + + provider.onModuleInit(); + + t.deepEqual( + await registry.getRequest('user.quota-state.get').handle(user, {}), + { + state: { + userId: 'u1', + plan: 'pro', + sourceEntitlementId: null, + blobLimit: 1, + storageQuota: 2, + usedStorageQuota: 3, + historyPeriodSeconds: 4, + copilotActionLimit: null, + flags: {}, + known: true, + stale: false, + lastReconciledAt: null, + staleAfter: null, + createdAt: new Date(0), + updatedAt: new Date(0), + }, + } + ); + const workspaceQuotaState = (await registry + .getRequest('workspace.quota-state.get') + .handle(user, { workspaceId: 'space' })) as { + state: WorkspaceQuotaStateSnapshot; + }; + t.is(workspaceQuotaState.state.memberCount, 4); + t.is( + registry + .getTopic('workspace.quota-state.changed') + .room(user, { workspaceId: 'space' }), + realtimeWorkspaceQuotaStateRoom('space') + ); }); test('copilot transcript realtime provider registers task live query handlers', async t => { @@ -234,7 +336,7 @@ test('copilot transcript realtime provider registers task live query handlers', }, }; }, - } as unknown as AccessController; + } as unknown as PermissionAccess; const transcript = { async queryTask( userId: string, diff --git a/packages/backend/server/src/core/realtime/index.ts b/packages/backend/server/src/core/realtime/index.ts index d21afcdfc0..e3f4e0ba4f 100644 --- a/packages/backend/server/src/core/realtime/index.ts +++ b/packages/backend/server/src/core/realtime/index.ts @@ -18,9 +18,11 @@ export { realtimeCommentRoom, realtimeNotificationRoom, realtimeTranscriptTaskRoom, + realtimeUserQuotaStateRoom, realtimeUserRoom, realtimeWorkspaceDocRoom, realtimeWorkspaceEmbeddingProgressRoom, + realtimeWorkspaceQuotaStateRoom, realtimeWorkspaceRoom, } from './rooms'; export type { RealtimeRequestHandler, RealtimeTopicHandler } from './types'; diff --git a/packages/backend/server/src/core/realtime/rooms.ts b/packages/backend/server/src/core/realtime/rooms.ts index 8a3e00f117..1363746eb2 100644 --- a/packages/backend/server/src/core/realtime/rooms.ts +++ b/packages/backend/server/src/core/realtime/rooms.ts @@ -32,3 +32,11 @@ export function realtimeCommentRoom(workspaceId: string, docId: string) { export function realtimeWorkspaceEmbeddingProgressRoom(workspaceId: string) { return realtimeWorkspaceRoom(workspaceId, 'embedding-progress'); } + +export function realtimeUserQuotaStateRoom(userId: string) { + return realtimeUserRoom(userId, 'quota-state'); +} + +export function realtimeWorkspaceQuotaStateRoom(workspaceId: string) { + return realtimeWorkspaceRoom(workspaceId, 'quota-state'); +} diff --git a/packages/backend/server/src/core/sync/gateway.ts b/packages/backend/server/src/core/sync/gateway.ts index 6f0b7da15a..cfcfb51db8 100644 --- a/packages/backend/server/src/core/sync/gateway.ts +++ b/packages/backend/server/src/core/sync/gateway.ts @@ -41,8 +41,8 @@ import { } from '../doc'; import { applyUpdatesWithNative } from '../doc/merge-updates'; import { - AccessController, type DocAction, + PermissionAccess, WorkspaceAction, } from '../permission'; import { DocID } from '../utils/doc'; @@ -223,7 +223,7 @@ export class SpaceSyncGateway private activeUsersFlushQueued = false; constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly event: EventBus, private readonly workspace: PgWorkspaceDocStorageAdapter, private readonly userspace: PgUserspaceDocStorageAdapter, @@ -899,7 +899,7 @@ class WorkspaceSyncAdapter extends SyncSocketAdapter { constructor( client: Socket, storage: DocStorageAdapter, - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly docReader: DocReader, private readonly models: Models ) { diff --git a/packages/backend/server/src/core/workspaces/controller.ts b/packages/backend/server/src/core/workspaces/controller.ts index be8938c35c..063427333e 100644 --- a/packages/backend/server/src/core/workspaces/controller.ts +++ b/packages/backend/server/src/core/workspaces/controller.ts @@ -29,7 +29,7 @@ import { buildPublicRootDoc } from '../../native'; import { CurrentUser, Public } from '../auth'; import { PgWorkspaceDocStorageAdapter } from '../doc'; import { DocReader } from '../doc/reader'; -import { AccessController, WorkspacePolicyService } from '../permission'; +import { PermissionAccess } from '../permission'; import { CommentAttachmentStorage, WorkspaceBlobStorage } from '../storage'; import { DocID } from '../utils/doc'; @@ -39,8 +39,7 @@ export class WorkspacesController { constructor( private readonly storage: WorkspaceBlobStorage, private readonly commentAttachmentStorage: CommentAttachmentStorage, - private readonly ac: AccessController, - private readonly workspacePolicy: WorkspacePolicyService, + private readonly ac: PermissionAccess, private readonly workspace: PgWorkspaceDocStorageAdapter, private readonly docReader: DocReader, private readonly models: Models @@ -113,7 +112,7 @@ export class WorkspacesController { .workspace(workspaceId) .can('Workspace.Read'); const canReadSharedWorkspaceBlobs = - await this.workspacePolicy.canReadWorkspaceBySharedDocs(workspaceId); + await this.canReadSharedWorkspaceBlobs(workspaceId); if (!canReadWorkspace && !canReadSharedWorkspaceBlobs) { throw new SpaceAccessDenied({ spaceId: workspaceId }); } @@ -163,6 +162,14 @@ export class WorkspacesController { body.pipe(res); } + private async canReadSharedWorkspaceBlobs(workspaceId: string) { + const [sharingEnabled, publicDocs] = await Promise.all([ + this.models.workspace.allowSharing(workspaceId), + this.models.docAccessPolicy.hasPublicExternal(workspaceId), + ]); + return sharingEnabled && publicDocs; + } + // get doc binary @Public() @Get('/:id/docs/:guid') diff --git a/packages/backend/server/src/core/workspaces/resolvers/admin.ts b/packages/backend/server/src/core/workspaces/resolvers/admin.ts index 1a59780545..20457a03e2 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/admin.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/admin.ts @@ -425,9 +425,6 @@ class AdminUpdateWorkspaceInput extends PartialType( ) { @Field() id!: string; - - @Field(() => [Feature], { nullable: true }) - features?: WorkspaceFeatureName[]; } @Injectable() @@ -617,28 +614,40 @@ export class AdminWorkspaceResolver { query, pagination ); - return list.map(({ user, status, type }) => ({ - id: user.id, - name: user.name, - email: user.email, - avatarUrl: user.avatarUrl, - role: type, - status, - })); + return list.flatMap(({ user, status, type }) => + user + ? [ + { + id: user.id, + name: user.name, + email: user.email, + avatarUrl: user.avatarUrl, + role: type, + status, + }, + ] + : [] + ); } const [list] = await this.models.workspaceUser.paginate( workspaceId, pagination ); - return list.map(({ user, status, type }) => ({ - id: user.id, - name: user.name, - email: user.email, - avatarUrl: user.avatarUrl, - role: type, - status, - })); + return list.flatMap(({ user, status, type }) => + user + ? [ + { + id: user.id, + name: user.name, + email: user.email, + avatarUrl: user.avatarUrl, + role: type, + status, + }, + ] + : [] + ); } @ResolveField(() => [AdminWorkspaceSharedLink], { @@ -654,7 +663,7 @@ export class AdminWorkspaceResolver { } @Mutation(() => AdminWorkspace, { - description: 'Update workspace flags and features for admin', + description: 'Update workspace flags for admin', nullable: true, }) async adminUpdateWorkspace( @@ -662,27 +671,12 @@ export class AdminWorkspaceResolver { input: AdminUpdateWorkspaceInput ) { this.assertCloudOnly(); - const { id, features, ...updates } = input; + const { id, ...updates } = input; if (Object.keys(updates).length) { await this.models.workspace.update(id, updates); } - if (features) { - const current = await this.models.workspaceFeature.list(id); - const toAdd = features.filter(feature => !current.includes(feature)); - const toRemove = current.filter(feature => !features.includes(feature)); - - await Promise.all([ - ...toAdd.map(feature => - this.models.workspaceFeature.add(id, feature, 'admin panel update') - ), - ...toRemove.map(feature => - this.models.workspaceFeature.remove(id, feature) - ), - ]); - } - const { rows } = await this.models.workspace.adminListWorkspaces({ first: 1, skip: 0, diff --git a/packages/backend/server/src/core/workspaces/resolvers/blob.ts b/packages/backend/server/src/core/workspaces/resolvers/blob.ts index b026a2cef4..ca50957038 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/blob.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/blob.ts @@ -25,7 +25,7 @@ import { } from '../../../base'; import { Models } from '../../../models'; import { CurrentUser } from '../../auth'; -import { AccessController, WorkspacePolicyService } from '../../permission'; +import { PermissionAccess } from '../../permission'; import { QuotaService } from '../../quota'; import { WorkspaceBlobStorage } from '../../storage'; import { @@ -125,8 +125,7 @@ class ListedBlob { export class WorkspaceBlobResolver { logger = new Logger(WorkspaceBlobResolver.name); constructor( - private readonly ac: AccessController, - private readonly policy: WorkspacePolicyService, + private readonly ac: PermissionAccess, private readonly quota: QuotaService, private readonly storage: WorkspaceBlobStorage, private readonly models: Models @@ -467,7 +466,10 @@ export class WorkspaceBlobResolver { return false; } - await this.policy.assertCanDeleteBlob(user.id, workspaceId); + await this.ac + .user(user.id) + .workspace(workspaceId) + .assert('Workspace.Blobs.Write'); await this.storage.delete(workspaceId, key, permanently); @@ -479,7 +481,10 @@ export class WorkspaceBlobResolver { @CurrentUser() user: CurrentUser, @Args('workspaceId') workspaceId: string ) { - await this.policy.assertCanDeleteBlob(user.id, workspaceId); + await this.ac + .user(user.id) + .workspace(workspaceId) + .assert('Workspace.Blobs.Write'); await this.storage.release(workspaceId); diff --git a/packages/backend/server/src/core/workspaces/resolvers/doc.ts b/packages/backend/server/src/core/workspaces/resolvers/doc.ts index 17b0abd009..95c43a1e01 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/doc.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/doc.ts @@ -11,7 +11,7 @@ import { ResolveField, Resolver, } from '@nestjs/graphql'; -import { PrismaClient } from '@prisma/client'; +import { Prisma, PrismaClient } from '@prisma/client'; import { SafeIntResolver } from 'graphql-scalars'; import { @@ -34,11 +34,11 @@ import { Models, PublicDocMode } from '../../../models'; import { CurrentUser } from '../../auth'; import { Editor } from '../../doc'; import { - AccessController, DOC_ACTIONS, DocAction, DocRole, - WorkspacePolicyService, + PermissionAccess, + PermissionService, } from '../../permission'; import { PublicUserType, WorkspaceUserType } from '../../user'; import { WorkspaceType } from '../types'; @@ -295,8 +295,8 @@ export class WorkspaceDocResolver { * @deprecated migrate to models */ private readonly prisma: PrismaClient, - private readonly ac: AccessController, - private readonly policy: WorkspacePolicyService, + private readonly ac: PermissionAccess, + private readonly permission: PermissionService, private readonly models: Models, private readonly cache: Cache ) {} @@ -361,16 +361,32 @@ export class WorkspaceDocResolver { @Parent() workspace: WorkspaceType, @Args('pagination', PaginationInput.decode) pagination: PaginationInput ): Promise { - const [count, rows] = await this.models.doc.paginateDocInfoByUpdatedAt( - workspace.id, - pagination - ); - const needs = await this.ac - .user(me.id) - .workspace(workspace.id) - .docs(rows, 'Doc.Read'); + const predicate = this.permission.docReadableSqlPredicate({ + userId: me.id, + workspaceId: workspace.id, + action: 'Doc.Read', + docIdColumn: Prisma.raw('"workspace_pages"."page_id"'), + }); + const fallbackPredicate = this.permission.fallbackDocReadableSqlPredicate({ + userId: me.id, + workspaceId: workspace.id, + action: 'Doc.Read', + docIdColumn: Prisma.raw('"workspace_pages"."page_id"'), + }); + const [count, rows] = await this.models.doc + .paginateDocInfoByUpdatedAt(workspace.id, pagination, predicate) + .catch(error => { + if (!fallbackPredicate) { + throw error; + } + return this.models.doc.paginateDocInfoByUpdatedAt( + workspace.id, + pagination, + fallbackPredicate + ); + }); - return paginate(needs, 'updatedAt', pagination, count); + return paginate(rows, 'updatedAt', pagination, count); } @ResolveField(() => DocType, { @@ -423,7 +439,7 @@ export class WorkspaceDocResolver { throw new ExpectToPublishDoc(); } - await this.policy.assertCanPublishDoc(user.id, workspaceId, docId); + await this.ac.user(user.id).doc(workspaceId, docId).assert('Doc.Publish'); const doc = await this.models.doc.publish(workspaceId, docId, mode); @@ -448,7 +464,7 @@ export class WorkspaceDocResolver { throw new ExpectToRevokePublicDoc('Expect doc not to be workspace'); } - await this.policy.assertCanUnpublishDoc(user.id, workspaceId, docId); + await this.ac.user(user.id).doc(workspaceId, docId).assert('Doc.Publish'); const doc = await this.models.doc.unpublish(workspaceId, docId); @@ -518,7 +534,7 @@ export class DocResolver { private readonly logger = new Logger(DocResolver.name); constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly models: Models ) {} diff --git a/packages/backend/server/src/core/workspaces/resolvers/history.ts b/packages/backend/server/src/core/workspaces/resolvers/history.ts index 193afb2075..b87dee5686 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/history.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/history.ts @@ -13,7 +13,7 @@ import type { SnapshotHistory } from '@prisma/client'; import { CurrentUser } from '../../auth'; import { PgWorkspaceDocStorageAdapter } from '../../doc'; -import { AccessController } from '../../permission'; +import { PermissionAccess } from '../../permission'; import { DocID } from '../../utils/doc'; import { WorkspaceType } from '../types'; import { EditorType } from './doc'; @@ -37,7 +37,7 @@ class DocHistoryType implements Partial { export class DocHistoryResolver { constructor( private readonly workspace: PgWorkspaceDocStorageAdapter, - private readonly ac: AccessController + private readonly ac: PermissionAccess ) {} @ResolveField(() => [DocHistoryType]) diff --git a/packages/backend/server/src/core/workspaces/resolvers/member.ts b/packages/backend/server/src/core/workspaces/resolvers/member.ts index 0e7916eda6..3ab6789b2a 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/member.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/member.ts @@ -25,8 +25,10 @@ import { mapAnyError, MemberNotFoundInSpace, NoMoreSeat, + OwnerCanNotLeaveWorkspace, QueryTooLong, RequestMutex, + SpaceAccessDenied, Throttle, TooManyRequest, URLHelper, @@ -35,7 +37,7 @@ import { import { Models } from '../../../models'; import { CurrentUser, Public } from '../../auth'; import { - AccessController, + PermissionAccess, WorkspacePolicyService, WorkspaceRole, } from '../../permission'; @@ -63,7 +65,7 @@ export class WorkspaceMemberResolver { private readonly cache: Cache, private readonly event: EventBus, private readonly url: URLHelper, - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly models: Models, private readonly mutex: RequestMutex, private readonly policy: WorkspacePolicyService, @@ -115,7 +117,8 @@ export class WorkspaceMemberResolver { return list.map(({ id, status, type, user }) => ({ ...user, - permission: type, + permission: Number(type), + role: Number(type), inviteId: id, status, })); @@ -127,7 +130,8 @@ export class WorkspaceMemberResolver { return list.map(({ id, status, type, user }) => ({ ...user, - permission: type, + permission: Number(type), + role: Number(type), inviteId: id, status, })); @@ -157,7 +161,7 @@ export class WorkspaceMemberResolver { } const quota = await this.quota.getWorkspaceSeatQuota(workspaceId); - const isTeam = await this.models.workspace.isTeamWorkspace(workspaceId); + const isTeam = await this.workspaceService.isTeamWorkspace(workspaceId); const results: InviteResult[] = []; @@ -307,7 +311,10 @@ export class WorkspaceMemberResolver { @CurrentUser() user: CurrentUser, @Args('workspaceId') workspaceId: string ) { - await this.policy.assertCanManageInviteLink(user.id, workspaceId); + await this.ac + .user(user.id) + .workspace(workspaceId) + .assert('Workspace.Users.Manage'); const cacheId = `workspace:inviteLink:${workspaceId}`; return await this.cache.delete(cacheId); @@ -324,7 +331,8 @@ export class WorkspaceMemberResolver { .workspace(workspaceId) .assert('Workspace.Users.Manage'); - const isTeam = await this.models.workspace.isTeamWorkspace(workspaceId); + const quota = await this.quota.getWorkspaceSeatQuota(workspaceId); + const isTeam = await this.workspaceService.isTeamWorkspace(workspaceId); const role = await this.models.workspaceUser.get(workspaceId, userId); if (role) { @@ -339,7 +347,6 @@ export class WorkspaceMemberResolver { } ); } else { - const quota = await this.quota.getWorkspaceSeatQuota(workspaceId); if (quota.memberCount >= quota.memberLimit) { throw new NoMoreSeat({ spaceId: workspaceId }); } else { @@ -454,7 +461,14 @@ export class WorkspaceMemberResolver { throw new MemberNotFoundInSpace({ spaceId: workspaceId }); } - await this.policy.assertCanRevokeMember(me.id, workspaceId, role.type); + await this.ac + .user(me.id) + .workspace(workspaceId) + .assert( + role.type === WorkspaceRole.Admin + ? 'Workspace.Administrators.Manage' + : 'Workspace.Users.Manage' + ); await this.models.workspaceUser.delete(workspaceId, userId); @@ -554,7 +568,16 @@ export class WorkspaceMemberResolver { }) _workspaceName?: string ) { - await this.policy.assertCanLeaveWorkspace(user.id, workspaceId); + const role = await this.models.workspaceUser.getActive( + workspaceId, + user.id + ); + if (!role) { + throw new MemberNotFoundInSpace({ spaceId: workspaceId }); + } + if (role.type === WorkspaceRole.Owner) { + throw new OwnerCanNotLeaveWorkspace(); + } await this.models.workspaceUser.delete(workspaceId, user.id); this.event.emit('workspace.members.leave', { @@ -571,7 +594,7 @@ export class WorkspaceMemberResolver { } private async acceptInvitationByEmail(role: WorkspaceUserRole) { - await this.policy.assertCanInviteMembers(role.workspaceId); + await this.assertWorkspaceAcceptsMemberChange(role.workspaceId); const hasSeat = await this.quota.tryCheckSeat(role.workspaceId, true); @@ -598,7 +621,7 @@ export class WorkspaceMemberResolver { workspaceId: string, inviterId: string ) { - await this.policy.assertCanInviteMembers(workspaceId); + await this.assertWorkspaceAcceptsMemberChange(workspaceId); let inviter = await this.models.user.getPublicUser(inviterId); if (!inviter) { @@ -619,4 +642,11 @@ export class WorkspaceMemberResolver { await this.workspaceService.sendReviewRequestNotification(role.id); return; } + + private async assertWorkspaceAcceptsMemberChange(workspaceId: string) { + const state = await this.policy.getWorkspaceState(workspaceId); + if (state.isReadonly) { + throw new SpaceAccessDenied({ spaceId: workspaceId }); + } + } } diff --git a/packages/backend/server/src/core/workspaces/resolvers/workspace.ts b/packages/backend/server/src/core/workspaces/resolvers/workspace.ts index e7f4b5578e..0c7f9a5942 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/workspace.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/workspace.ts @@ -20,7 +20,7 @@ import { import { Models } from '../../../models'; import { CurrentUser } from '../../auth'; import { - AccessController, + PermissionAccess, WORKSPACE_ACTIONS, WorkspaceAction, WorkspaceRole, @@ -79,7 +79,7 @@ export class WorkspaceRolePermissions { @Resolver(() => WorkspaceType) export class WorkspaceResolver { constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly quota: QuotaService, private readonly models: Models, private readonly workspaceService: WorkspaceService, diff --git a/packages/backend/server/src/core/workspaces/service.ts b/packages/backend/server/src/core/workspaces/service.ts index 60e1234084..3085668e82 100644 --- a/packages/backend/server/src/core/workspaces/service.ts +++ b/packages/backend/server/src/core/workspaces/service.ts @@ -10,6 +10,7 @@ import { import { DocReader } from '../doc'; import { Mailer } from '../mail'; import { WorkspaceRole } from '../permission'; +import { QuotaStateService } from '../quota/state'; import { WorkspaceBlobStorage } from '../storage'; export type InviteInfo = { @@ -30,7 +31,8 @@ export class WorkspaceService { private readonly doc: DocReader, private readonly blobStorage: WorkspaceBlobStorage, private readonly mailer: Mailer, - private readonly queue: JobQueue + private readonly queue: JobQueue, + private readonly quotaState: QuotaStateService ) {} async getInviteInfo(inviteId: string): Promise { @@ -99,7 +101,9 @@ export class WorkspaceService { // ================ Team ================ async isTeamWorkspace(workspaceId: string) { - return this.models.workspace.isTeamWorkspace(workspaceId); + const state = + await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); + return ['team', 'selfhost_team'].includes(state.plan); } async sendTeamWorkspaceUpgradedEmail(workspaceId: string) { diff --git a/packages/backend/server/src/data/migrations/1765500000000-backfill-permission-projection.ts b/packages/backend/server/src/data/migrations/1765500000000-backfill-permission-projection.ts new file mode 100644 index 0000000000..2e9d67f63e --- /dev/null +++ b/packages/backend/server/src/data/migrations/1765500000000-backfill-permission-projection.ts @@ -0,0 +1,28 @@ +import { ModuleRef } from '@nestjs/core'; +import { PrismaClient } from '@prisma/client'; + +import { WorkspacePolicyService } from '../../core/permission/policy'; +import { Models } from '../../models'; + +export class BackfillPermissionProjection1765500000000 { + static async up(_db: PrismaClient, ref: ModuleRef) { + const models = ref.get(Models, { strict: false }); + await models.permissionProjection.backfillLegacyProjection(); + + const policy = ref.get(WorkspacePolicyService, { strict: false }); + const workspaces = await _db.workspace.findMany({ + select: { id: true }, + }); + for (const workspace of workspaces) { + const state = await policy.getWorkspaceState(workspace.id); + await models.workspaceRuntimeState.upsert(workspace.id, { + readonly: state.isReadonly, + readonlyReasons: state.readonlyReasons, + known: true, + staleAfter: null, + }); + } + } + + static async down(_db: PrismaClient) {} +} diff --git a/packages/backend/server/src/data/migrations/1765600000000-backfill-entitlement-projection.ts b/packages/backend/server/src/data/migrations/1765600000000-backfill-entitlement-projection.ts new file mode 100644 index 0000000000..c039dc1dc3 --- /dev/null +++ b/packages/backend/server/src/data/migrations/1765600000000-backfill-entitlement-projection.ts @@ -0,0 +1,35 @@ +import { ModuleRef } from '@nestjs/core'; +import { PrismaClient } from '@prisma/client'; + +import { LegacyEntitlementProjectionService } from '../../core/entitlement'; +import { QuotaStateService } from '../../core/quota/state'; + +export class BackfillEntitlementProjection1765600000000 { + static async up(db: PrismaClient, ref: ModuleRef) { + const projection = ref.get(LegacyEntitlementProjectionService, { + strict: false, + }); + await projection.backfillEntitlementsAndQuotaStates(); + + const quota = ref.get(QuotaStateService, { strict: false }); + const [users, workspaces] = await Promise.all([ + db.user.findMany({ select: { id: true } }), + db.workspace.findMany({ select: { id: true } }), + ]); + + const tasks = [ + ...users.map(user => () => quota.reconcileUserQuotaState(user.id)), + ...workspaces.map( + workspace => () => quota.reconcileWorkspaceQuotaState(workspace.id) + ), + ]; + const batchSize = 16; + for (let index = 0; index < tasks.length; index += batchSize) { + await Promise.all( + tasks.slice(index, index + batchSize).map(task => task()) + ); + } + } + + static async down(_db: PrismaClient) {} +} diff --git a/packages/backend/server/src/data/migrations/index.ts b/packages/backend/server/src/data/migrations/index.ts index d3b11f2175..390b9fc149 100644 --- a/packages/backend/server/src/data/migrations/index.ts +++ b/packages/backend/server/src/data/migrations/index.ts @@ -4,3 +4,5 @@ export * from './1721299086340-refresh-unnamed-user'; export * from './1745211351719-create-indexer-tables'; export * from './1751966744168-correct-session-update-time'; export * from './1763800000000-rebuild-manticore-mixed-script-indexes'; +export * from './1765500000000-backfill-permission-projection'; +export * from './1765600000000-backfill-entitlement-projection'; diff --git a/packages/backend/server/src/models/__tests__/comment.spec.ts b/packages/backend/server/src/models/__tests__/comment.spec.ts index c55e81caa5..10ce7d587b 100644 --- a/packages/backend/server/src/models/__tests__/comment.spec.ts +++ b/packages/backend/server/src/models/__tests__/comment.spec.ts @@ -15,6 +15,9 @@ const workspace = await module.create(Mockers.Workspace, { owner, }); +const waitNextMillisecond = () => + new Promise(resolve => setTimeout(resolve, 1)); + test.after.always(async () => { await module.close(); }); @@ -77,7 +80,6 @@ test('should get a comment', async t => { docId, userId: owner.id, }); - const comment2 = await models.comment.get(comment1.id); t.deepEqual(comment2, comment1); t.deepEqual(comment2?.content, { @@ -146,7 +148,7 @@ test('should resolve a comment', async t => { docId, userId: owner.id, }); - + await waitNextMillisecond(); const comment2 = await models.comment.resolve({ id: comment.id, resolved: true, @@ -158,6 +160,7 @@ test('should resolve a comment', async t => { // updatedAt should be changed t.true(comment3!.updatedAt.getTime() > comment3!.createdAt.getTime()); + await waitNextMillisecond(); const comment4 = await models.comment.resolve({ id: comment.id, resolved: false, @@ -272,6 +275,7 @@ test('should update a reply', async t => { commentId: comment.id, }); + await waitNextMillisecond(); const reply2 = await models.comment.updateReply({ id: reply.id, content: { @@ -322,6 +326,7 @@ test('should list comments with replies', async t => { userId: owner.id, }); + await waitNextMillisecond(); const comment2 = await models.comment.create({ content: { type: 'paragraph', @@ -342,6 +347,7 @@ test('should list comments with replies', async t => { userId: owner.id, }); + await waitNextMillisecond(); const reply1 = await models.comment.createReply({ userId: owner.id, content: { @@ -351,6 +357,7 @@ test('should list comments with replies', async t => { commentId: comment1.id, }); + await waitNextMillisecond(); const reply2 = await models.comment.createReply({ userId: owner.id, content: { @@ -421,7 +428,9 @@ test('should list changes', async t => { docId, userId: owner.id, }); + const comment1Cursor = comment1.updatedAt; + await waitNextMillisecond(); const comment2 = await models.comment.create({ content: { type: 'paragraph', @@ -432,6 +441,7 @@ test('should list changes', async t => { userId: owner.id, }); + await waitNextMillisecond(); const reply1 = await models.comment.createReply({ userId: owner.id, content: { @@ -441,6 +451,7 @@ test('should list changes', async t => { commentId: comment1.id, }); + await waitNextMillisecond(); const reply2 = await models.comment.createReply({ userId: owner.id, content: { @@ -465,7 +476,7 @@ test('should list changes', async t => { t.is((changes1[2].item as Reply).commentId, comment1.id); const changes2 = await models.comment.listChanges(workspace.id, docId, { - commentUpdatedAt: comment1.updatedAt, + commentUpdatedAt: comment1Cursor, replyUpdatedAt: reply1.updatedAt, }); t.is(changes2.length, 2); @@ -476,6 +487,7 @@ test('should list changes', async t => { t.is(changes2[1].commentId, comment1.id); // update comment1 + await waitNextMillisecond(); const comment1Updated = await models.comment.update({ id: comment1.id, content: { @@ -493,8 +505,10 @@ test('should list changes', async t => { t.is(changes3[0].id, comment1Updated.id); // delete comment1 and reply1, update reply2 + await waitNextMillisecond(); await models.comment.delete(comment1.id); await models.comment.deleteReply(reply1.id); + await waitNextMillisecond(); await models.comment.updateReply({ id: reply2.id, content: { diff --git a/packages/backend/server/src/models/base.ts b/packages/backend/server/src/models/base.ts index 832be3c320..4e226474e0 100644 --- a/packages/backend/server/src/models/base.ts +++ b/packages/backend/server/src/models/base.ts @@ -19,4 +19,13 @@ export class BaseModel { // See https://papooch.github.io/nestjs-cls/plugins/available-plugins/transactional#using-the-injecttransaction-decorator return this.txHost.tx; } + + protected async withPermissionProjectionMetric(operation: Promise) { + try { + return await operation; + } catch (err) { + this.models.permissionProjection.recordTriggerErrorMetric(err); + throw err; + } + } } diff --git a/packages/backend/server/src/models/doc-user.ts b/packages/backend/server/src/models/doc-user.ts index f189fd5d01..1a25292306 100644 --- a/packages/backend/server/src/models/doc-user.ts +++ b/packages/backend/server/src/models/doc-user.ts @@ -3,11 +3,12 @@ import assert from 'node:assert'; import { Injectable } from '@nestjs/common'; import { Transactional } from '@nestjs-cls/transactional'; import type { TransactionalAdapterPrisma } from '@nestjs-cls/transactional-adapter-prisma'; -import { WorkspaceDocUserRole } from '@prisma/client'; +import { DocGrant, WorkspaceDocUserRole } from '@prisma/client'; import { CanNotBatchGrantDocOwnerPermissions, PaginationInput } from '../base'; import { BaseModel } from './base'; import { DocRole } from './common'; +import { docRoleFromNew } from './permission-write'; @Injectable() export class DocUserModel extends BaseModel { @@ -17,36 +18,7 @@ export class DocUserModel extends BaseModel { */ @Transactional({ timeout: 15000 }) async setOwner(workspaceId: string, docId: string, userId: string) { - await this.db.workspaceDocUserRole.updateMany({ - where: { - workspaceId, - docId, - type: DocRole.Owner, - userId: { not: userId }, - }, - data: { - type: DocRole.Manager, - }, - }); - - await this.db.workspaceDocUserRole.upsert({ - where: { - workspaceId_docId_userId: { - workspaceId, - docId, - userId, - }, - }, - update: { - type: DocRole.Owner, - }, - create: { - workspaceId, - docId, - userId, - type: DocRole.Owner, - }, - }); + await this.models.docGrant.setOwner(workspaceId, docId, userId); this.logger.log( `Set doc owner of [${workspaceId}/${docId}] to [${userId}]` ); @@ -62,34 +34,11 @@ export class DocUserModel extends BaseModel { // internal misuse, throw directly assert(role !== DocRole.Owner, 'Cannot set Owner role of a doc to a user.'); - const oldRole = await this.get(workspaceId, docId, userId); - - if (oldRole && oldRole.type === role) { - return oldRole; - } - - const newRole = await this.db.workspaceDocUserRole.upsert({ - where: { - workspaceId_docId_userId: { - workspaceId, - docId, - userId, - }, - }, - update: { - type: role, - }, - create: { - workspaceId, - docId, - userId, - type: role, - }, - }); - - return newRole; + await this.models.docGrant.set(workspaceId, docId, userId, role); + return await this.get(workspaceId, docId, userId); } + @Transactional() async batchSetUserRoles( workspaceId: string, docId: string, @@ -104,76 +53,83 @@ export class DocUserModel extends BaseModel { throw new CanNotBatchGrantDocOwnerPermissions(); } - const result = await this.db.workspaceDocUserRole.createMany({ - skipDuplicates: true, - data: userIds.map(userId => ({ - workspaceId, - docId, - userId, - type: role, - })), - }); - - return result.count; + return await this.models.docGrant.batchSetUserRoles( + workspaceId, + docId, + userIds, + role + ); } + @Transactional() async delete(workspaceId: string, docId: string, userId: string) { - await this.db.workspaceDocUserRole.deleteMany({ - where: { - workspaceId, - docId, - userId, - }, - }); + await this.models.docGrant.delete(workspaceId, docId, userId); } + @Transactional() async deleteByUserId(userId: string) { - await this.db.workspaceDocUserRole.deleteMany({ + await this.models.permissionProjection.markNewWriteOrigin(); + await this.db.docGrant.deleteMany({ where: { - userId, + principalType: 'user', + principalId: userId, }, }); + await this.withPermissionProjectionMetric( + this.db.workspaceDocUserRole.deleteMany({ + where: { + userId, + }, + }) + ); } async getOwner(workspaceId: string, docId: string) { - return await this.db.workspaceDocUserRole.findFirst({ + const grant = await this.db.docGrant.findFirst({ where: { workspaceId, docId, - type: DocRole.Owner, + principalType: 'user', + role: 'owner', }, }); + return grant ? this.docGrantToCompat(grant) : null; } async get(workspaceId: string, docId: string, userId: string) { - return await this.db.workspaceDocUserRole.findUnique({ + const grant = await this.db.docGrant.findUnique({ where: { - workspaceId_docId_userId: { + workspaceId_docId_principalType_principalId: { workspaceId, docId, - userId, + principalType: 'user', + principalId: userId, }, }, }); + return grant ? this.docGrantToCompat(grant) : null; } async findMany(workspaceId: string, docIds: string[], userId: string) { - return await this.db.workspaceDocUserRole.findMany({ + const grants = await this.db.docGrant.findMany({ where: { workspaceId, docId: { in: docIds, }, - userId, + principalType: 'user', + principalId: userId, }, }); + return grants.map(grant => this.docGrantToCompat(grant)); } count(workspaceId: string, docId: string) { - return this.db.workspaceDocUserRole.count({ + return this.db.docGrant.count({ where: { workspaceId, docId, + principalType: 'user', }, }); } @@ -183,11 +139,12 @@ export class DocUserModel extends BaseModel { docId: string, pagination: PaginationInput ): Promise<[WorkspaceDocUserRole[], number]> { - return await Promise.all([ - this.db.workspaceDocUserRole.findMany({ + const [grants, total] = await Promise.all([ + this.db.docGrant.findMany({ where: { workspaceId, docId, + principalType: 'user', createdAt: pagination.after ? { gte: pagination.after, @@ -202,5 +159,16 @@ export class DocUserModel extends BaseModel { }), this.count(workspaceId, docId), ]); + return [grants.map(grant => this.docGrantToCompat(grant)), total]; + } + + private docGrantToCompat(grant: DocGrant): WorkspaceDocUserRole { + return { + workspaceId: grant.workspaceId, + docId: grant.docId, + userId: grant.principalId, + type: docRoleFromNew(grant.role as never), + createdAt: grant.createdAt, + }; } } diff --git a/packages/backend/server/src/models/doc.ts b/packages/backend/server/src/models/doc.ts index 86adcdd991..da44ec259f 100644 --- a/packages/backend/server/src/models/doc.ts +++ b/packages/backend/server/src/models/doc.ts @@ -351,27 +351,44 @@ export class DocModel extends BaseModel { /** * Create or update the doc meta. */ + @Transactional() async upsertMeta( workspaceId: string, docId: string, data?: DocMetaUpsertInput ) { - const doc = await this.db.workspaceDoc.upsert({ - where: { - workspaceId_docId: { + if ( + data && + ('public' in data || 'defaultRole' in data || 'publishedAt' in data) + ) { + await this.models.docAccessPolicy.upsert(workspaceId, docId, { + public: data.public, + defaultRole: data.defaultRole, + publishedAt: + typeof data.publishedAt === 'string' + ? new Date(data.publishedAt) + : data.publishedAt, + }); + } + + const doc = await this.withPermissionProjectionMetric( + this.db.workspaceDoc.upsert({ + where: { + workspaceId_docId: { + workspaceId, + docId, + }, + }, + update: { + ...data, + }, + create: { + ...data, workspaceId, docId, }, - }, - update: { - ...data, - }, - create: { - ...data, - workspaceId, - docId, - }, - }); + }) + ); this.event.emit('doc.updated', { workspaceId, docId, @@ -643,13 +660,17 @@ export class DocModel extends BaseModel { async paginateDocInfoByUpdatedAt( workspaceId: string, - pagination: PaginationInput + pagination: PaginationInput, + readablePredicate: Prisma.Sql = Prisma.sql`TRUE` ) { - const count = await this.db.workspaceDoc.count({ - where: { - workspaceId, - }, - }); + const [countRow] = await this.db.$queryRaw<{ count: bigint | number }[]>` + SELECT COUNT(*) AS count + FROM "workspace_pages" + WHERE + "workspace_pages"."workspace_id" = ${workspaceId} + AND ${readablePredicate} + `; + const count = Number(countRow?.count ?? 0); const after = pagination.after ? Prisma.sql`AND "snapshots"."updated_at" < ${new Date(pagination.after)}` @@ -686,6 +707,7 @@ export class DocModel extends BaseModel { AND "workspace_pages"."page_id" = "snapshots"."guid" WHERE "workspace_pages"."workspace_id" = ${workspaceId} + AND ${readablePredicate} ${after} ORDER BY "snapshots"."updated_at" DESC diff --git a/packages/backend/server/src/models/index.ts b/packages/backend/server/src/models/index.ts index b92f655f23..8935ac8001 100644 --- a/packages/backend/server/src/models/index.ts +++ b/packages/backend/server/src/models/index.ts @@ -30,6 +30,14 @@ import { FeatureModel } from './feature'; import { HistoryModel } from './history'; import { MagicLinkOtpModel } from './magic-link-otp'; import { NotificationModel } from './notification'; +import { PermissionProjectionModel } from './permission-projection'; +import { + DocAccessPolicyModel, + DocGrantModel, + WorkspaceAccessPolicyModel, + WorkspaceInvitationModel, + WorkspaceMemberModel, +} from './permission-write'; import { MODELS_SYMBOL } from './provider'; import { SessionModel } from './session'; import { UserModel } from './user'; @@ -41,6 +49,7 @@ import { WorkspaceModel } from './workspace'; import { WorkspaceAnalyticsModel } from './workspace-analytics'; import { WorkspaceCalendarModel } from './workspace-calendar'; import { WorkspaceFeatureModel } from './workspace-feature'; +import { WorkspaceRuntimeStateModel } from './workspace-runtime-state'; import { WorkspaceUserModel } from './workspace-user'; const MODELS = { @@ -52,12 +61,19 @@ const MODELS = { workspace: WorkspaceModel, userFeature: UserFeatureModel, workspaceFeature: WorkspaceFeatureModel, + workspaceRuntimeState: WorkspaceRuntimeStateModel, doc: DocModel, userDoc: UserDocModel, workspaceUser: WorkspaceUserModel, docUser: DocUserModel, history: HistoryModel, notification: NotificationModel, + permissionProjection: PermissionProjectionModel, + workspaceMember: WorkspaceMemberModel, + workspaceInvitation: WorkspaceInvitationModel, + workspaceAccessPolicy: WorkspaceAccessPolicyModel, + docAccessPolicy: DocAccessPolicyModel, + docGrant: DocGrantModel, userSettings: UserSettingsModel, copilotSession: CopilotSessionModel, copilotUsage: CopilotUsageModel, @@ -150,6 +166,8 @@ export * from './feature'; export * from './history'; export * from './magic-link-otp'; export * from './notification'; +export * from './permission-projection'; +export * from './permission-write'; export * from './session'; export * from './user'; export * from './user-doc'; @@ -160,4 +178,5 @@ export * from './workspace'; export * from './workspace-analytics'; export * from './workspace-calendar'; export * from './workspace-feature'; +export * from './workspace-runtime-state'; export * from './workspace-user'; diff --git a/packages/backend/server/src/models/permission-projection.ts b/packages/backend/server/src/models/permission-projection.ts new file mode 100644 index 0000000000..3b2c5ab4e5 --- /dev/null +++ b/packages/backend/server/src/models/permission-projection.ts @@ -0,0 +1,555 @@ +import { Injectable, Optional } from '@nestjs/common'; +import { Prisma, PrismaClient } from '@prisma/client'; + +import { metrics } from '../base'; +import { BaseModel } from './base'; + +type CountRow = { count: bigint }; + +type ProjectionIssueRow = { + category: string; + count: bigint; +}; + +type ProjectionBackfillDb = { + $transaction: ( + callback: ( + tx: Pick + ) => Promise + ) => Promise; +}; + +export type PermissionProjectionCheckReport = { + oldWorkspacePolicyMismatch: number; + oldAcceptedMemberMismatch: number; + extraProjectedMember: number; + oldInvitationMismatch: number; + extraProjectedInvitation: number; + oldDocGrantMismatch: number; + extraProjectedDocGrant: number; + oldDocPolicyMismatch: number; + extraProjectedDocPolicy: number; + runtimeStateMissing: number; + runtimeStateMismatch: number; + ownerConflict: number; + oldNewDecisionMismatch: number; + invalidLegacyRows: Record; +}; + +export const PERMISSION_PROJECTION_TRIGGER_ERROR_CATEGORIES = [ + 'owner_conflict', + 'invalid_legacy_role', + 'foreign_key_missing', + 'projection_recursion_guard_missing', + 'unknown', +] as const; + +export function permissionProjectionTriggerErrorCategory(error: unknown) { + const message = + error instanceof Error ? error.message : String(error ?? 'unknown'); + + const match = message.match(/permission_projection_error:([^:]+):/); + const category = match?.[1]; + if (!category) { + return null; + } + + return PERMISSION_PROJECTION_TRIGGER_ERROR_CATEGORIES.includes( + category as (typeof PERMISSION_PROJECTION_TRIGGER_ERROR_CATEGORIES)[number] + ) + ? category + : 'unknown'; +} + +async function count(first: Promise) { + const rows = await first; + return Number(rows[0]?.count ?? 0); +} + +@Injectable() +export class PermissionProjectionModel extends BaseModel { + constructor(@Optional() private readonly prisma?: PrismaClient) { + super(); + } + + async backfillLegacyProjection() { + const db = (this.prisma ?? this.db) as unknown as ProjectionBackfillDb; + + await db.$transaction(async tx => { + await tx.$executeRaw` + DELETE FROM workspace_members projected + WHERE projected.legacy_permission_id IS NOT NULL + AND NOT EXISTS ( + SELECT 1 + FROM workspace_user_permissions old + WHERE old.id = projected.legacy_permission_id + AND old.status = 'Accepted'::"WorkspaceMemberStatus" + AND affine_permission_legacy_workspace_role(old.type) IS NOT NULL + ) + `; + + await tx.$executeRaw` + DELETE FROM workspace_invitations projected + WHERE projected.legacy_permission_id IS NOT NULL + AND NOT EXISTS ( + SELECT 1 + FROM workspace_user_permissions old + WHERE old.id = projected.legacy_permission_id + AND old.status <> 'Accepted'::"WorkspaceMemberStatus" + AND affine_permission_workspace_invitation_state(old.status) IS NOT NULL + AND affine_permission_legacy_workspace_role(old.type) IS NOT NULL + ) + `; + + await tx.$executeRaw` + DELETE FROM doc_grants projected + WHERE projected.principal_type = 'user' + AND NOT EXISTS ( + SELECT 1 + FROM workspace_page_user_permissions old + WHERE old.workspace_id = projected.workspace_id + AND old.page_id = projected.doc_id + AND old.user_id = projected.principal_id + AND affine_permission_legacy_doc_role(old.type) IS NOT NULL + ) + `; + + await tx.$executeRaw` + DELETE FROM doc_access_policies projected + WHERE NOT EXISTS ( + SELECT 1 + FROM workspace_pages old + WHERE old.workspace_id = projected.workspace_id + AND old.page_id = projected.doc_id + AND affine_permission_legacy_default_doc_role(old."defaultRole") IS NOT NULL + ) + `; + + await tx.$executeRaw` + DELETE FROM workspace_access_policies projected + WHERE NOT EXISTS ( + SELECT 1 + FROM workspaces old + WHERE old.id = projected.workspace_id + ) + `; + + await tx.$executeRaw` + INSERT INTO workspace_access_policies ( + workspace_id, + visibility, + sharing_enabled, + url_preview_enabled, + updated_at + ) + SELECT + id, + CASE WHEN public THEN 'public' ELSE 'private' END, + enable_sharing, + enable_url_preview, + now() + FROM workspaces + ON CONFLICT (workspace_id) + DO UPDATE SET + visibility = EXCLUDED.visibility, + sharing_enabled = EXCLUDED.sharing_enabled, + url_preview_enabled = EXCLUDED.url_preview_enabled, + updated_at = now() + `; + + await tx.$executeRaw` + INSERT INTO workspace_members ( + workspace_id, + user_id, + role, + state, + source, + legacy_permission_id, + created_at, + updated_at + ) + SELECT + workspace_id, + user_id, + affine_permission_legacy_workspace_role(type), + 'active', + CASE source + WHEN 'Email'::"WorkspaceMemberSource" THEN 'email' + WHEN 'Link'::"WorkspaceMemberSource" THEN 'link' + ELSE 'legacy' + END, + id, + created_at, + updated_at + FROM workspace_user_permissions + WHERE status = 'Accepted'::"WorkspaceMemberStatus" + AND affine_permission_legacy_workspace_role(type) IS NOT NULL + ON CONFLICT ("legacy_permission_id") WHERE "legacy_permission_id" IS NOT NULL + DO UPDATE SET + user_id = EXCLUDED.user_id, + role = EXCLUDED.role, + state = EXCLUDED.state, + source = EXCLUDED.source, + updated_at = EXCLUDED.updated_at + `; + + await tx.$executeRaw` + INSERT INTO workspace_invitations ( + workspace_id, + invitee_user_id, + inviter_user_id, + requested_role, + status, + kind, + legacy_permission_id, + created_at, + updated_at + ) + SELECT + workspace_id, + user_id, + inviter_id, + CASE WHEN affine_permission_legacy_workspace_role(type) = 'admin' THEN 'admin' ELSE 'member' END, + affine_permission_workspace_invitation_state(status), + CASE source + WHEN 'Link'::"WorkspaceMemberSource" THEN 'link' + ELSE 'email' + END, + id, + created_at, + updated_at + FROM workspace_user_permissions + WHERE status <> 'Accepted'::"WorkspaceMemberStatus" + AND affine_permission_workspace_invitation_state(status) IS NOT NULL + AND affine_permission_legacy_workspace_role(type) IS NOT NULL + ON CONFLICT ("legacy_permission_id") WHERE "legacy_permission_id" IS NOT NULL + DO UPDATE SET + invitee_user_id = EXCLUDED.invitee_user_id, + inviter_user_id = EXCLUDED.inviter_user_id, + requested_role = EXCLUDED.requested_role, + status = EXCLUDED.status, + kind = EXCLUDED.kind, + updated_at = EXCLUDED.updated_at + `; + + await tx.$executeRaw` + INSERT INTO doc_access_policies ( + workspace_id, + doc_id, + visibility, + public_role, + member_default_role, + published_at, + updated_at + ) + SELECT + workspace_id, + page_id, + CASE WHEN public THEN 'public' ELSE 'private' END, + CASE WHEN public THEN 'external' ELSE NULL END, + affine_permission_legacy_default_doc_role("defaultRole"), + published_at, + now() + FROM workspace_pages + WHERE affine_permission_legacy_default_doc_role("defaultRole") IS NOT NULL + ON CONFLICT (workspace_id, doc_id) + DO UPDATE SET + visibility = EXCLUDED.visibility, + public_role = EXCLUDED.public_role, + member_default_role = EXCLUDED.member_default_role, + published_at = EXCLUDED.published_at, + updated_at = now() + `; + + await tx.$executeRaw` + INSERT INTO doc_grants ( + workspace_id, + doc_id, + principal_type, + principal_id, + role, + legacy_workspace_id, + legacy_doc_id, + legacy_user_id, + created_at, + updated_at + ) + SELECT + workspace_id, + page_id, + 'user', + user_id, + affine_permission_legacy_doc_role(type), + workspace_id, + page_id, + user_id, + created_at, + now() + FROM workspace_page_user_permissions + WHERE affine_permission_legacy_doc_role(type) IS NOT NULL + ON CONFLICT (workspace_id, doc_id, principal_type, principal_id) + DO UPDATE SET + role = EXCLUDED.role, + updated_at = now() + `; + }); + } + + recordTriggerErrorMetric(error: unknown) { + const category = permissionProjectionTriggerErrorCategory(error); + if (!category) { + return null; + } + + metrics.permission + .counter('projection_trigger_errors', { + description: 'Permission projection trigger error count', + }) + .add(1, { category }); + return category; + } + + async checkLegacyProjection(): Promise { + const [ + oldWorkspacePolicyMismatch, + oldAcceptedMemberMismatch, + extraProjectedMember, + oldInvitationMismatch, + extraProjectedInvitation, + oldDocGrantMismatch, + extraProjectedDocGrant, + oldDocPolicyMismatch, + extraProjectedDocPolicy, + ownerConflict, + invalidLegacyRows, + ] = await Promise.all([ + count(this.db.$queryRaw` + SELECT COUNT(*)::bigint AS count + FROM workspaces old + LEFT JOIN workspace_access_policies projected + ON projected.workspace_id = old.id + WHERE projected.workspace_id IS NULL + OR projected.visibility <> CASE WHEN old.public THEN 'public' ELSE 'private' END + OR projected.sharing_enabled <> old.enable_sharing + OR projected.url_preview_enabled <> old.enable_url_preview + `), + count(this.db.$queryRaw` + SELECT COUNT(*)::bigint AS count + FROM workspace_user_permissions old + LEFT JOIN workspace_members projected + ON projected.legacy_permission_id = old.id + OR ( + projected.legacy_permission_id IS NULL + AND projected.workspace_id = old.workspace_id + AND projected.user_id = old.user_id + AND projected.state = 'active' + ) + WHERE old.status = 'Accepted'::"WorkspaceMemberStatus" + AND affine_permission_legacy_workspace_role(old.type) IS NOT NULL + AND ( + projected.id IS NULL OR + projected.workspace_id <> old.workspace_id OR + projected.user_id <> old.user_id OR + projected.role <> affine_permission_legacy_workspace_role(old.type) OR + projected.state <> 'active' + ) + `), + count(this.db.$queryRaw` + SELECT COUNT(*)::bigint AS count + FROM workspace_members projected + LEFT JOIN workspace_user_permissions old + ON old.id = projected.legacy_permission_id + OR ( + projected.legacy_permission_id IS NULL + AND old.workspace_id = projected.workspace_id + AND old.user_id = projected.user_id + AND old.status = 'Accepted'::"WorkspaceMemberStatus" + ) + WHERE + projected.state = 'active' + AND ( + old.id IS NULL OR + old.status <> 'Accepted'::"WorkspaceMemberStatus" OR + affine_permission_legacy_workspace_role(old.type) IS NULL + ) + `), + count(this.db.$queryRaw` + SELECT COUNT(*)::bigint AS count + FROM workspace_user_permissions old + LEFT JOIN workspace_invitations projected + ON projected.legacy_permission_id = old.id + OR ( + projected.legacy_permission_id IS NULL + AND projected.workspace_id = old.workspace_id + AND projected.invitee_user_id = old.user_id + ) + WHERE old.status <> 'Accepted'::"WorkspaceMemberStatus" + AND affine_permission_workspace_invitation_state(old.status) IS NOT NULL + AND affine_permission_legacy_workspace_role(old.type) IS NOT NULL + AND ( + projected.id IS NULL OR + projected.workspace_id <> old.workspace_id OR + projected.invitee_user_id <> old.user_id OR + projected.requested_role <> CASE WHEN affine_permission_legacy_workspace_role(old.type) = 'admin' THEN 'admin' ELSE 'member' END OR + projected.status <> affine_permission_workspace_invitation_state(old.status) + ) + `), + count(this.db.$queryRaw` + SELECT COUNT(*)::bigint AS count + FROM workspace_invitations projected + LEFT JOIN workspace_user_permissions old + ON old.id = projected.legacy_permission_id + OR ( + projected.legacy_permission_id IS NULL + AND old.workspace_id = projected.workspace_id + AND old.user_id = projected.invitee_user_id + AND old.status <> 'Accepted'::"WorkspaceMemberStatus" + ) + WHERE projected.invitee_user_id IS NOT NULL + AND ( + old.id IS NULL OR + old.status = 'Accepted'::"WorkspaceMemberStatus" OR + affine_permission_workspace_invitation_state(old.status) IS NULL OR + affine_permission_legacy_workspace_role(old.type) IS NULL + ) + `), + count(this.db.$queryRaw` + SELECT COUNT(*)::bigint AS count + FROM workspace_page_user_permissions old + LEFT JOIN doc_grants projected + ON projected.workspace_id = old.workspace_id + AND projected.doc_id = old.page_id + AND projected.principal_type = 'user' + AND projected.principal_id = old.user_id + WHERE affine_permission_legacy_doc_role(old.type) IS NOT NULL + AND ( + projected.workspace_id IS NULL OR + projected.role <> affine_permission_legacy_doc_role(old.type) + ) + `), + count(this.db.$queryRaw` + SELECT COUNT(*)::bigint AS count + FROM doc_grants projected + LEFT JOIN workspace_page_user_permissions old + ON old.workspace_id = projected.workspace_id + AND old.page_id = projected.doc_id + AND old.user_id = projected.principal_id + WHERE projected.principal_type = 'user' + AND ( + old.workspace_id IS NULL OR + affine_permission_legacy_doc_role(old.type) IS NULL + ) + `), + count(this.db.$queryRaw` + SELECT COUNT(*)::bigint AS count + FROM workspace_pages old + LEFT JOIN doc_access_policies projected + ON projected.workspace_id = old.workspace_id + AND projected.doc_id = old.page_id + WHERE affine_permission_legacy_default_doc_role(old."defaultRole") IS NOT NULL + AND ( + projected.workspace_id IS NULL OR + projected.visibility <> CASE WHEN old.public THEN 'public' ELSE 'private' END OR + projected.public_role IS DISTINCT FROM CASE WHEN old.public THEN 'external' ELSE NULL END OR + projected.member_default_role IS DISTINCT FROM affine_permission_legacy_default_doc_role(old."defaultRole") + ) + `), + count(this.db.$queryRaw` + SELECT COUNT(*)::bigint AS count + FROM doc_access_policies projected + LEFT JOIN workspace_pages old + ON old.workspace_id = projected.workspace_id + AND old.page_id = projected.doc_id + WHERE old.workspace_id IS NULL + `), + count(this.db.$queryRaw` + SELECT COALESCE(SUM(conflicts.count - 1), 0)::bigint AS count + FROM ( + SELECT workspace_id, COUNT(*)::bigint AS count + FROM workspace_members + WHERE state = 'active' + AND role = 'owner' + GROUP BY workspace_id + HAVING COUNT(*) > 1 + UNION ALL + SELECT workspace_id || ':' || doc_id AS workspace_id, COUNT(*)::bigint AS count + FROM doc_grants + WHERE principal_type = 'user' + AND role = 'owner' + GROUP BY workspace_id, doc_id + HAVING COUNT(*) > 1 + ) conflicts + `), + this.db.$queryRaw` + SELECT category, COUNT(*)::bigint AS count + FROM ( + SELECT 'unknown_workspace_role' AS category + FROM workspace_user_permissions + WHERE affine_permission_legacy_workspace_role(type) IS NULL + AND type <> -99 + UNION ALL + SELECT 'unknown_doc_role' AS category + FROM workspace_page_user_permissions + WHERE affine_permission_legacy_doc_role(type) IS NULL + AND type NOT IN (0, -32768) + UNION ALL + SELECT 'legacy_doc_external_row' AS category + FROM workspace_page_user_permissions + WHERE type = 0 + UNION ALL + SELECT 'legacy_doc_none_row' AS category + FROM workspace_page_user_permissions + WHERE type = -32768 + UNION ALL + SELECT 'doc_default_owner' AS category + FROM workspace_pages + WHERE "defaultRole" = 99 + ) issues + GROUP BY category + `, + ]); + + return { + oldWorkspacePolicyMismatch, + oldAcceptedMemberMismatch, + extraProjectedMember, + oldInvitationMismatch, + extraProjectedInvitation, + oldDocGrantMismatch, + extraProjectedDocGrant, + oldDocPolicyMismatch, + extraProjectedDocPolicy, + runtimeStateMissing: 0, + runtimeStateMismatch: 0, + ownerConflict, + oldNewDecisionMismatch: 0, + invalidLegacyRows: Object.fromEntries( + invalidLegacyRows.map(row => [row.category, Number(row.count)]) + ), + }; + } + + async lockWorkspaceOwnerTransfer(workspaceId: string) { + await this.db.$executeRaw` + SELECT pg_advisory_xact_lock(hashtextextended(${workspaceId}, 16)) + `; + } + + async lockDocOwnerTransfer(workspaceId: string, docId: string) { + await this.db.$executeRaw` + SELECT pg_advisory_xact_lock(hashtextextended(${`${workspaceId}:${docId}`}, 16)) + `; + } + + async markNewWriteOrigin() { + await this.db.$executeRaw` + SELECT set_config('affine.permission_sync_origin', 'new', true) + `; + } + + async markLegacyWriteOrigin() { + await this.db.$executeRaw` + SELECT set_config('affine.permission_sync_origin', 'legacy', true) + `; + } +} diff --git a/packages/backend/server/src/models/permission-write.ts b/packages/backend/server/src/models/permission-write.ts new file mode 100644 index 0000000000..941509f030 --- /dev/null +++ b/packages/backend/server/src/models/permission-write.ts @@ -0,0 +1,742 @@ +import assert from 'node:assert'; + +import { Injectable } from '@nestjs/common'; +import { Transactional } from '@nestjs-cls/transactional'; +import { WorkspaceMemberSource, WorkspaceMemberStatus } from '@prisma/client'; + +import { CanNotBatchGrantDocOwnerPermissions } from '../base'; +import { BaseModel } from './base'; +import { DocRole, WorkspaceRole } from './common'; + +type WorkspaceMemberRole = 'owner' | 'admin' | 'member'; +type WorkspaceInvitationStatus = 'pending' | 'waiting_review' | 'waiting_seat'; +type WorkspaceInvitationKind = 'email' | 'link'; +type PermissionSource = 'email' | 'link' | 'legacy'; +type DocGrantRole = 'owner' | 'manager' | 'editor' | 'commenter' | 'reader'; + +export function workspaceRoleToNew(role: WorkspaceRole): WorkspaceMemberRole { + switch (role) { + case WorkspaceRole.Owner: + return 'owner'; + case WorkspaceRole.Admin: + return 'admin'; + case WorkspaceRole.Collaborator: + return 'member'; + default: + throw new Error( + `Unsupported workspace role for new permission model: ${role}` + ); + } +} + +export function workspaceRoleFromNew(role: WorkspaceMemberRole): WorkspaceRole { + switch (role) { + case 'owner': + return WorkspaceRole.Owner; + case 'admin': + return WorkspaceRole.Admin; + case 'member': + return WorkspaceRole.Collaborator; + } +} + +export function workspaceStatusFromNew( + state: 'active' | WorkspaceInvitationStatus +): WorkspaceMemberStatus { + switch (state) { + case 'active': + return WorkspaceMemberStatus.Accepted; + case 'pending': + return WorkspaceMemberStatus.Pending; + case 'waiting_review': + return WorkspaceMemberStatus.UnderReview; + case 'waiting_seat': + return WorkspaceMemberStatus.NeedMoreSeat; + } +} + +export function workspaceSourceToNew( + source?: WorkspaceMemberSource +): PermissionSource { + switch (source) { + case WorkspaceMemberSource.Email: + return 'email'; + case WorkspaceMemberSource.Link: + return 'link'; + default: + return 'legacy'; + } +} + +export function workspaceSourceFromNew( + source?: PermissionSource | WorkspaceInvitationKind +): WorkspaceMemberSource { + return source === 'link' + ? WorkspaceMemberSource.Link + : WorkspaceMemberSource.Email; +} + +export function workspaceStatusToInvitationState( + status: WorkspaceMemberStatus +): WorkspaceInvitationStatus | null { + switch (status) { + case WorkspaceMemberStatus.Pending: + return 'pending'; + case WorkspaceMemberStatus.UnderReview: + return 'waiting_review'; + case WorkspaceMemberStatus.AllocatingSeat: + case WorkspaceMemberStatus.NeedMoreSeat: + case WorkspaceMemberStatus.NeedMoreSeatAndReview: + return 'waiting_seat'; + default: + return null; + } +} + +export function docRoleToNew(role: DocRole): DocGrantRole { + switch (role) { + case DocRole.Owner: + return 'owner'; + case DocRole.Manager: + return 'manager'; + case DocRole.Editor: + return 'editor'; + case DocRole.Commenter: + return 'commenter'; + case DocRole.Reader: + return 'reader'; + default: + throw new Error( + `Unsupported doc grant role for new permission model: ${role}` + ); + } +} + +function workspaceInvitationKindToNew( + source?: WorkspaceMemberSource +): WorkspaceInvitationKind { + return source === WorkspaceMemberSource.Link ? 'link' : 'email'; +} + +export function docRoleFromNew(role: DocGrantRole): DocRole { + switch (role) { + case 'owner': + return DocRole.Owner; + case 'manager': + return DocRole.Manager; + case 'editor': + return DocRole.Editor; + case 'commenter': + return DocRole.Commenter; + case 'reader': + return DocRole.Reader; + } +} + +@Injectable() +export class WorkspaceMemberModel extends BaseModel { + @Transactional() + async setOwner( + workspaceId: string, + userId: string, + fallbackRole: WorkspaceRole + ) { + await this.models.permissionProjection.markNewWriteOrigin(); + await this.models.permissionProjection.lockWorkspaceOwnerTransfer( + workspaceId + ); + const ownerCount = await this.db.workspaceMember.count({ + where: { workspaceId, role: 'owner', state: 'active' }, + }); + if (ownerCount > 0) { + const target = await this.db.workspaceMember.findFirst({ + where: { workspaceId, userId, state: 'active' }, + }); + if (!target) { + throw new Error('New workspace owner must be an active member.'); + } + } + + await this.db.workspaceMember.updateMany({ + where: { + workspaceId, + role: 'owner', + userId: { not: userId }, + state: 'active', + }, + data: { + role: workspaceRoleToNew(fallbackRole), + source: 'legacy', + }, + }); + + return await this.db.workspaceMember.upsert({ + where: { + workspaceId_userId_state: { + workspaceId, + userId, + state: 'active', + }, + }, + update: { + role: 'owner', + source: 'legacy', + }, + create: { + workspaceId, + userId, + role: 'owner', + state: 'active', + source: 'legacy', + }, + }); + } + + @Transactional() + async setActive( + workspaceId: string, + userId: string, + role: WorkspaceRole, + data: { legacyPermissionId?: string | null; source?: PermissionSource } = {} + ) { + await this.models.permissionProjection.markNewWriteOrigin(); + if (role === WorkspaceRole.Owner) { + throw new Error('Cannot grant Owner role of a workspace to a user.'); + } + + await this.db.workspaceInvitation.deleteMany({ + where: { workspaceId, inviteeUserId: userId }, + }); + + return await this.db.workspaceMember.upsert({ + where: { + workspaceId_userId_state: { + workspaceId, + userId, + state: 'active', + }, + }, + update: { + role: workspaceRoleToNew(role), + legacyPermissionId: data.legacyPermissionId, + source: data.source, + }, + create: { + workspaceId, + userId, + role: workspaceRoleToNew(role), + state: 'active', + source: data.source ?? 'legacy', + legacyPermissionId: data.legacyPermissionId, + }, + }); + } + + @Transactional() + async delete(workspaceId: string, userId: string) { + await this.models.permissionProjection.markNewWriteOrigin(); + await this.db.$queryRaw` + SELECT id + FROM workspace_members + WHERE workspace_id = ${workspaceId} + AND role = 'owner' + AND state = 'active' + FOR UPDATE + `; + const existingOwners = await this.db.workspaceMember.count({ + where: { + workspaceId, + role: 'owner', + state: 'active', + userId: { not: userId }, + }, + }); + const deletingOwner = await this.db.workspaceMember.count({ + where: { + workspaceId, + userId, + role: 'owner', + state: 'active', + }, + }); + + if (deletingOwner > 0 && existingOwners === 0) { + throw new Error('Cannot remove the last active workspace owner.'); + } + + return await this.db.workspaceMember.deleteMany({ + where: { workspaceId, userId, state: 'active' }, + }); + } +} + +@Injectable() +export class WorkspaceInvitationModel extends BaseModel { + private hasCurrentColumns?: Promise; + + @Transactional() + async set( + workspaceId: string, + userId: string, + role: WorkspaceRole, + status: WorkspaceMemberStatus, + data: { + source?: WorkspaceMemberSource; + inviterId?: string; + } = {} + ): Promise { + await this.models.permissionProjection.markNewWriteOrigin(); + if (role === WorkspaceRole.Owner) { + throw new Error('Cannot grant Owner role of a workspace to a user.'); + } + + const invitationStatus = workspaceStatusToInvitationState(status); + if (!invitationStatus) { + await this.models.workspaceMember.setActive(workspaceId, userId, role); + return; + } + + await this.db.workspaceMember.deleteMany({ + where: { workspaceId, userId, state: 'active' }, + }); + + await this.upsertInvitation({ + workspaceId, + userId, + inviterId: data.inviterId, + requestedRole: role === WorkspaceRole.Admin ? 'admin' : 'member', + status: invitationStatus, + kind: workspaceInvitationKindToNew(data.source), + source: workspaceSourceToNew(data.source), + }); + } + + @Transactional() + async setState( + workspaceId: string, + userId: string, + status: WorkspaceMemberStatus, + data: { + inviterId?: string; + } = {} + ) { + await this.models.permissionProjection.markNewWriteOrigin(); + const invitationStatus = workspaceStatusToInvitationState(status); + if (!invitationStatus) { + const invitation = await this.findInvitation(workspaceId, userId); + if (!invitation) { + throw new Error('Cannot activate a missing workspace invitation.'); + } + const role = + invitation.requestedRole === 'admin' + ? WorkspaceRole.Admin + : WorkspaceRole.Collaborator; + return await this.models.workspaceMember.setActive( + workspaceId, + userId, + role, + { + legacyPermissionId: invitation.legacyPermissionId, + source: invitation.source, + } + ); + } + + return await this.updateInvitationStatus({ + workspaceId, + userId, + status: invitationStatus, + inviterId: data.inviterId, + }); + } + + @Transactional() + async deleteNonAccepted(workspaceId: string) { + await this.models.permissionProjection.markNewWriteOrigin(); + return await this.db.workspaceInvitation.deleteMany({ + where: { workspaceId }, + }); + } + + private async supportsCurrentInvitationColumns() { + this.hasCurrentColumns ??= this.db.$queryRaw>` + SELECT EXISTS ( + SELECT 1 + FROM information_schema.columns + WHERE table_name = 'workspace_invitations' + AND column_name = 'requested_role' + ) AS "exists" + `.then(rows => rows[0]?.exists ?? false); + return await this.hasCurrentColumns; + } + + private async upsertInvitation(input: { + workspaceId: string; + userId: string; + inviterId?: string; + requestedRole: 'admin' | 'member'; + status: WorkspaceInvitationStatus; + kind: WorkspaceInvitationKind; + source: PermissionSource; + }) { + if (await this.supportsCurrentInvitationColumns()) { + return await this.db.$executeRaw` + INSERT INTO workspace_invitations ( + workspace_id, + invitee_user_id, + inviter_user_id, + requested_role, + status, + kind, + updated_at + ) + VALUES ( + ${input.workspaceId}, + ${input.userId}, + ${input.inviterId ?? null}, + ${input.requestedRole}, + ${input.status}, + ${input.kind}, + now() + ) + ON CONFLICT (workspace_id, invitee_user_id) + DO UPDATE SET + inviter_user_id = EXCLUDED.inviter_user_id, + requested_role = EXCLUDED.requested_role, + status = EXCLUDED.status, + kind = EXCLUDED.kind, + updated_at = now() + `; + } + + return await this.db.$executeRaw` + INSERT INTO workspace_invitations ( + workspace_id, + invitee_user_id, + inviter_id, + role, + state, + source, + updated_at + ) + VALUES ( + ${input.workspaceId}, + ${input.userId}, + ${input.inviterId ?? null}, + ${input.requestedRole}, + ${input.status}, + ${input.source}, + now() + ) + ON CONFLICT (workspace_id, invitee_user_id) + DO UPDATE SET + inviter_id = EXCLUDED.inviter_id, + role = EXCLUDED.role, + state = EXCLUDED.state, + source = EXCLUDED.source, + updated_at = now() + `; + } + + private async findInvitation(workspaceId: string, userId: string) { + if (await this.supportsCurrentInvitationColumns()) { + const rows = await this.db.$queryRaw< + Array<{ + requestedRole: 'admin' | 'member'; + legacyPermissionId: string | null; + source: PermissionSource; + }> + >` + SELECT + requested_role AS "requestedRole", + legacy_permission_id AS "legacyPermissionId", + kind AS source + FROM workspace_invitations + WHERE workspace_id = ${workspaceId} + AND invitee_user_id = ${userId} + LIMIT 1 + `; + return rows[0] ?? null; + } + + const rows = await this.db.$queryRaw< + Array<{ + requestedRole: 'admin' | 'member'; + legacyPermissionId: string | null; + source: PermissionSource; + }> + >` + SELECT + role AS "requestedRole", + legacy_permission_id AS "legacyPermissionId", + source + FROM workspace_invitations + WHERE workspace_id = ${workspaceId} + AND invitee_user_id = ${userId} + LIMIT 1 + `; + return rows[0] ?? null; + } + + private async updateInvitationStatus(input: { + workspaceId: string; + userId: string; + status: WorkspaceInvitationStatus; + inviterId?: string; + }) { + if (await this.supportsCurrentInvitationColumns()) { + return await this.db.$executeRaw` + UPDATE workspace_invitations + SET + status = ${input.status}, + inviter_user_id = ${input.inviterId ?? null}, + updated_at = now() + WHERE workspace_id = ${input.workspaceId} + AND invitee_user_id = ${input.userId} + `; + } + + return await this.db.$executeRaw` + UPDATE workspace_invitations + SET + state = ${input.status}, + inviter_id = ${input.inviterId ?? null}, + updated_at = now() + WHERE workspace_id = ${input.workspaceId} + AND invitee_user_id = ${input.userId} + `; + } +} + +@Injectable() +export class WorkspaceAccessPolicyModel extends BaseModel { + @Transactional() + async upsert( + workspaceId: string, + policy: { + public?: boolean; + enableSharing?: boolean; + enableUrlPreview?: boolean; + } + ) { + await this.models.permissionProjection.markNewWriteOrigin(); + return await this.db.workspaceAccessPolicy.upsert({ + where: { workspaceId }, + update: { + visibility: + policy.public === undefined + ? undefined + : policy.public + ? 'public' + : 'private', + sharingEnabled: policy.enableSharing, + urlPreviewEnabled: policy.enableUrlPreview, + }, + create: { + workspaceId, + visibility: policy.public ? 'public' : 'private', + sharingEnabled: policy.enableSharing ?? true, + urlPreviewEnabled: policy.enableUrlPreview ?? false, + }, + }); + } +} + +@Injectable() +export class DocAccessPolicyModel extends BaseModel { + async hasPublicExternal(workspaceId: string) { + const count = await this.db.docAccessPolicy.count({ + where: { + workspaceId, + visibility: 'public', + publicRole: 'external', + }, + }); + return count > 0; + } + + @Transactional() + async upsert( + workspaceId: string, + docId: string, + policy: { + public?: boolean; + defaultRole?: DocRole; + publishedAt?: Date | null; + urlPreviewEnabled?: boolean; + } + ) { + await this.models.permissionProjection.markNewWriteOrigin(); + const publicRole = policy.public ? 'external' : null; + return await this.db.docAccessPolicy.upsert({ + where: { workspaceId_docId: { workspaceId, docId } }, + update: { + visibility: + policy.public === undefined + ? undefined + : policy.public + ? 'public' + : 'private', + publicRole: policy.public === undefined ? undefined : publicRole, + memberDefaultRole: + policy.defaultRole === undefined + ? undefined + : policy.defaultRole === DocRole.None + ? 'none' + : docRoleToNew(policy.defaultRole), + publishedAt: policy.publishedAt, + urlPreviewEnabled: policy.urlPreviewEnabled, + }, + create: { + workspaceId, + docId, + visibility: policy.public ? 'public' : 'private', + publicRole, + memberDefaultRole: + policy.defaultRole === undefined + ? null + : policy.defaultRole === DocRole.None + ? 'none' + : docRoleToNew(policy.defaultRole), + publishedAt: policy.publishedAt, + urlPreviewEnabled: policy.urlPreviewEnabled ?? false, + }, + }); + } +} + +@Injectable() +export class DocGrantModel extends BaseModel { + @Transactional() + async setOwner(workspaceId: string, docId: string, userId: string) { + await this.models.permissionProjection.markNewWriteOrigin(); + await this.models.permissionProjection.lockDocOwnerTransfer( + workspaceId, + docId + ); + await this.db.docGrant.updateMany({ + where: { + workspaceId, + docId, + principalType: 'user', + role: 'owner', + principalId: { not: userId }, + }, + data: { role: 'manager' }, + }); + + return await this.set(workspaceId, docId, userId, DocRole.Owner); + } + + @Transactional() + async set(workspaceId: string, docId: string, userId: string, role: DocRole) { + await this.models.permissionProjection.markNewWriteOrigin(); + assert(role !== DocRole.None && role !== DocRole.External); + + return await this.db.docGrant.upsert({ + where: { + workspaceId_docId_principalType_principalId: { + workspaceId, + docId, + principalType: 'user', + principalId: userId, + }, + }, + update: { + role: docRoleToNew(role), + }, + create: { + workspaceId, + docId, + principalType: 'user', + principalId: userId, + role: docRoleToNew(role), + }, + }); + } + + @Transactional() + async batchSetUserRoles( + workspaceId: string, + docId: string, + userIds: string[], + role: DocRole + ) { + await this.models.permissionProjection.markNewWriteOrigin(); + if (role === DocRole.Owner) { + throw new CanNotBatchGrantDocOwnerPermissions(); + } + if (userIds.length === 0) { + return 0; + } + + const grantRole = docRoleToNew(role); + for (const userId of userIds) { + await this.db.docGrant.upsert({ + where: { + workspaceId_docId_principalType_principalId: { + workspaceId, + docId, + principalType: 'user', + principalId: userId, + }, + }, + update: { + role: grantRole, + }, + create: { + workspaceId, + docId, + principalType: 'user', + principalId: userId, + role: grantRole, + }, + }); + } + return userIds.length; + } + + @Transactional() + async delete(workspaceId: string, docId: string, userId: string) { + await this.models.permissionProjection.markNewWriteOrigin(); + await this.db.$queryRaw` + SELECT 1 + FROM doc_grants + WHERE workspace_id = ${workspaceId} + AND doc_id = ${docId} + AND principal_type = 'user' + AND role = 'owner' + FOR UPDATE + `; + const deletingOwner = await this.db.docGrant.count({ + where: { + workspaceId, + docId, + principalType: 'user', + principalId: userId, + role: 'owner', + }, + }); + const otherOwners = await this.db.docGrant.count({ + where: { + workspaceId, + docId, + principalType: 'user', + principalId: { not: userId }, + role: 'owner', + }, + }); + if (deletingOwner > 0 && otherOwners === 0) { + throw new Error('Cannot remove the last doc owner grant.'); + } + + return await this.db.docGrant.deleteMany({ + where: { + workspaceId, + docId, + principalType: 'user', + principalId: userId, + }, + }); + } +} diff --git a/packages/backend/server/src/models/workspace-runtime-state.ts b/packages/backend/server/src/models/workspace-runtime-state.ts new file mode 100644 index 0000000000..856b4632b1 --- /dev/null +++ b/packages/backend/server/src/models/workspace-runtime-state.ts @@ -0,0 +1,223 @@ +import { Injectable } from '@nestjs/common'; + +import { BaseModel } from './base'; + +export type WorkspaceRuntimeState = { + workspaceId: string; + known: boolean; + stale: boolean; + readonly: boolean; + readonlyReasons: string[]; + updatedAt: Date | null; + lastReconciledAt: Date | null; + staleAfter: Date | null; +}; + +type WorkspaceRuntimeStateRow = { + workspaceId: string; + known: boolean; + readonly: boolean; + readonlyReasons: string[]; + updatedAt: Date; + lastReconciledAt: Date | null; + staleAfter: Date | null; +}; + +type LegacyWorkspaceRuntimeStateRow = { + workspaceId: string; + readonly: boolean; + readonlyReasons: string[]; + updatedAt: Date; + staleAt: Date | null; +}; + +function isMissingRuntimeStateColumn(error: unknown) { + const meta = (error as { meta?: { code?: string } })?.meta; + return meta?.code === '42703'; +} + +@Injectable() +export class WorkspaceRuntimeStateModel extends BaseModel { + private hasCurrentColumns?: Promise; + + async get(workspaceId: string): Promise { + const rows = await this.loadRows(workspaceId); + const row = rows[0]; + + if (!row) { + return { + workspaceId, + known: false, + stale: true, + readonly: false, + readonlyReasons: [], + updatedAt: null, + lastReconciledAt: null, + staleAfter: null, + }; + } + + return { + workspaceId, + known: row.known, + stale: + !row.known || (row.staleAfter !== null && row.staleAfter <= new Date()), + readonly: row.readonly, + readonlyReasons: row.readonlyReasons, + updatedAt: row.updatedAt, + lastReconciledAt: row.lastReconciledAt, + staleAfter: row.staleAfter, + }; + } + + async upsert( + workspaceId: string, + state: { + readonly: boolean; + readonlyReasons: string[]; + known?: boolean; + lastReconciledAt?: Date | null; + staleAfter?: Date | null; + } + ) { + if (await this.supportsCurrentRuntimeStateColumns()) { + await this.upsertCurrent(workspaceId, state); + } else { + await this.upsertLegacy(workspaceId, state); + } + } + + private async loadRows(workspaceId: string) { + if (!(await this.supportsCurrentRuntimeStateColumns())) { + return await this.loadLegacyRows(workspaceId); + } + + try { + return await this.db.$queryRaw` + SELECT + workspace_id AS "workspaceId", + known, + readonly, + readonly_reasons AS "readonlyReasons", + updated_at AS "updatedAt", + last_reconciled_at AS "lastReconciledAt", + stale_after AS "staleAfter" + FROM workspace_runtime_states + WHERE workspace_id = ${workspaceId} + LIMIT 1 + `; + } catch (error) { + if (!isMissingRuntimeStateColumn(error)) { + throw error; + } + return await this.loadLegacyRows(workspaceId); + } + } + + private async supportsCurrentRuntimeStateColumns() { + this.hasCurrentColumns ??= this.db.$queryRaw>` + SELECT EXISTS ( + SELECT 1 + FROM information_schema.columns + WHERE table_name = 'workspace_runtime_states' + AND column_name = 'known' + ) AS "exists" + `.then(rows => rows[0]?.exists ?? false); + return await this.hasCurrentColumns; + } + + private async loadLegacyRows(workspaceId: string) { + const rows = await this.db.$queryRaw` + SELECT + workspace_id AS "workspaceId", + readonly, + readonly_reasons AS "readonlyReasons", + updated_at AS "updatedAt", + stale_at AS "staleAt" + FROM workspace_runtime_states + WHERE workspace_id = ${workspaceId} + LIMIT 1 + `; + return rows.map(row => ({ + workspaceId: row.workspaceId, + known: true, + readonly: row.readonly, + readonlyReasons: row.readonlyReasons, + updatedAt: row.updatedAt, + lastReconciledAt: row.updatedAt, + staleAfter: row.staleAt, + })); + } + + private async upsertCurrent( + workspaceId: string, + state: { + readonly: boolean; + readonlyReasons: string[]; + known?: boolean; + lastReconciledAt?: Date | null; + staleAfter?: Date | null; + } + ) { + await this.db.$executeRaw` + INSERT INTO workspace_runtime_states ( + workspace_id, + known, + readonly, + readonly_reasons, + last_reconciled_at, + stale_after, + updated_at + ) + VALUES ( + ${workspaceId}, + ${state.known ?? true}, + ${state.readonly}, + ${state.readonlyReasons}, + ${state.lastReconciledAt ?? new Date()}, + ${state.staleAfter ?? null}, + now() + ) + ON CONFLICT (workspace_id) + DO UPDATE SET + known = EXCLUDED.known, + readonly = EXCLUDED.readonly, + readonly_reasons = EXCLUDED.readonly_reasons, + last_reconciled_at = EXCLUDED.last_reconciled_at, + stale_after = EXCLUDED.stale_after, + updated_at = now() + `; + } + + private async upsertLegacy( + workspaceId: string, + state: { + readonly: boolean; + readonlyReasons: string[]; + staleAfter?: Date | null; + } + ) { + await this.db.$executeRaw` + INSERT INTO workspace_runtime_states ( + workspace_id, + readonly, + readonly_reasons, + stale_at, + updated_at + ) + VALUES ( + ${workspaceId}, + ${state.readonly}, + ${state.readonlyReasons}, + ${state.staleAfter ?? null}, + now() + ) + ON CONFLICT (workspace_id) + DO UPDATE SET + readonly = EXCLUDED.readonly, + readonly_reasons = EXCLUDED.readonly_reasons, + stale_at = EXCLUDED.stale_at, + updated_at = now() + `; + } +} diff --git a/packages/backend/server/src/models/workspace-user-compat.ts b/packages/backend/server/src/models/workspace-user-compat.ts new file mode 100644 index 0000000000..7cb91054ae --- /dev/null +++ b/packages/backend/server/src/models/workspace-user-compat.ts @@ -0,0 +1,389 @@ +import { + Prisma, + PrismaClient, + User, + WorkspaceInvitation, + WorkspaceMember, + WorkspaceMemberSource, + WorkspaceMemberStatus, + WorkspaceUserRole, +} from '@prisma/client'; +import { groupBy } from 'lodash-es'; + +import { WorkspaceRole, workspaceUserSelect } from './common'; +import { + workspaceRoleFromNew, + workspaceSourceFromNew, + workspaceStatusFromNew, +} from './permission-write'; + +export type WorkspaceUserCompat = WorkspaceUserRole & { + user?: Pick; +}; + +export type WorkspaceMemberWithUser = WorkspaceMember & { + user?: Pick; +}; + +export type WorkspaceInvitationWithUser = WorkspaceInvitation & { + inviteeUser?: Pick | null; +}; + +type WorkspaceUserCompatRow = { + id: string; + workspaceId: string; + userId: string; + type: number; + status: WorkspaceMemberStatus; + source: WorkspaceMemberSource; + inviterId: string | null; + createdAt: Date; + updatedAt: Date; + userName: string; + userEmail: string; + userAvatarUrl: string | null; +}; + +type WorkspaceUserCompatDb = Pick< + PrismaClient, + 'workspaceMember' | 'workspaceInvitation' | '$queryRaw' +>; + +export function workspaceMemberToCompat( + member: WorkspaceMemberWithUser +): WorkspaceUserCompat { + return { + id: member.legacyPermissionId ?? member.id, + workspaceId: member.workspaceId, + userId: member.userId, + type: workspaceRoleFromNew(member.role as never), + status: WorkspaceMemberStatus.Accepted, + source: workspaceSourceFromNew(member.source as never), + inviterId: null, + createdAt: member.createdAt, + updatedAt: member.updatedAt, + ...(member.user ? { user: member.user } : {}), + }; +} + +export function workspaceInvitationToCompat( + invitation: WorkspaceInvitationWithUser +): WorkspaceUserCompat { + return { + id: invitation.legacyPermissionId ?? invitation.id, + workspaceId: invitation.workspaceId, + userId: invitation.inviteeUserId ?? '', + type: workspaceRoleFromNew(invitation.requestedRole as never), + status: workspaceStatusFromNew(invitation.status as never), + source: workspaceSourceFromNew(invitation.kind as never), + inviterId: invitation.inviterUserId, + createdAt: invitation.createdAt, + updatedAt: invitation.updatedAt, + ...(invitation.inviteeUser ? { user: invitation.inviteeUser } : {}), + }; +} + +function rawCompatRowToCompat( + row: WorkspaceUserCompatRow +): WorkspaceUserCompat { + return { + id: row.id, + workspaceId: row.workspaceId, + userId: row.userId, + type: row.type, + status: row.status, + source: row.source, + inviterId: row.inviterId, + createdAt: row.createdAt, + updatedAt: row.updatedAt, + user: { + id: row.userId, + name: row.userName, + email: row.userEmail, + avatarUrl: row.userAvatarUrl, + }, + }; +} + +export async function queryCompatRows( + db: WorkspaceUserCompatDb, + workspaceId: string, + pagination: { first: number; offset: number; after?: string | Date } +) { + const after = pagination.after + ? Prisma.sql`AND created_at >= ${pagination.after}::timestamptz` + : Prisma.empty; + const rows = await db.$queryRaw` + SELECT * + FROM ( + SELECT + COALESCE(wm.legacy_permission_id, wm.id) AS id, + wm.workspace_id AS "workspaceId", + wm.user_id AS "userId", + CASE wm.role + WHEN 'owner' THEN ${WorkspaceRole.Owner} + WHEN 'admin' THEN ${WorkspaceRole.Admin} + ELSE ${WorkspaceRole.Collaborator} + END AS type, + 'Accepted'::"WorkspaceMemberStatus" AS status, + CASE wm.source + WHEN 'link' THEN 'Link'::"WorkspaceMemberSource" + ELSE 'Email'::"WorkspaceMemberSource" + END AS source, + NULL::varchar AS "inviterId", + wm.created_at AS "createdAt", + wm.updated_at AS "updatedAt", + u.name AS "userName", + u.email AS "userEmail", + u.avatar_url AS "userAvatarUrl", + wm.created_at AS created_at + FROM workspace_members wm + INNER JOIN users u ON u.id = wm.user_id + WHERE wm.workspace_id = ${workspaceId} + AND wm.state = 'active' + UNION ALL + SELECT + COALESCE(wi.legacy_permission_id, wi.id) AS id, + wi.workspace_id AS "workspaceId", + wi.invitee_user_id AS "userId", + CASE wi.requested_role + WHEN 'admin' THEN ${WorkspaceRole.Admin} + ELSE ${WorkspaceRole.Collaborator} + END AS type, + CASE wi.status + WHEN 'waiting_review' THEN 'UnderReview'::"WorkspaceMemberStatus" + WHEN 'waiting_seat' THEN 'NeedMoreSeat'::"WorkspaceMemberStatus" + ELSE 'Pending'::"WorkspaceMemberStatus" + END AS status, + CASE wi.kind + WHEN 'link' THEN 'Link'::"WorkspaceMemberSource" + ELSE 'Email'::"WorkspaceMemberSource" + END AS source, + wi.inviter_user_id AS "inviterId", + wi.created_at AS "createdAt", + wi.updated_at AS "updatedAt", + u.name AS "userName", + u.email AS "userEmail", + u.avatar_url AS "userAvatarUrl", + wi.created_at AS created_at + FROM workspace_invitations wi + INNER JOIN users u ON u.id = wi.invitee_user_id + WHERE wi.workspace_id = ${workspaceId} + ) roles + WHERE true + ${after} + ORDER BY created_at ASC + OFFSET ${pagination.offset} + LIMIT ${pagination.first} + `; + return rows.map(row => rawCompatRowToCompat(row)); +} + +export async function searchCompatRows( + db: WorkspaceUserCompatDb, + workspaceId: string, + query: string, + pagination: { first: number; offset: number; after?: string | Date } +) { + const after = pagination.after + ? Prisma.sql`AND wm.created_at >= ${pagination.after}::timestamptz` + : Prisma.empty; + const rows = await db.$queryRaw` + SELECT + COALESCE(wm.legacy_permission_id, wm.id) AS id, + wm.workspace_id AS "workspaceId", + wm.user_id AS "userId", + CASE wm.role + WHEN 'owner' THEN ${WorkspaceRole.Owner} + WHEN 'admin' THEN ${WorkspaceRole.Admin} + ELSE ${WorkspaceRole.Collaborator} + END AS type, + 'Accepted'::"WorkspaceMemberStatus" AS status, + CASE wm.source + WHEN 'link' THEN 'Link'::"WorkspaceMemberSource" + ELSE 'Email'::"WorkspaceMemberSource" + END AS source, + NULL::varchar AS "inviterId", + wm.created_at AS "createdAt", + wm.updated_at AS "updatedAt", + u.name AS "userName", + u.email AS "userEmail", + u.avatar_url AS "userAvatarUrl" + FROM workspace_members wm + INNER JOIN users u ON u.id = wm.user_id + WHERE wm.workspace_id = ${workspaceId} + AND wm.state = 'active' + AND (u.email ILIKE ${`%${query}%`} OR u.name ILIKE ${`%${query}%`}) + ${after} + ORDER BY wm.created_at ASC + OFFSET ${pagination.offset} + LIMIT ${pagination.first} + `; + return rows.map(row => rawCompatRowToCompat(row)); +} + +export async function countWorkspaceUsers( + db: WorkspaceUserCompatDb, + workspaceId: string +) { + const [members, invitations] = await Promise.all([ + db.workspaceMember.count({ + where: { workspaceId, state: 'active' }, + }), + db.workspaceInvitation.count({ where: { workspaceId } }), + ]); + return members + invitations; +} + +export async function countChargedWorkspaceUsers( + db: WorkspaceUserCompatDb, + workspaceId: string +) { + const [members, invitations] = await Promise.all([ + db.workspaceMember.count({ + where: { workspaceId, state: 'active' }, + }), + db.workspaceInvitation.count({ + where: { + workspaceId, + status: { + not: 'waiting_review', + }, + }, + }), + ]); + return members + invitations; +} + +export async function findUserActiveWorkspaceRoles( + db: WorkspaceUserCompatDb, + userId: string, + filter: { role?: WorkspaceRole } = {} +) { + const roles = await db.workspaceMember.findMany({ + where: { + userId, + state: 'active', + role: filter.role ? workspaceRoleToNewFilter(filter.role) : undefined, + }, + }); + return roles.map(role => workspaceMemberToCompat(role)); +} + +export async function hasSharedWorkspace( + db: WorkspaceUserCompatDb, + userId: string, + otherUserId: string +) { + if (userId === otherUserId) { + return true; + } + + const shared = await db.$queryRaw<{ id: string }[]>` + SELECT mine.id + FROM workspace_members mine + INNER JOIN workspace_members other + ON other.workspace_id = mine.workspace_id + AND other.user_id = ${otherUserId} + AND other.state = 'active' + WHERE mine.user_id = ${userId} + AND mine.state = 'active' + LIMIT 1 + `; + + return shared.length > 0; +} + +export async function allocateWorkspaceSeats( + db: WorkspaceUserCompatDb, + models: { + permissionProjection: { markNewWriteOrigin(): Promise }; + workspaceMember: { + setActive( + workspaceId: string, + userId: string, + role: WorkspaceRole + ): Promise; + }; + }, + workspaceId: string, + limit: number +) { + await models.permissionProjection.markNewWriteOrigin(); + const [activeCount, pendingCount] = await Promise.all([ + db.workspaceMember.count({ + where: { + workspaceId, + state: 'active', + }, + }), + db.workspaceInvitation.count({ + where: { + workspaceId, + status: 'pending', + }, + }), + ]); + const usedCount = activeCount + pendingCount; + + if (limit <= usedCount) { + return []; + } + + const invitationsToAllocate = await db.workspaceInvitation.findMany({ + where: { + workspaceId, + status: 'waiting_seat', + inviteeUserId: { + not: null, + }, + }, + orderBy: { createdAt: 'asc' }, + take: limit - usedCount, + }); + + const groups = groupBy(invitationsToAllocate, invitation => + workspaceSourceFromNew(invitation.kind as never) + ) as Record; + + if (groups.Email?.length > 0) { + await db.workspaceInvitation.updateMany({ + where: { id: { in: groups.Email.map(invitation => invitation.id) } }, + data: { status: 'pending' }, + }); + } + + if (groups.Link?.length > 0) { + await Promise.all( + groups.Link.map(invitation => + models.workspaceMember.setActive( + invitation.workspaceId, + invitation.inviteeUserId as string, + invitation.requestedRole === 'admin' + ? WorkspaceRole.Admin + : WorkspaceRole.Collaborator + ) + ) + ); + } + + return (groups.Email ?? []).map(invitation => + workspaceInvitationToCompat({ + ...invitation, + status: 'pending', + }) + ); +} + +export function workspaceRoleToNewFilter(role: WorkspaceRole) { + switch (role) { + case WorkspaceRole.Owner: + return 'owner'; + case WorkspaceRole.Admin: + return 'admin'; + case WorkspaceRole.Collaborator: + return 'member'; + default: + return undefined; + } +} diff --git a/packages/backend/server/src/models/workspace-user.ts b/packages/backend/server/src/models/workspace-user.ts index b46632cfc0..132e0179de 100644 --- a/packages/backend/server/src/models/workspace-user.ts +++ b/packages/backend/server/src/models/workspace-user.ts @@ -5,11 +5,21 @@ import { WorkspaceMemberStatus, WorkspaceUserRole, } from '@prisma/client'; -import { groupBy } from 'lodash-es'; import { EventBus, NewOwnerIsNotActiveMember, PaginationInput } from '../base'; import { BaseModel } from './base'; import { WorkspaceRole, workspaceUserSelect } from './common'; +import { + allocateWorkspaceSeats, + countChargedWorkspaceUsers, + countWorkspaceUsers, + findUserActiveWorkspaceRoles, + hasSharedWorkspace, + queryCompatRows, + searchCompatRows, + workspaceInvitationToCompat, + workspaceMemberToCompat, +} from './workspace-user-compat'; export { WorkspaceMemberStatus }; @@ -41,68 +51,50 @@ export class WorkspaceUserModel extends BaseModel { */ @Transactional() async setOwner(workspaceId: string, userId: string) { - const oldOwner = await this.db.workspaceUserRole.findFirst({ + const oldOwner = await this.db.workspaceMember.findFirst({ + include: { + user: { + select: workspaceUserSelect, + }, + }, where: { workspaceId, - type: WorkspaceRole.Owner, + role: 'owner', + state: 'active', }, }); + const fallbackRole = (await this.models.workspace.isTeamWorkspace( + workspaceId + )) + ? WorkspaceRole.Admin + : WorkspaceRole.Collaborator; - // If there is already an owner, we need to change the old owner to admin - if (oldOwner) { - const newOwnerOldRole = await this.db.workspaceUserRole.findFirst({ - where: { - workspaceId, - userId, - }, - }); - + try { + await this.models.workspaceMember.setOwner( + workspaceId, + userId, + fallbackRole + ); + } catch (error) { if ( - !newOwnerOldRole || - newOwnerOldRole.status !== WorkspaceMemberStatus.Accepted + error instanceof Error && + error.message === 'New workspace owner must be an active member.' ) { throw new NewOwnerIsNotActiveMember(); } + throw error; + } - const fallbackRole = (await this.models.workspace.isTeamWorkspace( - workspaceId - )) - ? WorkspaceRole.Admin - : WorkspaceRole.Collaborator; - - await this.db.workspaceUserRole.update({ - where: { - id: oldOwner.id, - }, - data: { - type: fallbackRole, - }, - }); - await this.db.workspaceUserRole.update({ - where: { - id: newOwnerOldRole.id, - }, - data: { - type: WorkspaceRole.Owner, - }, - }); + if (oldOwner?.user && oldOwner.user.id !== userId) { this.event.emit('workspace.owner.changed', { workspaceId, - from: oldOwner.userId, + from: oldOwner.user.id, to: userId, }); this.logger.log( - `Transfer workspace owner of [${workspaceId}] from [${oldOwner.userId}] to [${userId}]` + `Transfer workspace owner of [${workspaceId}] from [${oldOwner.user.id}] to [${userId}]` ); } else { - await this.db.workspaceUserRole.create({ - data: { - workspaceId, - userId, - type: WorkspaceRole.Owner, - status: WorkspaceMemberStatus.Accepted, - }, - }); this.logger.log(`Set workspace owner of [${workspaceId}] to [${userId}]`); } } @@ -123,21 +115,33 @@ export class WorkspaceUserModel extends BaseModel { inviterId?: string; } = {} ) { - if (role === WorkspaceRole.Owner) { - throw new Error('Cannot grant Owner role of a workspace to a user.'); - } - const oldRole = await this.get(workspaceId, userId); + if (role === WorkspaceRole.External) { + return await this.setExternal(workspaceId, userId, oldRole); + } + if (oldRole) { if (oldRole.type === role) { return oldRole; } - const newRole = await this.db.workspaceUserRole.update({ - where: { id: oldRole.id }, - data: { type: role }, - }); + const status = defaultData.status ?? oldRole.status; + if (status === WorkspaceMemberStatus.Accepted) { + await this.models.workspaceMember.setActive(workspaceId, userId, role); + } else { + await this.models.workspaceInvitation.set( + workspaceId, + userId, + role, + status, + { + source: defaultData.source ?? oldRole.source, + inviterId: defaultData.inviterId ?? oldRole.inviterId ?? undefined, + } + ); + } + const newRole = await this.mustGet(workspaceId, userId); if (oldRole.status === WorkspaceMemberStatus.Accepted) { this.event.emit('workspace.members.roleChanged', { @@ -155,17 +159,60 @@ export class WorkspaceUserModel extends BaseModel { inviterId, } = defaultData; - return await this.db.workspaceUserRole.create({ + if (status === WorkspaceMemberStatus.Accepted) { + await this.models.workspaceMember.setActive(workspaceId, userId, role); + } else { + await this.models.workspaceInvitation.set( + workspaceId, + userId, + role, + status, + { + source, + inviterId, + } + ); + } + return await this.mustGet(workspaceId, userId); + } + } + + private async setExternal( + workspaceId: string, + userId: string, + oldRole: WorkspaceUserRole | null + ) { + await this.models.permissionProjection.markLegacyWriteOrigin(); + await this.db.workspaceMember.deleteMany({ + where: { workspaceId, userId, state: 'active' }, + }); + await this.db.workspaceInvitation.deleteMany({ + where: { workspaceId, inviteeUserId: userId }, + }); + + await this.models.permissionProjection.markNewWriteOrigin(); + if (oldRole) { + return await this.withPermissionProjectionMetric( + this.db.workspaceUserRole.update({ + where: { id: oldRole.id }, + data: { + type: WorkspaceRole.External, + status: WorkspaceMemberStatus.Accepted, + }, + }) + ); + } + + return await this.withPermissionProjectionMetric( + this.db.workspaceUserRole.create({ data: { workspaceId, userId, - type: role, - status, - source, - inviterId, + type: WorkspaceRole.External, + status: WorkspaceMemberStatus.Accepted, }, - }); - } + }) + ); } async setStatus( @@ -177,68 +224,129 @@ export class WorkspaceUserModel extends BaseModel { } = {} ) { const { inviterId } = data; - return await this.db.workspaceUserRole.update({ - where: { - workspaceId_userId: { + await this.models.workspaceInvitation.setState( + workspaceId, + userId, + status, + { + inviterId, + } + ); + return await this.mustGet(workspaceId, userId); + } + + private async mustGet(workspaceId: string, userId: string) { + const role = await this.get(workspaceId, userId); + if (!role) { + throw new Error( + `Workspace permission ${workspaceId}/${userId} not found after write.` + ); + } + return role; + } + + @Transactional() + async delete(workspaceId: string, userId: string) { + await this.models.workspaceMember.delete(workspaceId, userId); + await this.db.workspaceInvitation.deleteMany({ + where: { workspaceId, inviteeUserId: userId }, + }); + await this.withPermissionProjectionMetric( + this.db.workspaceUserRole.deleteMany({ + where: { workspaceId, userId, }, - }, - data: { - status, - inviterId, - }, - }); - } - - async delete(workspaceId: string, userId: string) { - await this.db.workspaceUserRole.deleteMany({ - where: { - workspaceId, - userId, - }, - }); + }) + ); } + @Transactional() async deleteByUserId(userId: string) { - await this.db.workspaceUserRole.deleteMany({ - where: { - userId, - }, + await this.models.permissionProjection.markNewWriteOrigin(); + await this.db.workspaceMember.deleteMany({ + where: { userId }, }); + await this.db.workspaceInvitation.deleteMany({ + where: { inviteeUserId: userId }, + }); + await this.withPermissionProjectionMetric( + this.db.workspaceUserRole.deleteMany({ + where: { + userId, + }, + }) + ); } async deleteNonAccepted(workspaceId: string) { - return await this.db.workspaceUserRole.deleteMany({ - where: { workspaceId, status: { not: WorkspaceMemberStatus.Accepted } }, - }); + return await this.models.workspaceInvitation.deleteNonAccepted(workspaceId); } + @Transactional() async demoteAcceptedAdmins(workspaceId: string) { - return await this.db.workspaceUserRole.updateMany({ - where: { - workspaceId, - status: WorkspaceMemberStatus.Accepted, - type: WorkspaceRole.Admin, - }, - data: { - type: WorkspaceRole.Collaborator, - }, + await this.models.permissionProjection.markNewWriteOrigin(); + return await this.db.workspaceMember.updateMany({ + where: { workspaceId, role: 'admin', state: 'active' }, + data: { role: 'member' }, }); } async get(workspaceId: string, userId: string) { - return await this.db.workspaceUserRole.findUnique({ + const active = await this.db.workspaceMember.findFirst({ where: { - workspaceId_userId: { + workspaceId, + userId, + state: 'active', + }, + }); + if (active) { + return workspaceMemberToCompat(active); + } + + const invitation = await this.db.workspaceInvitation.findUnique({ + where: { + workspaceId_inviteeUserId: { workspaceId, - userId, + inviteeUserId: userId, }, }, }); + if (invitation) { + return workspaceInvitationToCompat(invitation); + } + + return await this.db.workspaceUserRole.findFirst({ + where: { + workspaceId, + userId, + type: WorkspaceRole.External, + }, + }); } async getById(id: string) { + const member = await this.db.workspaceMember.findFirst({ + where: { + OR: [{ id }, { legacyPermissionId: id }], + }, + }); + if (member) { + return workspaceMemberToCompat(member); + } + + const invitation = await this.db.workspaceInvitation.findFirst({ + where: { + OR: [{ id }, { legacyPermissionId: id }], + inviteeUserId: { + not: null, + }, + }, + }); + if (invitation) { + return workspaceInvitationToCompat(invitation); + } + return await this.db.workspaceUserRole.findUnique({ where: { id }, }); @@ -248,16 +356,18 @@ export class WorkspaceUserModel extends BaseModel { * Get the **accepted** Role of a user in a workspace. */ async getActive(workspaceId: string, userId: string) { - return await this.db.workspaceUserRole.findUnique({ + const active = await this.db.workspaceMember.findFirst({ where: { - workspaceId_userId: { workspaceId, userId }, - status: WorkspaceMemberStatus.Accepted, + workspaceId, + userId, + state: 'active', }, }); + return active ? workspaceMemberToCompat(active) : null; } async getOwner(workspaceId: string) { - const role = await this.db.workspaceUserRole.findFirst({ + const role = await this.db.workspaceMember.findFirst({ include: { user: { select: workspaceUserSelect, @@ -265,11 +375,12 @@ export class WorkspaceUserModel extends BaseModel { }, where: { workspaceId, - type: WorkspaceRole.Owner, + role: 'owner', + state: 'active', }, }); - if (!role) { + if (!role?.user) { throw new Error('Workspace owner not found'); } @@ -277,7 +388,7 @@ export class WorkspaceUserModel extends BaseModel { } async getAdmins(workspaceId: string) { - const list = await this.db.workspaceUserRole.findMany({ + const list = await this.db.workspaceMember.findMany({ include: { user: { select: workspaceUserSelect, @@ -285,8 +396,8 @@ export class WorkspaceUserModel extends BaseModel { }, where: { workspaceId, - type: WorkspaceRole.Admin, - status: WorkspaceMemberStatus.Accepted, + role: 'admin', + state: 'active', }, }); @@ -294,87 +405,34 @@ export class WorkspaceUserModel extends BaseModel { } async count(workspaceId: string) { - return this.db.workspaceUserRole.count({ - where: { - workspaceId, - }, - }); + return await countWorkspaceUsers(this.db, workspaceId); } /** * Get the number of users those in the status should be charged in billing system in a workspace. */ async chargedCount(workspaceId: string) { - return this.db.workspaceUserRole.count({ - where: { - workspaceId, - status: { - not: WorkspaceMemberStatus.UnderReview, - }, - }, - }); + return await countChargedWorkspaceUsers(this.db, workspaceId); } async getUserActiveRoles( userId: string, filter: { role?: WorkspaceRole } = {} ) { - return await this.db.workspaceUserRole.findMany({ - where: { - userId, - status: WorkspaceMemberStatus.Accepted, - type: filter.role, - }, - }); + return await findUserActiveWorkspaceRoles(this.db, userId, filter); } async hasSharedWorkspace(userId: string, otherUserId: string) { - if (userId === otherUserId) { - return true; - } - - const shared = await this.db.workspaceUserRole.findFirst({ - select: { id: true }, - where: { - userId, - status: WorkspaceMemberStatus.Accepted, - workspace: { - permissions: { - some: { - userId: otherUserId, - }, - }, - }, - }, - }); - - return !!shared; + return await hasSharedWorkspace(this.db, userId, otherUserId); } async paginate(workspaceId: string, pagination: PaginationInput) { - return await Promise.all([ - this.db.workspaceUserRole.findMany({ - include: { - user: { - select: workspaceUserSelect, - }, - }, - where: { - workspaceId, - createdAt: pagination.after - ? { - gte: pagination.after, - } - : undefined, - }, - orderBy: { - createdAt: 'asc', - }, - take: pagination.first, - skip: pagination.offset + (pagination.after ? 1 : 0), - }), - this.count(workspaceId), - ]); + const rows = await queryCompatRows(this.db, workspaceId, { + first: pagination.first, + offset: pagination.offset + (pagination.after ? 1 : 0), + after: pagination.after ?? undefined, + }); + return [rows, await this.count(workspaceId)] as const; } async search( @@ -382,91 +440,20 @@ export class WorkspaceUserModel extends BaseModel { query: string, pagination: PaginationInput ) { - return await this.db.workspaceUserRole.findMany({ - include: { user: { select: workspaceUserSelect } }, - where: { - workspaceId, - status: WorkspaceMemberStatus.Accepted, - user: { - OR: [ - { - email: { - contains: query, - mode: 'insensitive', - }, - }, - { - name: { - contains: query, - mode: 'insensitive', - }, - }, - ], - }, - }, - orderBy: { createdAt: 'asc' }, - take: pagination.first, - skip: pagination.offset + (pagination.after ? 1 : 0), + return await searchCompatRows(this.db, workspaceId, query, { + first: pagination.first, + offset: pagination.offset + (pagination.after ? 1 : 0), + after: pagination.after ?? undefined, }); } @Transactional() async allocateSeats(workspaceId: string, limit: number) { - const usedCount = await this.db.workspaceUserRole.count({ - where: { - workspaceId, - status: { - in: [WorkspaceMemberStatus.Accepted, WorkspaceMemberStatus.Pending], - }, - }, - }); - - if (limit <= usedCount) { - return []; - } - - const membersToBeAllocated = await this.db.workspaceUserRole.findMany({ - where: { - workspaceId, - status: { - in: [ - WorkspaceMemberStatus.AllocatingSeat, - WorkspaceMemberStatus.NeedMoreSeat, - ], - }, - }, - orderBy: { createdAt: 'asc' }, - take: limit - usedCount, - }); - - const groups = groupBy( - membersToBeAllocated, - member => member.source - ) as Record; - - if (groups.Email?.length > 0) { - await this.db.workspaceUserRole.updateMany({ - where: { id: { in: groups.Email.map(m => m.id) } }, - data: { status: WorkspaceMemberStatus.Pending }, - }); - } - - if (groups.Link?.length > 0) { - await this.db.workspaceUserRole.updateMany({ - where: { id: { in: groups.Link.map(m => m.id) } }, - data: { status: WorkspaceMemberStatus.Accepted }, - }); - } - - // after allocating, all rests should be `NeedMoreSeat` - await this.db.workspaceUserRole.updateMany({ - where: { - workspaceId, - status: WorkspaceMemberStatus.AllocatingSeat, - }, - data: { status: WorkspaceMemberStatus.NeedMoreSeat }, - }); - - return groups.Email ?? []; + return await allocateWorkspaceSeats( + this.db, + this.models, + workspaceId, + limit + ); } } diff --git a/packages/backend/server/src/models/workspace.ts b/packages/backend/server/src/models/workspace.ts index 46d0e23a55..084d5de7ad 100644 --- a/packages/backend/server/src/models/workspace.ts +++ b/packages/backend/server/src/models/workspace.ts @@ -90,9 +90,11 @@ export class WorkspaceModel extends BaseModel { */ @Transactional() async create(userId: string) { - const workspace = await this.db.workspace.create({ - data: { public: false }, - }); + const workspace = await this.withPermissionProjectionMetric( + this.db.workspace.create({ + data: { public: false }, + }) + ); this.logger.log(`Workspace created with id ${workspace.id}`); await this.models.workspaceUser.setOwner(workspace.id, userId); return workspace; @@ -106,12 +108,26 @@ export class WorkspaceModel extends BaseModel { data: UpdateWorkspaceInput, notifyUpdate = true ) { - const workspace = await this.db.workspace.update({ - where: { - id: workspaceId, - }, - data, - }); + if ( + data.public !== undefined || + data.enableSharing !== undefined || + data.enableUrlPreview !== undefined + ) { + await this.models.workspaceAccessPolicy.upsert(workspaceId, { + public: data.public, + enableSharing: data.enableSharing, + enableUrlPreview: data.enableUrlPreview, + }); + } + + const workspace = await this.withPermissionProjectionMetric( + this.db.workspace.update({ + where: { + id: workspaceId, + }, + data, + }) + ); this.logger.debug( `Updated workspace ${workspaceId} with data ${JSON.stringify(data)}` ); @@ -155,11 +171,13 @@ export class WorkspaceModel extends BaseModel { } async delete(workspaceId: string) { - const rawResult = await this.db.workspace.deleteMany({ - where: { - id: workspaceId, - }, - }); + const rawResult = await this.withPermissionProjectionMetric( + this.db.workspace.deleteMany({ + where: { + id: workspaceId, + }, + }) + ); if (rawResult.count > 0) { this.event.emit('workspace.deleted', { id: workspaceId }); @@ -183,7 +201,23 @@ export class WorkspaceModel extends BaseModel { } async isTeamWorkspace(workspaceId: string) { - return this.models.workspaceFeature.has(workspaceId, 'team_plan_v1'); + const now = new Date(); + const count = await this.db.entitlement.count({ + where: { + targetType: 'workspace', + targetId: workspaceId, + plan: { in: ['team', 'selfhost_team'] }, + OR: [ + { + status: 'active', + OR: [{ expiresAt: null }, { expiresAt: { gt: now } }], + }, + { status: 'grace', graceUntil: { gt: now } }, + ], + }, + }); + + return count > 0; } // #endregion diff --git a/packages/backend/server/src/plugins/calendar/resolver.ts b/packages/backend/server/src/plugins/calendar/resolver.ts index 75ed602b52..5fed19033b 100644 --- a/packages/backend/server/src/plugins/calendar/resolver.ts +++ b/packages/backend/server/src/plugins/calendar/resolver.ts @@ -11,7 +11,7 @@ import { import { ActionForbidden, AuthenticationRequired, Config } from '../../base'; import { CurrentUser } from '../../core/auth'; import { ServerConfigType } from '../../core/config/types'; -import { AccessController } from '../../core/permission'; +import { PermissionAccess } from '../../core/permission'; import { UserType } from '../../core/user'; import { WorkspaceType } from '../../core/workspaces'; import { Models } from '../../models'; @@ -113,7 +113,7 @@ export class CalendarAccountResolver { export class WorkspaceCalendarResolver { constructor( private readonly calendar: CalendarService, - private readonly access: AccessController + private readonly access: PermissionAccess ) {} @ResolveField(() => [WorkspaceCalendarObjectType]) @@ -133,7 +133,7 @@ export class WorkspaceCalendarResolver { export class WorkspaceCalendarEventsResolver { constructor( private readonly calendar: CalendarService, - private readonly access: AccessController + private readonly access: PermissionAccess ) {} @ResolveField(() => [CalendarEventObjectType]) @@ -162,7 +162,7 @@ export class CalendarMutationResolver { private readonly calendar: CalendarService, private readonly oauth: CalendarOAuthService, private readonly models: Models, - private readonly access: AccessController + private readonly access: PermissionAccess ) {} @Mutation(() => String) diff --git a/packages/backend/server/src/plugins/copilot/byok/policy.ts b/packages/backend/server/src/plugins/copilot/byok/policy.ts index 45f78b6ab7..ef288a36fe 100644 --- a/packages/backend/server/src/plugins/copilot/byok/policy.ts +++ b/packages/backend/server/src/plugins/copilot/byok/policy.ts @@ -1,24 +1,24 @@ import { Injectable } from '@nestjs/common'; import { ActionForbidden } from '../../../base'; +import { QuotaStateService } from '../../../core/quota/state'; import { Models, WorkspaceRole } from '../../../models'; @Injectable() export class ByokEntitlementPolicy { - constructor(private readonly models: Models) {} - - private isUserPlanEntitled(features: string[]) { - return ( - features.includes('pro_plan_v1') || - features.includes('lifetime_pro_plan_v1') || - features.includes('unlimited_copilot') - ); - } + constructor( + private readonly models: Models, + private readonly quotaState: QuotaStateService + ) {} async hasAiPlan(userId?: string) { if (!userId) return false; - const features = await this.models.userFeature.list(userId); - return this.isUserPlanEntitled(features); + const state = await this.quotaState.reconcileUserQuotaState(userId); + const flags = state.flags as { unlimitedCopilot?: boolean }; + return ( + flags.unlimitedCopilot || + ['pro', 'lifetime_pro', 'ai'].includes(state.plan) + ); } async hasManagementAccess(workspaceId: string, userId?: string) { @@ -59,7 +59,7 @@ export class ByokEntitlementPolicy { async hasLocalEntitlement(workspaceId: string, userId?: string) { if (env.selfhosted) return true; - if (await this.models.workspaceFeature.has(workspaceId, 'team_plan_v1')) { + if (await this.hasWorkspaceTeamPlan(workspaceId)) { return true; } @@ -73,7 +73,7 @@ export class ByokEntitlementPolicy { async hasServerEntitlement(workspaceId: string) { if (env.selfhosted) return true; - if (await this.models.workspaceFeature.has(workspaceId, 'team_plan_v1')) { + if (await this.hasWorkspaceTeamPlan(workspaceId)) { return true; } @@ -102,4 +102,20 @@ export class ByokEntitlementPolicy { throw new ActionForbidden('BYOK requires Pro, Team, or Believer.'); } } + + private async hasWorkspaceTeamPlan(workspaceId: string) { + try { + const state = + await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); + return ['team', 'selfhost_team'].includes(state.plan); + } catch (error) { + if ( + error instanceof Error && + error.message === 'Workspace owner not found' + ) { + return false; + } + throw error; + } + } } diff --git a/packages/backend/server/src/plugins/copilot/byok/resolver.ts b/packages/backend/server/src/plugins/copilot/byok/resolver.ts index 65438cd1f1..ec2246bba9 100644 --- a/packages/backend/server/src/plugins/copilot/byok/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/byok/resolver.ts @@ -13,7 +13,7 @@ import { SafeIntResolver } from 'graphql-scalars'; import { Throttle } from '../../../base'; import { CurrentUser } from '../../../core/auth'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { WorkspaceType } from '../../../core/workspaces'; import { ByokEntitlementPolicy } from './policy'; import { ByokKeyConfig, ByokLocalLeaseProvider, ByokService } from './service'; @@ -259,7 +259,7 @@ class CreateWorkspaceByokLocalLeaseInput { @Resolver(() => WorkspaceType) export class WorkspaceByokResolver { constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly entitlement: ByokEntitlementPolicy, private readonly byok: ByokService ) {} diff --git a/packages/backend/server/src/plugins/copilot/context/realtime.ts b/packages/backend/server/src/plugins/copilot/context/realtime.ts index 23648f6219..84a7529208 100644 --- a/packages/backend/server/src/plugins/copilot/context/realtime.ts +++ b/packages/backend/server/src/plugins/copilot/context/realtime.ts @@ -2,7 +2,7 @@ import { Injectable, OnModuleInit } from '@nestjs/common'; import { z } from 'zod'; import { OnEvent } from '../../../base'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { RealtimePublisher, RealtimeRegistry, @@ -19,7 +19,7 @@ export function workspaceEmbeddingRoom(workspaceId: string) { @Injectable() export class CopilotEmbeddingRealtimeProvider implements OnModuleInit { constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly models: Models, private readonly context: CopilotContextService, private readonly registry: RealtimeRegistry, diff --git a/packages/backend/server/src/plugins/copilot/context/resolver.ts b/packages/backend/server/src/plugins/copilot/context/resolver.ts index 23caa4c79a..bef1b67c36 100644 --- a/packages/backend/server/src/plugins/copilot/context/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/context/resolver.ts @@ -37,10 +37,7 @@ import { UserFriendlyError, } from '../../../base'; import { CurrentUser } from '../../../core/auth'; -import { - AccessController, - WorkspacePolicyService, -} from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { ContextBlob, ContextCategories, @@ -60,7 +57,7 @@ import { getSignal, MAX_EMBEDDABLE_SIZE, readStream } from '../utils'; import { CopilotContextService } from './service'; async function assertAccess( - ac: AccessController, + ac: PermissionAccess, userId: string, workspaceId: string ) { @@ -73,7 +70,7 @@ async function assertAccess( async function getSession( context: CopilotContextService, - ac: AccessController, + ac: PermissionAccess, userId: string, contextId: string, options: { workspaceId?: string; sessionId?: string } = {} @@ -292,7 +289,7 @@ class ContextMatchedDocChunk implements DocChunkSimilarity { @Resolver(() => CopilotType) export class CopilotContextRootResolver { constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly event: EventBus, private readonly mutex: RequestMutex, private readonly chatSession: ChatSessionService, @@ -441,8 +438,7 @@ export class CopilotContextRootResolver { @Resolver(() => CopilotContextType) export class CopilotContextResolver { constructor( - private readonly ac: AccessController, - private readonly policy: WorkspacePolicyService, + private readonly ac: PermissionAccess, private readonly models: Models, private readonly mutex: RequestMutex, private readonly context: CopilotContextService, @@ -745,7 +741,11 @@ export class CopilotContextResolver { const blobId = createHash('sha256').update(buffer).digest('base64url'); const { filename, mimetype } = content; - await this.policy.assertCanUploadBlob(user.id, session.workspaceId); + await this.ac + .user(user.id) + .workspace(session.workspaceId) + .allowLocal() + .assert('Workspace.Blobs.Write'); await this.storage.put(user.id, session.workspaceId, blobId, buffer); const file = await session.addFile( blobId, diff --git a/packages/backend/server/src/plugins/copilot/conversation/inbox.ts b/packages/backend/server/src/plugins/copilot/conversation/inbox.ts index 1aa5a46ee9..c0df7e798c 100644 --- a/packages/backend/server/src/plugins/copilot/conversation/inbox.ts +++ b/packages/backend/server/src/plugins/copilot/conversation/inbox.ts @@ -7,7 +7,7 @@ import { ImageFormatNotSupported, sniffMime, } from '../../../base'; -import { WorkspacePolicyService } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { processImage } from '../../../native'; import { CompatSubmissionStore } from '../compat/submission-store'; import type { PromptMessage } from '../providers/types'; @@ -29,7 +29,7 @@ type CreateInboxMessage = { export class ConversationInboxService { constructor( private readonly chatSession: ChatSessionService, - private readonly policy: WorkspacePolicyService, + private readonly ac: PermissionAccess, private readonly storage: CopilotStorage, private readonly submissions: CompatSubmissionStore ) {} @@ -49,7 +49,11 @@ export class ConversationInboxService { ); if (blobs.length) { - await this.policy.assertCanUploadBlob(userId, session.config.workspaceId); + await this.ac + .user(userId) + .workspace(session.config.workspaceId) + .allowLocal() + .assert('Workspace.Blobs.Write'); } for (const blob of blobs) { diff --git a/packages/backend/server/src/plugins/copilot/conversation/policy.ts b/packages/backend/server/src/plugins/copilot/conversation/policy.ts index fc679f911c..1ed3a889d8 100644 --- a/packages/backend/server/src/plugins/copilot/conversation/policy.ts +++ b/packages/backend/server/src/plugins/copilot/conversation/policy.ts @@ -14,16 +14,8 @@ export class ConversationPolicy { ) {} async getQuota(userId: string) { - const isCopilotUser = await this.models.userFeature.has( - userId, - 'unlimited_copilot' - ); - - let limit: number | undefined; - if (!isCopilotUser) { - const quota = await this.quota.getUserQuota(userId); - limit = quota.copilotActionLimit; - } + const quota = await this.quota.getUserQuota(userId); + const limit = quota.copilotActionLimit; const used = await this.models.copilotSession.countUserMessages(userId); diff --git a/packages/backend/server/src/plugins/copilot/mcp/provider.ts b/packages/backend/server/src/plugins/copilot/mcp/provider.ts index 73f14d4f86..3551704700 100644 --- a/packages/backend/server/src/plugins/copilot/mcp/provider.ts +++ b/packages/backend/server/src/plugins/copilot/mcp/provider.ts @@ -3,7 +3,7 @@ import { pick } from 'lodash-es'; import z from 'zod/v3'; import { DocReader, DocWriter } from '../../../core/doc'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { clearEmbeddingChunk } from '../../../models'; import { IndexerService } from '../../indexer'; import { CopilotContextService } from '../context/service'; @@ -99,7 +99,7 @@ function defineTool( @Injectable() export class WorkspaceMcpProvider { constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly reader: DocReader, private readonly writer: DocWriter, private readonly context: CopilotContextService, diff --git a/packages/backend/server/src/plugins/copilot/resolver.ts b/packages/backend/server/src/plugins/copilot/resolver.ts index 63c775142b..3d882c85c2 100644 --- a/packages/backend/server/src/plugins/copilot/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/resolver.ts @@ -28,7 +28,7 @@ import { TooManyRequest, } from '../../base'; import { CurrentUser } from '../../core/auth'; -import { AccessController, DocAction } from '../../core/permission'; +import { DocAction, PermissionAccess } from '../../core/permission'; import { UserType } from '../../core/user'; import type { ListSessionOptions, UpdateChatSession } from '../../models'; import { CompatHistoryProjector } from './compat/history-projector'; @@ -365,7 +365,7 @@ export class CopilotResolver { private readonly modelNames = new Map(); constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly mutex: RequestMutex, private readonly prompt: PromptService, private readonly chatSession: ChatSessionService, @@ -815,7 +815,7 @@ export class CopilotResolver { @Throttle() @Resolver(() => UserType) export class UserCopilotResolver { - constructor(private readonly ac: AccessController) {} + constructor(private readonly ac: PermissionAccess) {} @ResolveField(() => CopilotType) async copilot( diff --git a/packages/backend/server/src/plugins/copilot/runtime/hosts/capability-policy-host.ts b/packages/backend/server/src/plugins/copilot/runtime/hosts/capability-policy-host.ts index 9a3411932e..86f0f5e28c 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/hosts/capability-policy-host.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/hosts/capability-policy-host.ts @@ -2,8 +2,7 @@ import { Injectable } from '@nestjs/common'; import { ModuleRef } from '@nestjs/core'; import { ServerFeature, ServerService } from '../../../../core'; -import { SubscriptionService } from '../../../payment/service'; -import { SubscriptionPlan, SubscriptionStatus } from '../../../payment/types'; +import { QuotaStateService } from '../../../../core/quota/state'; import type { ChatSession } from '../../session'; import { type ToolsConfig } from '../../types'; import { getTools } from '../../utils'; @@ -47,15 +46,14 @@ export class CapabilityPolicyHost { } try { - const subscription = await this.moduleRef - .get(SubscriptionService, { strict: false }) - .select(SubscriptionPlan.AI) - .getSubscription({ - userId, - plan: SubscriptionPlan.AI, - } as never); - - return subscription?.status === SubscriptionStatus.Active; + const state = await this.moduleRef + .get(QuotaStateService, { strict: false }) + .reconcileUserQuotaState(userId); + const flags = state.flags as { unlimitedCopilot?: boolean }; + return ( + !!flags.unlimitedCopilot || + ['pro', 'lifetime_pro', 'ai'].includes(state.plan) + ); } catch { return false; } diff --git a/packages/backend/server/src/plugins/copilot/runtime/task-policy.ts b/packages/backend/server/src/plugins/copilot/runtime/task-policy.ts index 70495bc788..1948e826e8 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/task-policy.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/task-policy.ts @@ -1,6 +1,6 @@ import { Injectable } from '@nestjs/common'; -import { Models } from '../../../models'; +import { QuotaStateService } from '../../../core/quota/state'; import { PromptService } from '../prompt/service'; export const DEFAULT_EMBEDDING_MODEL = 'gemini-embedding-001'; @@ -9,7 +9,7 @@ export const DEFAULT_RERANK_MODEL = 'gpt-4o-mini'; @Injectable() export class TaskPolicy { constructor( - private readonly models: Models, + private readonly quotaState: QuotaStateService, private readonly prompts: PromptService ) {} @@ -25,10 +25,11 @@ export class TaskPolicy { const prompt = await this.prompts.get('Transcript audio'); if (!prompt) return; - const hasAccess = await this.models.userFeature.has( - userId, - 'unlimited_copilot' - ); + const state = await this.quotaState.reconcileUserQuotaState(userId); + const flags = state.flags as { unlimitedCopilot?: boolean }; + const hasAccess = + !!flags.unlimitedCopilot || + ['pro', 'lifetime_pro', 'ai'].includes(state.plan); return prompt.optionalModels[hasAccess ? 1 : 0] ?? prompt.model; } } diff --git a/packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts b/packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts index e371cfd731..15ebf27808 100644 --- a/packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts +++ b/packages/backend/server/src/plugins/copilot/runtime/tool-runtime.ts @@ -2,7 +2,7 @@ import { Injectable } from '@nestjs/common'; import { Config } from '../../../base'; import { DocReader, DocWriter } from '../../../core/doc'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { Models } from '../../../models'; import { IndexerService } from '../../indexer'; import type { NodeTextMiddleware } from '../config'; @@ -48,7 +48,7 @@ export type ProviderSpecificToolResolver = ( export class ToolRuntime { constructor( private readonly config: Config, - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly context: CopilotContextService, private readonly docReader: DocReader, private readonly docWriter: DocWriter, diff --git a/packages/backend/server/src/plugins/copilot/tools/blob-read.ts b/packages/backend/server/src/plugins/copilot/tools/blob-read.ts index 06e30e7e77..58d6762708 100644 --- a/packages/backend/server/src/plugins/copilot/tools/blob-read.ts +++ b/packages/backend/server/src/plugins/copilot/tools/blob-read.ts @@ -1,7 +1,7 @@ import { Logger } from '@nestjs/common'; import { z } from 'zod'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { toolError } from './error'; import { defineTool } from './tool'; import type { ContextSession, CopilotChatOptions } from './types'; @@ -9,7 +9,7 @@ import type { ContextSession, CopilotChatOptions } from './types'; const logger = new Logger('ContextBlobReadTool'); export const buildBlobContentGetter = ( - ac: AccessController, + ac: PermissionAccess, context: ContextSession | null ) => { const getBlobContent = async ( diff --git a/packages/backend/server/src/plugins/copilot/tools/doc-keyword-search.ts b/packages/backend/server/src/plugins/copilot/tools/doc-keyword-search.ts index aad296b21d..91e0a40239 100644 --- a/packages/backend/server/src/plugins/copilot/tools/doc-keyword-search.ts +++ b/packages/backend/server/src/plugins/copilot/tools/doc-keyword-search.ts @@ -1,6 +1,6 @@ import { z } from 'zod'; -import type { AccessController } from '../../../core/permission'; +import type { PermissionAccess } from '../../../core/permission'; import type { Models } from '../../../models'; import type { IndexerService, SearchDoc } from '../../indexer'; import { workspaceSyncRequiredError } from './doc-sync'; @@ -9,7 +9,7 @@ import { defineTool } from './tool'; import type { CopilotChatOptions } from './types'; export const buildDocKeywordSearchGetter = ( - ac: AccessController, + ac: PermissionAccess, indexerService: IndexerService, models: Models ) => { diff --git a/packages/backend/server/src/plugins/copilot/tools/doc-read.ts b/packages/backend/server/src/plugins/copilot/tools/doc-read.ts index 2cba8f1064..3b581087bc 100644 --- a/packages/backend/server/src/plugins/copilot/tools/doc-read.ts +++ b/packages/backend/server/src/plugins/copilot/tools/doc-read.ts @@ -2,7 +2,7 @@ import { Logger } from '@nestjs/common'; import { z } from 'zod'; import { DocReader } from '../../../core/doc'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { Models } from '../../../models'; import { documentSyncPendingError, @@ -18,7 +18,7 @@ const isToolError = (result: ToolError | object): result is ToolError => 'type' in result && result.type === 'error'; export const buildDocContentGetter = ( - ac: AccessController, + ac: PermissionAccess, docReader: DocReader, models: Models ) => { diff --git a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts index 4958775559..5d39e58dcd 100644 --- a/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts +++ b/packages/backend/server/src/plugins/copilot/tools/doc-semantic-search.ts @@ -1,7 +1,7 @@ import { omit } from 'lodash-es'; import { z } from 'zod'; -import type { AccessController } from '../../../core/permission'; +import type { PermissionAccess } from '../../../core/permission'; import { type ChunkSimilarity, clearEmbeddingChunk, @@ -19,7 +19,7 @@ const getEmbeddingRouteContext = (options: CopilotChatOptions) => ({ }); export const buildDocSearchGetter = ( - ac: AccessController, + ac: PermissionAccess, context: CopilotContextService, sessionId: string | undefined, models: Models diff --git a/packages/backend/server/src/plugins/copilot/tools/doc-write.ts b/packages/backend/server/src/plugins/copilot/tools/doc-write.ts index 7f89d7bbc4..322c5e6bad 100644 --- a/packages/backend/server/src/plugins/copilot/tools/doc-write.ts +++ b/packages/backend/server/src/plugins/copilot/tools/doc-write.ts @@ -2,7 +2,7 @@ import { Logger } from '@nestjs/common'; import { z } from 'zod'; import { DocWriter } from '../../../core/doc'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { toolError } from './error'; import { defineTool } from './tool'; import type { CopilotChatOptions } from './types'; @@ -15,7 +15,7 @@ const stripLeadingH1 = (content: string) => const sanitizeTitle = (title: string) => title.replace(/[\r\n]+/g, ' ').trim(); export const buildDocCreateHandler = ( - ac: AccessController, + ac: PermissionAccess, writer: DocWriter ) => { return async ( @@ -57,7 +57,7 @@ export const buildDocCreateHandler = ( }; export const buildDocUpdateHandler = ( - ac: AccessController, + ac: PermissionAccess, writer: DocWriter ) => { return async ( @@ -95,7 +95,7 @@ export const buildDocUpdateHandler = ( }; export const buildDocUpdateMetaHandler = ( - ac: AccessController, + ac: PermissionAccess, writer: DocWriter ) => { return async (options: CopilotChatOptions, docId: string, title: string) => { diff --git a/packages/backend/server/src/plugins/copilot/transcript/realtime.ts b/packages/backend/server/src/plugins/copilot/transcript/realtime.ts index e6f5ca19d0..1b3db913e5 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/realtime.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/realtime.ts @@ -2,7 +2,7 @@ import { Injectable, OnModuleInit } from '@nestjs/common'; import { z } from 'zod'; import { CopilotTranscriptionJobNotFound } from '../../../base'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { RealtimeRegistry, realtimeTranscriptTaskRoom, @@ -13,7 +13,7 @@ import { CopilotTranscriptionReader } from './reader'; @Injectable() export class CopilotTranscriptRealtimeProvider implements OnModuleInit { constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly transcript: CopilotTranscriptionReader, private readonly registry: RealtimeRegistry ) {} diff --git a/packages/backend/server/src/plugins/copilot/transcript/resolver.ts b/packages/backend/server/src/plugins/copilot/transcript/resolver.ts index 07de9ade57..941fe5aa43 100644 --- a/packages/backend/server/src/plugins/copilot/transcript/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/transcript/resolver.ts @@ -21,7 +21,7 @@ import { type FileUpload, } from '../../../base'; import { CurrentUser } from '../../../core/auth'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { CopilotType } from '../resolver'; import type { TranscriptionJob } from './job'; import { buildLegacyProjection } from './projection'; @@ -295,7 +295,7 @@ const FinishedStatus: Set = new Set([ @Resolver(() => CopilotType) export class CopilotTranscriptionResolver { constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly transcript: CopilotTranscriptionService ) {} diff --git a/packages/backend/server/src/plugins/copilot/workspace/resolver.ts b/packages/backend/server/src/plugins/copilot/workspace/resolver.ts index 936d931386..3f4e6b436d 100644 --- a/packages/backend/server/src/plugins/copilot/workspace/resolver.ts +++ b/packages/backend/server/src/plugins/copilot/workspace/resolver.ts @@ -24,7 +24,7 @@ import { UserFriendlyError, } from '../../../base'; import { CurrentUser } from '../../../core/auth'; -import { AccessController } from '../../../core/permission'; +import { PermissionAccess } from '../../../core/permission'; import { WorkspaceType } from '../../../core/workspaces'; import { COPILOT_LOCKER } from '../resolver'; import { MAX_EMBEDDABLE_SIZE } from '../utils'; @@ -49,7 +49,7 @@ export class CopilotWorkspaceConfigType { */ @Resolver(() => WorkspaceType) export class CopilotWorkspaceEmbeddingResolver { - constructor(private readonly ac: AccessController) {} + constructor(private readonly ac: PermissionAccess) {} @ResolveField(() => CopilotWorkspaceConfigType, { complexity: 2, @@ -70,7 +70,7 @@ export class CopilotWorkspaceEmbeddingResolver { @Resolver(() => CopilotWorkspaceConfigType) export class CopilotWorkspaceEmbeddingConfigResolver { constructor( - private readonly ac: AccessController, + private readonly ac: PermissionAccess, private readonly mutex: Mutex, private readonly copilotWorkspace: CopilotWorkspaceService ) {} diff --git a/packages/backend/server/src/plugins/indexer/__tests__/service.spec.ts b/packages/backend/server/src/plugins/indexer/__tests__/service.spec.ts index e90a781e02..e34ac55c1c 100644 --- a/packages/backend/server/src/plugins/indexer/__tests__/service.spec.ts +++ b/packages/backend/server/src/plugins/indexer/__tests__/service.spec.ts @@ -2313,6 +2313,67 @@ test('should search docs by keyword work', async t => { ); }); +test('should search docs by keyword with doc id filter', async t => { + const workspaceId = workspace.id; + const docId1 = randomUUID(); + const docId2 = randomUUID(); + + await module.create(Mockers.DocMeta, { + workspaceId, + docId: docId1, + title: 'hello filtered 1', + }); + await module.create(Mockers.DocMeta, { + workspaceId, + docId: docId2, + title: 'hello filtered 2', + }); + + await indexerService.write( + SearchTable.block, + [ + { + workspaceId, + docId: docId1, + blockId: 'filtered-block1', + content: 'hello filtered', + flavour: 'affine:text', + createdByUserId: user.id, + updatedByUserId: user.id, + createdAt: new Date('2025-06-20T00:00:00.000Z'), + updatedAt: new Date('2025-06-20T00:00:00.000Z'), + }, + { + workspaceId, + docId: docId2, + blockId: 'filtered-block2', + content: 'hello filtered', + flavour: 'affine:text', + createdByUserId: user.id, + updatedByUserId: user.id, + createdAt: new Date('2025-06-20T00:00:01.000Z'), + updatedAt: new Date('2025-06-20T00:00:01.000Z'), + }, + ], + { + refresh: true, + } + ); + + const rows = await indexerService.searchDocsByKeyword( + workspaceId, + 'hello filtered', + { + docIds: [docId2], + } + ); + + t.deepEqual( + rows.map(row => row.docId), + [docId2] + ); +}); + // #endregion test('should rebuild manticore indexes and requeue workspaces', async t => { diff --git a/packages/backend/server/src/plugins/indexer/index.ts b/packages/backend/server/src/plugins/indexer/index.ts index a550d84130..6e2afc817c 100644 --- a/packages/backend/server/src/plugins/indexer/index.ts +++ b/packages/backend/server/src/plugins/indexer/index.ts @@ -4,6 +4,7 @@ import { Module } from '@nestjs/common'; import { ServerConfigModule } from '../../core/config'; import { PermissionModule } from '../../core/permission'; +import { QuotaServiceModule } from '../../core/quota'; import { IndexerEvent } from './event'; import { SearchProviderFactory } from './factory'; import { IndexerJob } from './job'; @@ -12,7 +13,7 @@ import { IndexerResolver } from './resolver'; import { IndexerService } from './service'; @Module({ - imports: [ServerConfigModule, PermissionModule], + imports: [ServerConfigModule, PermissionModule, QuotaServiceModule], providers: [ IndexerResolver, IndexerService, diff --git a/packages/backend/server/src/plugins/indexer/resolver.ts b/packages/backend/server/src/plugins/indexer/resolver.ts index 4b8c6acd9a..3c1e06799c 100644 --- a/packages/backend/server/src/plugins/indexer/resolver.ts +++ b/packages/backend/server/src/plugins/indexer/resolver.ts @@ -1,18 +1,19 @@ import { Args, Parent, ResolveField, Resolver } from '@nestjs/graphql'; +import { Prisma, PrismaClient } from '@prisma/client'; import { CurrentUser } from '../../core/auth'; -import { AccessController } from '../../core/permission'; +import { PermissionAccess, PermissionService } from '../../core/permission'; +import { QuotaStateService } from '../../core/quota/state'; import { UserType } from '../../core/user'; import { WorkspaceType } from '../../core/workspaces'; -import { Models } from '../../models'; -import { AggregateBucket } from './providers'; -import { IndexerService, SearchNodeWithMeta } from './service'; +import { IndexerService } from './service'; import { AggregateInput, AggregateResultObjectType, SearchDocObjectType, SearchDocsInput, SearchInput, + SearchQuery, SearchQueryOccur, SearchQueryType, SearchResultObjectType, @@ -22,8 +23,10 @@ import { export class IndexerResolver { constructor( private readonly indexer: IndexerService, - private readonly ac: AccessController, - private readonly models: Models + private readonly ac: PermissionAccess, + private readonly db: PrismaClient, + private readonly permission: PermissionService, + private readonly quotaState: QuotaStateService ) {} @ResolveField(() => SearchResultObjectType, { @@ -37,18 +40,22 @@ export class IndexerResolver { // currentUser can read the workspace await this.ac.user(me.id).workspace(workspace.id).assert('Workspace.Read'); this.#addWorkspaceFilter(workspace, input); + if (!(await this.#addReadableDocFilter(workspace, me, input))) { + return { + nodes: [], + pagination: { + count: 0, + hasMore: false, + }, + }; + } const result = await this.indexer.search(input); - const nodes = await this.#filterUserReadableDocs( - workspace, - me, - result.nodes - ); return { - nodes, + nodes: result.nodes, pagination: { count: result.total, - hasMore: nodes.length > 0, + hasMore: result.nodes.length > 0, nextCursor: result.nextCursor, }, }; @@ -65,24 +72,22 @@ export class IndexerResolver { // currentUser can read the workspace await this.ac.user(me.id).workspace(workspace.id).assert('Workspace.Read'); this.#addWorkspaceFilter(workspace, input); + if (!(await this.#addReadableDocFilter(workspace, me, input))) { + return { + buckets: [], + pagination: { + count: 0, + hasMore: false, + }, + }; + } const result = await this.indexer.aggregate(input); - const needs: AggregateBucket[] = []; - for (const bucket of result.buckets) { - bucket.hits.nodes = await this.#filterUserReadableDocs( - workspace, - me, - bucket.hits.nodes as SearchNodeWithMeta[] - ); - if (bucket.hits.nodes.length > 0) { - needs.push(bucket); - } - } return { - buckets: needs, + buckets: result.buckets, pagination: { count: result.total, - hasMore: needs.length > 0, + hasMore: result.buckets.length > 0, nextCursor: result.nextCursor, }, }; @@ -96,19 +101,17 @@ export class IndexerResolver { @Parent() workspace: WorkspaceType, @Args('input') input: SearchDocsInput ): Promise { + const readableDocIds = await this.#readableDocIdsForSearch(workspace, me); const docs = await this.indexer.searchDocsByKeyword( workspace.id, input.keyword, { limit: input.limit, + docIds: readableDocIds ?? undefined, } ); - const needs = await this.ac - .user(me.id) - .workspace(workspace.id) - .docs(docs, 'Doc.Read'); - return needs; + return docs; } #addWorkspaceFilter( @@ -130,36 +133,84 @@ export class IndexerResolver { }; } - /** - * filter user readable docs on team workspace - */ - async #filterUserReadableDocs( + async #addReadableDocFilter( workspace: WorkspaceType, user: UserType, - nodes: SearchNodeWithMeta[] + input: SearchInput | AggregateInput ) { - if (nodes.length === 0) { - return nodes; + const docIds = await this.#readableDocIdsForSearch(workspace, user); + if (docIds === null) { + return true; } - const isTeamWorkspace = await this.models.workspaceFeature.has( - workspace.id, - 'team_plan_v1' + if (docIds.length === 0) { + return false; + } + + input.query = { + type: SearchQueryType.boolean, + occur: SearchQueryOccur.must, + queries: [input.query, this.#docIdFilterQuery(docIds)], + }; + return true; + } + + async #readableDocIdsForSearch(workspace: WorkspaceType, user: UserType) { + const state = await this.quotaState.reconcileWorkspaceQuotaState( + workspace.id ); + const isTeamWorkspace = + state.plan === 'team' || state.plan === 'selfhost_team'; if (!isTeamWorkspace) { - return nodes; + return null; } - const needs = await this.ac - .user(user.id) - .workspace(workspace.id) - .docs( - nodes.map(node => ({ - node, - docId: node._source.docId, - })), - 'Doc.Read' - ); - return needs.map(node => node.node); + return await this.#listReadableDocIds(workspace, user); + } + + #docIdFilterQuery(docIds: string[]): SearchQuery { + return { + type: SearchQueryType.boolean, + occur: SearchQueryOccur.should, + queries: docIds.map(docId => ({ + type: SearchQueryType.match, + field: 'docId', + match: docId, + })), + }; + } + + async #listReadableDocIds(workspace: WorkspaceType, user: UserType) { + const input = { + userId: user.id, + workspaceId: workspace.id, + action: 'Doc.Read', + docIdColumn: Prisma.raw('candidate_docs.doc_id'), + } as const; + const predicate = this.permission.docReadableSqlPredicate(input); + const fallbackPredicate = + this.permission.fallbackDocReadableSqlPredicate(input); + const query = (predicate: Prisma.Sql) => + this.db.$queryRaw<{ docId: string }[]>` + WITH candidate_docs AS ( + SELECT "workspace_pages"."page_id" AS doc_id + FROM "workspace_pages" + WHERE "workspace_pages"."workspace_id" = ${workspace.id} + UNION + SELECT "snapshots"."guid" AS doc_id + FROM "snapshots" + WHERE "snapshots"."workspace_id" = ${workspace.id} + ) + SELECT candidate_docs.doc_id AS "docId" + FROM candidate_docs + WHERE ${predicate} + `; + const rows = await query(predicate).catch(error => { + if (!fallbackPredicate) { + throw error; + } + return query(fallbackPredicate); + }); + return rows.map(row => row.docId); } } diff --git a/packages/backend/server/src/plugins/indexer/service.ts b/packages/backend/server/src/plugins/indexer/service.ts index 8e75dde571..2cf943d10e 100644 --- a/packages/backend/server/src/plugins/indexer/service.ts +++ b/packages/backend/server/src/plugins/indexer/service.ts @@ -497,8 +497,13 @@ export class IndexerService { keyword: string, options?: { limit?: number; + docIds?: string[]; } ): Promise { + if (options?.docIds?.length === 0) { + return []; + } + const limit = options?.limit ?? 20; const result = await this.aggregate({ table: SearchTable.block, @@ -512,6 +517,19 @@ export class IndexerService { field: 'workspaceId', match: workspaceId, }, + ...(options?.docIds + ? [ + { + type: SearchQueryType.boolean, + occur: SearchQueryOccur.should, + queries: options.docIds.map(docId => ({ + type: SearchQueryType.match, + field: 'docId', + match: docId, + })), + }, + ] + : []), { type: SearchQueryType.boolean, occur: SearchQueryOccur.must, diff --git a/packages/backend/server/src/plugins/license/index.ts b/packages/backend/server/src/plugins/license/index.ts index 5aef0d83fa..bc1cc5d006 100644 --- a/packages/backend/server/src/plugins/license/index.ts +++ b/packages/backend/server/src/plugins/license/index.ts @@ -1,5 +1,6 @@ import { Module } from '@nestjs/common'; +import { EntitlementModule } from '../../core/entitlement'; import { FeatureModule } from '../../core/features'; import { PermissionModule } from '../../core/permission'; import { QuotaModule } from '../../core/quota'; @@ -8,7 +9,13 @@ import { AdminLicenseResolver, LicenseResolver } from './resolver'; import { LicenseService } from './service'; @Module({ - imports: [FeatureModule, QuotaModule, PermissionModule, WorkspaceModule], + imports: [ + FeatureModule, + QuotaModule, + PermissionModule, + WorkspaceModule, + EntitlementModule, + ], providers: [LicenseService, LicenseResolver, AdminLicenseResolver], }) export class LicenseModule {} diff --git a/packages/backend/server/src/plugins/license/resolver.ts b/packages/backend/server/src/plugins/license/resolver.ts index f0f809f1b3..dab7bec695 100644 --- a/packages/backend/server/src/plugins/license/resolver.ts +++ b/packages/backend/server/src/plugins/license/resolver.ts @@ -15,7 +15,7 @@ import GraphQLUpload, { import { toBuffer, UseNamedGuard } from '../../base'; import { CurrentUser } from '../../core/auth'; import { Admin } from '../../core/common'; -import { AccessController } from '../../core/permission'; +import { PermissionAccess } from '../../core/permission'; import { WorkspaceType } from '../../core/workspaces'; import { SubscriptionPlan, @@ -86,7 +86,7 @@ export class AdminLicensePreview { export class LicenseResolver { constructor( private readonly service: LicenseService, - private readonly ac: AccessController + private readonly ac: PermissionAccess ) {} @ResolveField(() => License, { diff --git a/packages/backend/server/src/plugins/license/service.ts b/packages/backend/server/src/plugins/license/service.ts index 41ea08352d..3a95214430 100644 --- a/packages/backend/server/src/plugins/license/service.ts +++ b/packages/backend/server/src/plugins/license/service.ts @@ -1,9 +1,6 @@ -import { createDecipheriv, createVerify } from 'node:crypto'; - import { Injectable, Logger } from '@nestjs/common'; import { Cron, CronExpression } from '@nestjs/schedule'; import { InstalledLicense, PrismaClient } from '@prisma/client'; -import { z } from 'zod'; import { CryptoHelper, @@ -16,8 +13,11 @@ import { UserFriendlyError, WorkspaceLicenseAlreadyExists, } from '../../base'; +import { EntitlementService } from '../../core/entitlement'; import { WorkspacePolicyService } from '../../core/permission'; +import { QuotaStateService } from '../../core/quota/state'; import { Models } from '../../models'; +import { ResolvedEntitlement, resolveEntitlementV1 } from '../../native'; import { SubscriptionPlan, SubscriptionRecurring, @@ -31,6 +31,8 @@ interface License { endAt: number; } +const AFFINE_PRO_REQUEST_TIMEOUT = 10_000; + export interface LicensePreview { id: string; workspaceId: string; @@ -45,27 +47,6 @@ export interface LicensePreview { valid: boolean; } -const BaseLicenseSchema = z.object({ - entity: z.string().nonempty(), - issuer: z.string().nonempty(), - issuedAt: z.string().datetime(), - expiresAt: z.string().datetime(), -}); - -const TeamLicenseSchema = z - .object({ - subject: z.literal(SubscriptionPlan.SelfHostedTeam), - data: z.object({ - id: z.string().nonempty(), - workspaceId: z.string().nonempty(), - plan: z.literal(SubscriptionPlan.SelfHostedTeam), - recurring: z.nativeEnum(SubscriptionRecurring), - quantity: z.number().positive(), - endAt: z.string().datetime(), - }), - }) - .extend(BaseLicenseSchema.shape); - @Injectable() export class LicenseService { private readonly logger = new Logger(LicenseService.name); @@ -75,37 +56,11 @@ export class LicenseService { private readonly event: EventBus, private readonly models: Models, private readonly crypto: CryptoHelper, - private readonly policy: WorkspacePolicyService + private readonly policy: WorkspacePolicyService, + private readonly entitlement: EntitlementService, + private readonly quotaState: QuotaStateService ) {} - @OnEvent('workspace.subscription.activated') - async onWorkspaceSubscriptionUpdated({ - workspaceId, - plan, - recurring, - quantity, - }: Events['workspace.subscription.activated']) { - switch (plan) { - case SubscriptionPlan.SelfHostedTeam: - await this.models.workspaceFeature.add( - workspaceId, - 'team_plan_v1', - `${recurring} team subscription activated`, - { - memberLimit: quantity, - } - ); - this.event.emit('workspace.members.allocateSeats', { - workspaceId, - quantity, - }); - await this.policy.reconcileWorkspaceQuotaState(workspaceId); - break; - default: - break; - } - } - @OnEvent('workspace.subscription.canceled') async onWorkspaceSubscriptionCanceled({ workspaceId, @@ -137,74 +92,54 @@ export class LicenseService { } async installLicense(workspaceId: string, license: Buffer) { - const payload = this.decryptWorkspaceTeamLicense(workspaceId, license); - const data = payload.data; - const now = new Date(); - - if (new Date(payload.expiresAt) < now || new Date(data.endAt) < now) { + const resolved = this.resolveWorkspaceTeamLicense(workspaceId, license); + if (!resolved.valid) { throw new LicenseExpired(); } - const installed = await this.db.installedLicense.upsert({ - where: { - workspaceId, - }, - update: { - key: data.id, - expiredAt: new Date(data.endAt), - validatedAt: new Date(), - recurring: data.recurring, - quantity: data.quantity, - variant: SubscriptionVariant.Onetime, - license, - }, - create: { - key: data.id, - workspaceId, - expiredAt: new Date(data.endAt), - validateKey: '', - validatedAt: new Date(), - recurring: data.recurring, - quantity: data.quantity, - variant: SubscriptionVariant.Onetime, - license, - }, - }); + const validatedAt = new Date(); await this.event.emitAsync('workspace.subscription.activated', { workspaceId, - plan: data.plan, - recurring: data.recurring, - quantity: data.quantity, + plan: SubscriptionPlan.SelfHostedTeam, + recurring: this.licenseRecurring(resolved), + quantity: this.licenseQuantity(resolved), + }); + await this.entitlement.upsertFromSelfhostLicense({ + workspaceId, + recurring: this.licenseRecurring(resolved), + quantity: this.licenseQuantity(resolved), + expiresAt: this.licenseExpiresAt(resolved), + validatedAt, + variant: SubscriptionVariant.Onetime, + license, }); - return installed; + return this.db.installedLicense.findUniqueOrThrow({ + where: { workspaceId }, + }); } previewLicense(license: Buffer): LicensePreview { - const payload = this.decryptWorkspaceTeamLicensePayload(license); - const data = payload.data; - const now = new Date(); - const expiresAt = new Date(payload.expiresAt); - const endAt = new Date(data.endAt); - - if (expiresAt < now || endAt < now) { + const resolved = this.resolveWorkspaceTeamLicense(null, license); + if (!resolved.valid) { throw new InvalidLicenseToActivate({ reason: 'Invalid license.', }); } + const expiresAt = this.licenseExpiresAt(resolved); return { - id: data.id, - workspaceId: data.workspaceId, - plan: data.plan, - recurring: data.recurring, - quantity: data.quantity, - issuedAt: new Date(payload.issuedAt), + id: this.licenseSubjectId(resolved), + workspaceId: this.licenseWorkspaceId(resolved), + plan: SubscriptionPlan.SelfHostedTeam, + recurring: this.licenseRecurring(resolved), + quantity: this.licenseQuantity(resolved), + issuedAt: new Date(resolved.issuedAt ?? ''), expiresAt, - endAt, - entity: payload.entity, - issuer: payload.issuer, + endAt: expiresAt, + entity: resolved.entity ?? '', + issuer: resolved.issuer ?? '', valid: true, }; } @@ -215,6 +150,12 @@ export class LicenseService { if (installedLicense) { throw new WorkspaceLicenseAlreadyExists(); } + const occupiedLicense = await this.db.installedLicense.findUnique({ + where: { key: licenseKey }, + }); + if (occupiedLicense) { + throw new WorkspaceLicenseAlreadyExists(); + } const data = await this.fetchAffinePro( `/api/team/licenses/${licenseKey}/activate`, @@ -223,28 +164,9 @@ export class LicenseService { } ); - const license = await this.db.installedLicense.upsert({ - where: { - workspaceId, - }, - update: { - key: licenseKey, - validatedAt: new Date(), - validateKey: data.res.headers.get('x-next-validate-key') ?? '', - expiredAt: new Date(data.endAt), - recurring: data.recurring, - quantity: data.quantity, - }, - create: { - workspaceId, - key: licenseKey, - expiredAt: new Date(data.endAt), - validatedAt: new Date(), - validateKey: data.res.headers.get('x-next-validate-key') ?? '', - recurring: data.recurring, - quantity: data.quantity, - }, - }); + const validatedAt = new Date(); + const expiresAt = this.remoteLicenseExpiresAt(data); + const validateKey = data.res.headers.get('x-next-validate-key') ?? ''; this.event.emit('workspace.subscription.activated', { workspaceId, @@ -252,7 +174,40 @@ export class LicenseService { recurring: data.recurring, quantity: data.quantity, }); - return license; + await this.entitlement.upsertFromValidatedSelfhostLicense({ + workspaceId, + licenseKey, + recurring: data.recurring, + quantity: data.quantity, + expiresAt, + validatedAt, + validateKey, + }); + + return this.db.installedLicense.upsert({ + where: { workspaceId }, + update: { + key: licenseKey, + quantity: data.quantity, + recurring: data.recurring, + variant: null, + validateKey, + validatedAt, + expiredAt: expiresAt, + license: null, + }, + create: { + workspaceId, + key: licenseKey, + quantity: data.quantity, + recurring: data.recurring, + variant: null, + validateKey, + validatedAt, + expiredAt: expiresAt, + license: null, + }, + }); } async removeTeamLicense(workspaceId: string) { @@ -271,6 +226,7 @@ export class LicenseService { workspaceId: license.workspaceId, }, }); + await this.entitlement.revokeBySubject('selfhost_license', license.key); if (license.variant !== SubscriptionVariant.Onetime) { await this.deactivateTeamLicense(license); @@ -334,9 +290,11 @@ export class LicenseService { } if (license.variant === SubscriptionVariant.Onetime) { + const state = + await this.quotaState.reconcileWorkspaceQuotaState(workspaceId); this.event.emit('workspace.members.allocateSeats', { workspaceId, - quantity: license.quantity, + quantity: state.seatLimit ?? 0, }); return; @@ -389,7 +347,7 @@ export class LicenseService { for (const license of licenses) { if (license.variant === SubscriptionVariant.Onetime) { - this.revalidateOnetimeLicense(license); + await this.revalidateOnetimeLicense(license); } else { await this.revalidateRecurringLicense(license); } @@ -407,18 +365,9 @@ export class LicenseService { } ); - await this.db.installedLicense.update({ - where: { - key: license.key, - }, - data: { - validatedAt: new Date(), - validateKey: res.res.headers.get('x-next-validate-key') ?? '', - quantity: res.quantity, - recurring: res.recurring, - expiredAt: new Date(res.endAt), - }, - }); + const validatedAt = new Date(); + const expiresAt = this.remoteLicenseExpiresAt(res); + const validateKey = res.res.headers.get('x-next-validate-key') ?? ''; this.event.emit('workspace.subscription.activated', { workspaceId: license.workspaceId, @@ -426,6 +375,39 @@ export class LicenseService { recurring: res.recurring, quantity: res.quantity, }); + await this.db.installedLicense.update({ + where: { key: license.key }, + data: { + quantity: res.quantity, + recurring: res.recurring, + validateKey, + validatedAt, + expiredAt: expiresAt, + }, + }); + + if (license.license) { + await this.entitlement.upsertFromSelfhostLicense({ + workspaceId: license.workspaceId, + licenseKey: license.key, + recurring: res.recurring, + quantity: res.quantity, + expiresAt, + validatedAt, + validateKey, + license: Buffer.from(license.license), + }); + } else { + await this.entitlement.upsertFromValidatedSelfhostLicense({ + workspaceId: license.workspaceId, + licenseKey: license.key, + recurring: res.recurring, + quantity: res.quantity, + expiresAt, + validatedAt, + validateKey, + }); + } return res; } catch (e) { @@ -441,12 +423,21 @@ export class LicenseService { plan: SubscriptionPlan.SelfHostedTeam, recurring: SubscriptionRecurring.Monthly, }); + await this.entitlement.revokeBySubject('selfhost_license', license.key); } return null; } } + private remoteLicenseExpiresAt(license: Pick) { + const expiresAt = new Date(license.endAt); + if (!Number.isFinite(expiresAt.getTime()) || expiresAt <= new Date()) { + throw new LicenseExpired(); + } + return expiresAt; + } + private async fetchAffinePro( path: string, init?: RequestInit @@ -455,8 +446,14 @@ export class LicenseService { process.env.AFFINE_PRO_SERVER_ENDPOINT ?? 'https://app.affine.pro'; try { + const signal = + init?.signal ?? + (AbortSignal.timeout + ? AbortSignal.timeout(AFFINE_PRO_REQUEST_TIMEOUT) + : undefined); const res = await fetch(endpoint + path, { ...init, + signal, headers: { 'Content-Type': 'application/json', ...init?.headers, @@ -486,25 +483,33 @@ export class LicenseService { } } - private revalidateOnetimeLicense(license: InstalledLicense) { + private async revalidateOnetimeLicense(license: InstalledLicense) { const buf = license.license; let valid = !!buf; if (buf) { try { - const { data } = this.decryptWorkspaceTeamLicense( + const resolved = this.resolveWorkspaceTeamLicense( license.workspaceId, Buffer.from(buf) ); - if (new Date(data.endAt) < new Date()) { + if (!resolved.valid) { valid = false; } else { this.event.emit('workspace.subscription.activated', { workspaceId: license.workspaceId, - plan: data.plan, - recurring: data.recurring, - quantity: data.quantity, + plan: SubscriptionPlan.SelfHostedTeam, + recurring: this.licenseRecurring(resolved), + quantity: this.licenseQuantity(resolved), + }); + await this.entitlement.upsertFromSelfhostLicense({ + workspaceId: license.workspaceId, + recurring: this.licenseRecurring(resolved), + quantity: this.licenseQuantity(resolved), + expiresAt: this.licenseExpiresAt(resolved), + validatedAt: new Date(), + license: Buffer.from(buf), }); } } catch { @@ -518,116 +523,101 @@ export class LicenseService { plan: SubscriptionPlan.SelfHostedTeam, recurring: SubscriptionRecurring.Monthly, }); + await this.entitlement.revokeBySubject('selfhost_license', license.key); } } - private decryptWorkspaceTeamLicense(workspaceId: string, buf: Buffer) { - const payload = this.decryptWorkspaceTeamLicensePayload(buf); - - if (new Date(payload.expiresAt) < new Date()) { - throw new InvalidLicenseToActivate({ - reason: - 'License file has expired. Please contact with Affine support to fetch a latest one.', - }); - } - - if (payload.data.workspaceId !== workspaceId) { - throw new InvalidLicenseToActivate({ - reason: 'Workspace mismatched with license.', - }); - } - - return payload; - } - - private decryptWorkspaceTeamLicensePayload(buf: Buffer) { + private resolveWorkspaceTeamLicense(workspaceId: string | null, buf: Buffer) { if (!this.crypto.AFFiNEProPublicKey) { throw new InternalServerError( 'License public key is not loaded. Please contact with Affine support.' ); } - // we use workspace id as aes key hash plain text content - // verify signature to make sure the payload or signature is not forged - const { payload: payloadStr, signature, iv } = this.decryptLicense(buf); - - const verifier = createVerify('rsa-sha256'); - verifier.update(iv); - verifier.update(payloadStr); - const valid = verifier.verify( - this.crypto.AFFiNEProPublicKey, - signature, - 'hex' - ); - if (!valid) { - throw new InvalidLicenseToActivate({ - reason: 'Invalid license signature.', - }); - } - - const payload = JSON.parse(payloadStr); - - const parseResult = TeamLicenseSchema.safeParse(payload); - - if (!parseResult.success) { - throw new InvalidLicenseToActivate({ - reason: 'Invalid license payload.', - }); - } - - return parseResult.data; - } - - private decryptLicense(buf: Buffer) { if (!this.crypto.AFFiNEProLicenseAESKey) { throw new InternalServerError( 'License AES key is not loaded. Please contact with Affine support.' ); } - if (buf.length < 2) { + const resolved = resolveEntitlementV1({ + deploymentType: 'selfhosted', + targetType: 'workspace', + targetId: workspaceId ?? undefined, + signedPayload: buf, + publicKey: this.crypto.AFFiNEProPublicKey.toString(), + licenseAesKey: this.crypto.AFFiNEProLicenseAESKey.toString('hex'), + now: new Date().toISOString(), + }); + + if (resolved.errorCode === 'workspace_mismatch') { throw new InvalidLicenseToActivate({ - reason: 'Invalid license file.', + reason: 'Workspace mismatched with license.', }); } - try { - const ivLength = buf.readUint8(0); - const authTagLength = buf.readUInt8(1); + if (!resolved.valid && resolved.errorCode === 'expired_end_at') { + throw new LicenseExpired(); + } - const iv = buf.subarray(2, 2 + ivLength); - const tag = buf.subarray(2 + ivLength, 2 + ivLength + authTagLength); - const payload = buf.subarray(2 + ivLength + authTagLength); - - const decipher = createDecipheriv( - 'aes-256-gcm', - this.crypto.AFFiNEProLicenseAESKey, - iv, - { - authTagLength, - } - ); - decipher.setAuthTag(tag); - - const decrypted = Buffer.concat([ - decipher.update(payload), - decipher.final(), - ]); - - const data = JSON.parse(decrypted.toString('utf-8')) as { - payload: string; - signature: string; - }; - - return { - ...data, - iv, - }; - } catch { - // we use workspace id as aes key hash plain text content + if (!resolved.valid && resolved.status === 'expired') { throw new InvalidLicenseToActivate({ - reason: 'Failed to verify the license.', + reason: + 'License file has expired. Please contact with Affine support to fetch a latest one.', }); } + + if (!resolved.valid && resolved.status !== 'expired') { + throw new InvalidLicenseToActivate({ + reason: resolved.errorMessage ?? 'Failed to verify the license.', + }); + } + + return resolved; + } + + private licenseSubjectId(resolved: ResolvedEntitlement) { + if (!resolved.subjectId) { + throw new InvalidLicenseToActivate({ + reason: 'Invalid license payload.', + }); + } + return resolved.subjectId; + } + + private licenseWorkspaceId(resolved: ResolvedEntitlement) { + if (!resolved.targetId) { + throw new InvalidLicenseToActivate({ + reason: 'Invalid license payload.', + }); + } + return resolved.targetId; + } + + private licenseRecurring(resolved: ResolvedEntitlement) { + if (!resolved.recurring) { + throw new InvalidLicenseToActivate({ + reason: 'Invalid license payload.', + }); + } + return resolved.recurring as SubscriptionRecurring; + } + + private licenseQuantity(resolved: ResolvedEntitlement) { + if (!resolved.quantity) { + throw new InvalidLicenseToActivate({ + reason: 'Invalid license payload.', + }); + } + return resolved.quantity; + } + + private licenseExpiresAt(resolved: ResolvedEntitlement) { + if (!resolved.expiresAt) { + throw new InvalidLicenseToActivate({ + reason: 'Invalid license payload.', + }); + } + return new Date(resolved.expiresAt); } } diff --git a/packages/backend/server/src/plugins/payment/cron.ts b/packages/backend/server/src/plugins/payment/cron.ts index cead2f1150..d908606166 100644 --- a/packages/backend/server/src/plugins/payment/cron.ts +++ b/packages/backend/server/src/plugins/payment/cron.ts @@ -3,6 +3,7 @@ import { Cron, CronExpression } from '@nestjs/schedule'; import { PrismaClient, Provider } from '@prisma/client'; import { EventBus, JobQueue, OneHour, OnJob } from '../../base'; +import { EntitlementService } from '../../core/entitlement'; import { RevenueCatWebhookHandler } from './revenuecat'; import { SubscriptionService } from './service'; import { StripeFactory } from './stripe'; @@ -34,7 +35,8 @@ export class SubscriptionCronJobs { private readonly queue: JobQueue, private readonly rcHandler: RevenueCatWebhookHandler, private readonly stripeFactory: StripeFactory, - private readonly subscription: SubscriptionService + private readonly subscription: SubscriptionService, + private readonly entitlement: EntitlementService ) {} private getDateRange(after: number, base: number | Date = Date.now()) { @@ -157,6 +159,12 @@ export class SubscriptionCronJobs { }); for (const subscription of subscriptions) { + await this.entitlement.revokeCloudSubscription({ + targetId: subscription.targetId, + plan: subscription.plan, + subscriptionId: subscription.id, + stripeSubscriptionId: subscription.stripeSubscriptionId, + }); await this.db.subscription.delete({ where: { targetId_plan: { diff --git a/packages/backend/server/src/plugins/payment/event.ts b/packages/backend/server/src/plugins/payment/event.ts index ab5d4ab0d8..61a56d5a32 100644 --- a/packages/backend/server/src/plugins/payment/event.ts +++ b/packages/backend/server/src/plugins/payment/event.ts @@ -2,47 +2,32 @@ import { Injectable } from '@nestjs/common'; import { EventBus, OnEvent } from '../../base'; import { WorkspacePolicyService } from '../../core/permission'; +import { QuotaStateService } from '../../core/quota/state'; import { WorkspaceService } from '../../core/workspaces'; -import { Models } from '../../models'; -import { SubscriptionPlan, SubscriptionRecurring } from './types'; +import { SubscriptionPlan } from './types'; @Injectable() export class PaymentEventHandlers { constructor( private readonly workspace: WorkspaceService, - private readonly models: Models, - private readonly event: EventBus, - private readonly policy: WorkspacePolicyService + private readonly policy: WorkspacePolicyService, + private readonly quotaState: QuotaStateService, + private readonly event: EventBus ) {} @OnEvent('workspace.subscription.activated') async onWorkspaceSubscriptionUpdated({ workspaceId, plan, - recurring, - quantity, }: Events['workspace.subscription.activated']) { switch (plan) { case 'team': { const isTeam = await this.workspace.isTeamWorkspace(workspaceId); - await this.models.workspaceFeature.add( - workspaceId, - 'team_plan_v1', - `${recurring} team subscription activated`, - { - memberLimit: quantity, - } - ); - this.event.emit('workspace.members.allocateSeats', { - workspaceId, - quantity, - }); if (!isTeam) { // this event will triggered when subscription is activated or changed // we only send emails when the team workspace is activated await this.workspace.sendTeamWorkspaceUpgradedEmail(workspaceId); } - await this.policy.reconcileWorkspaceQuotaState(workspaceId); break; } default: @@ -68,22 +53,10 @@ export class PaymentEventHandlers { async onUserSubscriptionUpdated({ userId, plan, - recurring, }: Events['user.subscription.activated']) { switch (plan) { case SubscriptionPlan.AI: - await this.models.userFeature.add( - userId, - 'unlimited_copilot', - 'subscription activated' - ); - break; case SubscriptionPlan.Pro: - await this.models.userFeature.switchQuota( - userId, - recurring === 'lifetime' ? 'lifetime_pro_plan_v1' : 'pro_plan_v1', - 'subscription activated' - ); await this.policy.reconcileOwnedWorkspaces(userId); break; default: @@ -95,43 +68,35 @@ export class PaymentEventHandlers { async onUserSubscriptionCanceled({ userId, plan, - recurring, }: Events['user.subscription.canceled']) { switch (plan) { case SubscriptionPlan.AI: - await this.models.userFeature.remove(userId, 'unlimited_copilot'); - break; case SubscriptionPlan.Pro: { - // if user disputed a lifetime plan, we just switch them to free plan directly - if (recurring === SubscriptionRecurring.Lifetime) { - await this.models.userFeature.switchQuota( - userId, - 'free_plan_v1', - 'lifetime subscription canceled' - ); - await this.policy.reconcileOwnedWorkspaces(userId); - break; - } - - // edge case: when user switch from recurring Pro plan to `Lifetime` plan, - // a subscription canceled event will be triggered because `Lifetime` plan is not subscription based - const isLifetimeUser = await this.models.userFeature.has( - userId, - 'lifetime_pro_plan_v1' - ); - - if (!isLifetimeUser) { - await this.models.userFeature.switchQuota( - userId, - 'free_plan_v1', - 'subscription canceled' - ); - await this.policy.reconcileOwnedWorkspaces(userId); - } + await this.policy.reconcileOwnedWorkspaces(userId); break; } default: break; } } + + @OnEvent('entitlement.changed') + async onEntitlementChanged({ + targetType, + targetId, + }: Events['entitlement.changed']) { + if (targetType !== 'workspace') { + return; + } + + const state = await this.quotaState.reconcileWorkspaceQuotaState(targetId); + if (state.plan !== 'team' && state.plan !== 'selfhost_team') { + return; + } + + this.event.emit('workspace.members.allocateSeats', { + workspaceId: targetId, + quantity: state.seatLimit ?? 0, + }); + } } diff --git a/packages/backend/server/src/plugins/payment/index.ts b/packages/backend/server/src/plugins/payment/index.ts index bca6936c30..5940d31777 100644 --- a/packages/backend/server/src/plugins/payment/index.ts +++ b/packages/backend/server/src/plugins/payment/index.ts @@ -3,6 +3,7 @@ import './config'; import { Module } from '@nestjs/common'; import { ServerConfigModule } from '../../core'; +import { EntitlementModule } from '../../core/entitlement'; import { FeatureModule } from '../../core/features'; import { MailModule } from '../../core/mail'; import { PermissionModule } from '../../core/permission'; @@ -41,6 +42,7 @@ import { StripeWebhook } from './webhook'; WorkspaceModule, MailModule, ServerConfigModule, + EntitlementModule, ], providers: [ StripeFactory, diff --git a/packages/backend/server/src/plugins/payment/manager/user.ts b/packages/backend/server/src/plugins/payment/manager/user.ts index e0647644b3..38bf257613 100644 --- a/packages/backend/server/src/plugins/payment/manager/user.ts +++ b/packages/backend/server/src/plugins/payment/manager/user.ts @@ -19,6 +19,7 @@ import { TooManyRequest, URLHelper, } from '../../../base'; +import { EntitlementService } from '../../../core/entitlement'; import { EarlyAccessType, FeatureService } from '../../../core/features'; import { StripeFactory } from '../stripe'; import { @@ -64,7 +65,8 @@ export class UserSubscriptionManager extends SubscriptionManager { private readonly feature: FeatureService, private readonly event: EventBus, private readonly url: URLHelper, - private readonly mutex: Mutex + private readonly mutex: Mutex, + private readonly entitlement: EntitlementService ) { super(stripeProvider, db); } @@ -254,7 +256,7 @@ export class UserSubscriptionManager extends SubscriptionManager { const subscriptionData = this.transformSubscription(subscription); - return this.db.subscription.upsert({ + const saved = await this.db.subscription.upsert({ where: { stripeSubscriptionId: stripeSubscription.id, }, @@ -270,6 +272,8 @@ export class UserSubscriptionManager extends SubscriptionManager { ...omit(subscriptionData, ['provider', 'iapStore']), }, }); + await this.entitlement.upsertFromCloudSubscription(saved); + return saved; } async deleteStripeSubscription({ @@ -285,6 +289,11 @@ export class UserSubscriptionManager extends SubscriptionManager { }); if (result.count > 0) { + await this.entitlement.revokeCloudSubscription({ + targetId: userId, + plan: lookupKey.plan, + stripeSubscriptionId: stripeSubscription.id, + }); this.event.emit('user.subscription.canceled', { userId, plan: lookupKey.plan, @@ -416,7 +425,7 @@ export class UserSubscriptionManager extends SubscriptionManager { if (prevSubscription) { if (prevSubscription.stripeSubscriptionId) { - await this.db.subscription.update({ + const subscription = await this.db.subscription.update({ where: { id: prevSubscription.id, }, @@ -431,6 +440,7 @@ export class UserSubscriptionManager extends SubscriptionManager { nextBillAt: null, }, }); + await this.entitlement.upsertFromCloudSubscription(subscription); await this.stripe.subscriptions.cancel( prevSubscription.stripeSubscriptionId, @@ -440,7 +450,7 @@ export class UserSubscriptionManager extends SubscriptionManager { ); } } else { - await this.db.subscription.create({ + const subscription = await this.db.subscription.create({ data: { targetId: knownInvoice.userId, stripeSubscriptionId: null, @@ -452,6 +462,7 @@ export class UserSubscriptionManager extends SubscriptionManager { nextBillAt: null, }, }); + await this.entitlement.upsertFromCloudSubscription(subscription); } this.event.emit('user.subscription.activated', { @@ -552,6 +563,10 @@ export class UserSubscriptionManager extends SubscriptionManager { plan: lookupKey.plan, recurring: lookupKey.recurring, }); + await this.entitlement.upsertFromCloudSubscription({ + ...subscription, + targetId: userId, + }); return subscription; } @@ -582,6 +597,12 @@ export class UserSubscriptionManager extends SubscriptionManager { canceledAt: new Date(), }, }); + await this.entitlement.revokeCloudSubscription({ + targetId: userId, + plan: lookupKey.plan, + subscriptionId: subscription.id, + stripeSubscriptionId: subscription.stripeSubscriptionId, + }); this.event.emit('user.subscription.canceled', { userId, @@ -621,7 +642,7 @@ export class UserSubscriptionManager extends SubscriptionManager { } if (subscription) { - await this.db.subscription.update({ + const saved = await this.db.subscription.update({ where: { id: subscription.id }, data: { status: SubscriptionStatus.Active, @@ -631,8 +652,9 @@ export class UserSubscriptionManager extends SubscriptionManager { end, }, }); + await this.entitlement.upsertFromCloudSubscription(saved); } else { - await this.db.subscription.create({ + const saved = await this.db.subscription.create({ data: { targetId: userId, stripeSubscriptionId: null, @@ -643,6 +665,7 @@ export class UserSubscriptionManager extends SubscriptionManager { nextBillAt: null, }, }); + await this.entitlement.upsertFromCloudSubscription(saved); } this.event.emit('user.subscription.activated', { diff --git a/packages/backend/server/src/plugins/payment/manager/workspace.ts b/packages/backend/server/src/plugins/payment/manager/workspace.ts index 6501ed740b..30ad1320e0 100644 --- a/packages/backend/server/src/plugins/payment/manager/workspace.ts +++ b/packages/backend/server/src/plugins/payment/manager/workspace.ts @@ -10,6 +10,7 @@ import { SubscriptionPlanNotFound, URLHelper, } from '../../../base'; +import { EntitlementService } from '../../../core/entitlement'; import { Models } from '../../../models'; import { StripeFactory } from '../stripe'; import { @@ -50,7 +51,8 @@ export class WorkspaceSubscriptionManager extends SubscriptionManager { db: PrismaClient, private readonly url: URLHelper, private readonly event: EventBus, - private readonly models: Models + private readonly models: Models, + private readonly entitlement: EntitlementService ) { super(stripeProvider, db); } @@ -155,7 +157,7 @@ export class WorkspaceSubscriptionManager extends SubscriptionManager { }); } - return this.db.subscription.upsert({ + const saved = await this.db.subscription.upsert({ where: { provider: Provider.stripe, stripeSubscriptionId: stripeSubscription.id, @@ -175,6 +177,8 @@ export class WorkspaceSubscriptionManager extends SubscriptionManager { ...omit(subscriptionData, 'provider', 'iapStore'), }, }); + await this.entitlement.upsertFromCloudSubscription(saved); + return saved; } async deleteStripeSubscription({ @@ -194,6 +198,11 @@ export class WorkspaceSubscriptionManager extends SubscriptionManager { }); if (result.count > 0) { + await this.entitlement.revokeCloudSubscription({ + targetId: workspaceId, + plan: lookupKey.plan, + stripeSubscriptionId: stripeSubscription.id, + }); this.event.emit('workspace.subscription.canceled', { workspaceId, plan: lookupKey.plan, diff --git a/packages/backend/server/src/plugins/payment/resolver.ts b/packages/backend/server/src/plugins/payment/resolver.ts index d6a93f402f..b09ef50f57 100644 --- a/packages/backend/server/src/plugins/payment/resolver.ts +++ b/packages/backend/server/src/plugins/payment/resolver.ts @@ -27,7 +27,7 @@ import { WorkspaceIdRequiredToUpdateTeamSubscription, } from '../../base'; import { CurrentUser, Public } from '../../core/auth'; -import { AccessController } from '../../core/permission'; +import { PermissionAccess } from '../../core/permission'; import { UserType } from '../../core/user'; import { WorkspaceType } from '../../core/workspaces'; import { Invoice, Subscription, WorkspaceSubscriptionManager } from './manager'; @@ -665,7 +665,7 @@ export class WorkspaceSubscriptionResolver { constructor( private readonly service: WorkspaceSubscriptionManager, private readonly db: PrismaClient, - private readonly ac: AccessController + private readonly ac: PermissionAccess ) {} @ResolveField(() => SubscriptionType, { diff --git a/packages/backend/server/src/plugins/payment/revenuecat/webhook.ts b/packages/backend/server/src/plugins/payment/revenuecat/webhook.ts index 17f7b96792..b6dabe91cd 100644 --- a/packages/backend/server/src/plugins/payment/revenuecat/webhook.ts +++ b/packages/backend/server/src/plugins/payment/revenuecat/webhook.ts @@ -11,6 +11,7 @@ import { OnJob, sleep, } from '../../../base'; +import { EntitlementService } from '../../../core/entitlement'; import { SubscriptionPlan, SubscriptionRecurring, @@ -32,7 +33,8 @@ export class RevenueCatWebhookHandler { private readonly db: PrismaClient, private readonly config: Config, private readonly event: EventBus, - private readonly queue: JobQueue + private readonly queue: JobQueue, + private readonly entitlement: EntitlementService ) {} @OnEvent('revenuecat.webhook') @@ -197,6 +199,10 @@ export class RevenueCatWebhookHandler { }, }); if (result.count > 0) { + await this.entitlement.revokeCloudSubscription({ + targetId: appUserId, + plan: mapping.plan, + }); this.event.emit('user.subscription.canceled', { userId: appUserId, plan: mapping.plan, @@ -207,7 +213,7 @@ export class RevenueCatWebhookHandler { continue; } - await this.db.subscription.upsert({ + const saved = await this.db.subscription.upsert({ where: { targetId_plan: { targetId: appUserId, plan: mapping.plan }, }, @@ -252,6 +258,7 @@ export class RevenueCatWebhookHandler { trialEnd: null, }, }); + await this.entitlement.upsertFromCloudSubscription(saved); if ( status === SubscriptionStatus.Active || @@ -278,6 +285,12 @@ export class RevenueCatWebhookHandler { if (toBeCleanup.length) { for (const sub of toBeCleanup) { await this.db.subscription.deleteMany({ where: { id: sub.id } }); + await this.entitlement.revokeCloudSubscription({ + targetId: appUserId, + plan: sub.plan as SubscriptionPlan, + subscriptionId: sub.id, + stripeSubscriptionId: sub.stripeSubscriptionId, + }); this.event.emit('user.subscription.canceled', { userId: appUserId, plan: sub.plan as SubscriptionPlan, diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index ac1e4e47c1..e882eb0167 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -129,7 +129,6 @@ input AdminUpdateWorkspaceInput { enableDocEmbedding: Boolean enableSharing: Boolean enableUrlPreview: Boolean - features: [FeatureType!] id: String! name: String public: Boolean @@ -874,6 +873,7 @@ type DocPermissions { Doc_Comments_Delete: Boolean! Doc_Comments_Read: Boolean! Doc_Comments_Resolve: Boolean! + Doc_Comments_Update: Boolean! Doc_Copy: Boolean! Doc_Delete: Boolean! Doc_Duplicate: Boolean! @@ -1562,9 +1562,8 @@ type Mutation { """Update workspace embedding files""" addWorkspaceEmbeddingFiles(blob: Upload!, workspaceId: String!): CopilotWorkspaceFile! - addWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Boolean! - """Update workspace flags and features for admin""" + """Update workspace flags for admin""" adminUpdateWorkspace(input: AdminUpdateWorkspaceInput!): AdminWorkspace approveMember(userId: String!, workspaceId: String!): Boolean! @@ -1633,6 +1632,7 @@ type Mutation { forkCopilotSession(options: ForkChatSessionInput!): String! generateLicenseKey(sessionId: String!): String! generateUserAccessToken(input: GenerateAccessTokenInput!): RevealedAccessToken! + grantCommercialEntitlement(plan: String!, quantity: Int, targetId: String!, targetType: String!): Boolean! grantDocUserRoles(input: GrantDocUserRolesInput!): Boolean! grantMember(permission: Permission!, userId: String!, workspaceId: String!): Boolean! @@ -1680,7 +1680,6 @@ type Mutation { """Remove workspace embedding files""" removeWorkspaceEmbeddingFiles(fileId: String!, workspaceId: String!): Boolean! - removeWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Boolean! reorderWorkspaceByokConfigs(input: ReorderWorkspaceByokConfigsInput!): [WorkspaceByokKeyConfigType!]! """Request to apply the subscription in advance""" @@ -1690,6 +1689,7 @@ type Mutation { resolveComment(input: CommentResolveInput!): Boolean! resumeSubscription(idempotencyKey: String @deprecated(reason: "use header `Idempotency-Key`"), plan: SubscriptionPlan = Pro, workspaceId: String): SubscriptionType! retryTranscriptTask(taskId: String!, workspaceId: String!): TranscriptionResultType + revokeCommercialEntitlement(targetId: String!, targetType: String!): Boolean! revokeDocUserRoles(input: RevokeDocUserRoleInput!): Boolean! revokeInviteLink(workspaceId: String!): Boolean! revokeMember(userId: String!, workspaceId: String!): Boolean! diff --git a/packages/common/graphql/src/schema.ts b/packages/common/graphql/src/schema.ts index bd50352535..2835f15ca7 100644 --- a/packages/common/graphql/src/schema.ts +++ b/packages/common/graphql/src/schema.ts @@ -172,7 +172,6 @@ export interface AdminUpdateWorkspaceInput { enableDocEmbedding?: InputMaybe; enableSharing?: InputMaybe; enableUrlPreview?: InputMaybe; - features?: InputMaybe>; id: Scalars['String']['input']; name?: InputMaybe; public?: InputMaybe; @@ -1021,6 +1020,7 @@ export interface DocPermissions { Doc_Comments_Delete: Scalars['Boolean']['output']; Doc_Comments_Read: Scalars['Boolean']['output']; Doc_Comments_Resolve: Scalars['Boolean']['output']; + Doc_Comments_Update: Scalars['Boolean']['output']; Doc_Copy: Scalars['Boolean']['output']; Doc_Delete: Scalars['Boolean']['output']; Doc_Duplicate: Scalars['Boolean']['output']; @@ -1779,8 +1779,7 @@ export interface Mutation { addContextFile: CopilotContextFile; /** Update workspace embedding files */ addWorkspaceEmbeddingFiles: CopilotWorkspaceFile; - addWorkspaceFeature: Scalars['Boolean']['output']; - /** Update workspace flags and features for admin */ + /** Update workspace flags for admin */ adminUpdateWorkspace: Maybe; approveMember: Scalars['Boolean']['output']; /** Ban an user */ @@ -1836,6 +1835,7 @@ export interface Mutation { forkCopilotSession: Scalars['String']['output']; generateLicenseKey: Scalars['String']['output']; generateUserAccessToken: RevealedAccessToken; + grantCommercialEntitlement: Scalars['Boolean']['output']; grantDocUserRoles: Scalars['Boolean']['output']; grantMember: Scalars['Boolean']['output']; /** import users */ @@ -1871,7 +1871,6 @@ export interface Mutation { removeContextFile: Scalars['Boolean']['output']; /** Remove workspace embedding files */ removeWorkspaceEmbeddingFiles: Scalars['Boolean']['output']; - removeWorkspaceFeature: Scalars['Boolean']['output']; reorderWorkspaceByokConfigs: Array; /** Request to apply the subscription in advance */ requestApplySubscription: Array; @@ -1879,6 +1878,7 @@ export interface Mutation { resolveComment: Scalars['Boolean']['output']; resumeSubscription: SubscriptionType; retryTranscriptTask: Maybe; + revokeCommercialEntitlement: Scalars['Boolean']['output']; revokeDocUserRoles: Scalars['Boolean']['output']; revokeInviteLink: Scalars['Boolean']['output']; revokeMember: Scalars['Boolean']['output']; @@ -1966,11 +1966,6 @@ export interface MutationAddWorkspaceEmbeddingFilesArgs { workspaceId: Scalars['String']['input']; } -export interface MutationAddWorkspaceFeatureArgs { - feature: FeatureType; - workspaceId: Scalars['String']['input']; -} - export interface MutationAdminUpdateWorkspaceArgs { input: AdminUpdateWorkspaceInput; } @@ -2127,6 +2122,13 @@ export interface MutationGenerateUserAccessTokenArgs { input: GenerateAccessTokenInput; } +export interface MutationGrantCommercialEntitlementArgs { + plan: Scalars['String']['input']; + quantity?: InputMaybe; + targetId: Scalars['String']['input']; + targetType: Scalars['String']['input']; +} + export interface MutationGrantDocUserRolesArgs { input: GrantDocUserRolesInput; } @@ -2219,11 +2221,6 @@ export interface MutationRemoveWorkspaceEmbeddingFilesArgs { workspaceId: Scalars['String']['input']; } -export interface MutationRemoveWorkspaceFeatureArgs { - feature: FeatureType; - workspaceId: Scalars['String']['input']; -} - export interface MutationReorderWorkspaceByokConfigsArgs { input: ReorderWorkspaceByokConfigsInput; } @@ -2247,6 +2244,11 @@ export interface MutationRetryTranscriptTaskArgs { workspaceId: Scalars['String']['input']; } +export interface MutationRevokeCommercialEntitlementArgs { + targetId: Scalars['String']['input']; + targetType: Scalars['String']['input']; +} + export interface MutationRevokeDocUserRolesArgs { input: RevokeDocUserRoleInput; } diff --git a/packages/common/realtime/src/index.ts b/packages/common/realtime/src/index.ts index d7312788bd..ea8c385255 100644 --- a/packages/common/realtime/src/index.ts +++ b/packages/common/realtime/src/index.ts @@ -36,14 +36,56 @@ export interface RealtimeRequestMap { }; 'user.quota-state.get': { input: Record; - output: { state: unknown }; + output: { state: UserQuotaStateSnapshot }; }; 'workspace.quota-state.get': { input: { workspaceId: string }; - output: { state: unknown }; + output: { state: WorkspaceQuotaStateSnapshot }; }; } +export interface UserQuotaStateSnapshot { + userId: string; + plan: string; + sourceEntitlementId: string | null; + blobLimit: number; + storageQuota: number; + usedStorageQuota: number; + historyPeriodSeconds: number; + copilotActionLimit: number | null; + flags: Record; + known: boolean; + stale: boolean; + lastReconciledAt: string | Date | null; + staleAfter: string | Date | null; + createdAt: string | Date; + updatedAt: string | Date; +} + +export interface WorkspaceQuotaStateSnapshot { + workspaceId: string; + plan: string; + sourceEntitlementId: string | null; + ownerUserId: string | null; + usesOwnerQuota: boolean; + seatLimit: number; + memberCount: number; + overcapacityMemberCount: number; + blobLimit: number; + storageQuota: number; + usedStorageQuota: number; + historyPeriodSeconds: number; + readonly: boolean; + readonlyReasons: string[]; + flags: Record; + known: boolean; + stale: boolean; + lastReconciledAt: string | Date | null; + staleAfter: string | Date | null; + createdAt: string | Date; + updatedAt: string | Date; +} + export type NotificationCountChangedReason = | 'created' | 'read' diff --git a/packages/frontend/admin/src/config.json b/packages/frontend/admin/src/config.json index 630c564141..95930fc08d 100644 --- a/packages/frontend/admin/src/config.json +++ b/packages/frontend/admin/src/config.json @@ -171,6 +171,18 @@ "desc": "The minimum time interval in milliseconds of creating a new history snapshot when doc get updated." } }, + "permission": { + "readModel": { + "type": "String", + "desc": "Permission data source for Rust evaluation", + "env": "AFFINE_PERMISSION_READ_MODEL" + }, + "fallbackLegacyLoader": { + "type": "Boolean", + "desc": "Fallback from projection loader to legacy loader when projection input loading fails", + "env": "AFFINE_PERMISSION_FALLBACK_LEGACY_LOADER" + } + }, "storages": { "avatar.publicPath": { "type": "String", diff --git a/packages/frontend/admin/src/modules/settings/config.ts b/packages/frontend/admin/src/modules/settings/config.ts index 49fa7e27d3..e44ce454a3 100644 --- a/packages/frontend/admin/src/modules/settings/config.ts +++ b/packages/frontend/admin/src/modules/settings/config.ts @@ -150,7 +150,6 @@ export const KNOWN_CONFIG_GROUPS = [ module: 'copilot', fields: [ 'enabled', - 'scenarios', 'providers.openai', 'providers.gemini', 'providers.anthropic', diff --git a/packages/frontend/admin/src/modules/workspaces/components/workspace-panel.tsx b/packages/frontend/admin/src/modules/workspaces/components/workspace-panel.tsx index d571557aed..cffd51e3f8 100644 --- a/packages/frontend/admin/src/modules/workspaces/components/workspace-panel.tsx +++ b/packages/frontend/admin/src/modules/workspaces/components/workspace-panel.tsx @@ -7,7 +7,6 @@ import { Input } from '@affine/admin/components/ui/input'; import { Label } from '@affine/admin/components/ui/label'; import { Separator } from '@affine/admin/components/ui/separator'; import { Switch } from '@affine/admin/components/ui/switch'; -import type { FeatureType } from '@affine/graphql'; import { adminUpdateWorkspaceMutation, adminWorkspaceQuery, @@ -17,10 +16,8 @@ import { AccountIcon } from '@blocksuite/icons/rc'; import { useCallback, useEffect, useMemo, useState } from 'react'; import { toast } from 'sonner'; -import { FeatureToggleList } from '../../../components/shared/feature-toggle-list'; import { useMutateQueryResource, useMutation } from '../../../use-mutation'; import { useQuery } from '../../../use-query'; -import { useServerConfig } from '../../common'; import { RightPanelHeader } from '../../header'; import { useRightPanel } from '../../panel/context'; import type { WorkspaceDetail } from '../schema'; @@ -69,7 +66,6 @@ function WorkspacePanelContent({ workspace: WorkspaceDetail; onClose: () => void; }) { - const serverConfig = useServerConfig(); const { setHasDirtyChanges } = useRightPanel(); const revalidate = useMutateQueryResource(); const { trigger: updateWorkspace, isMutating } = useMutation({ @@ -78,7 +74,6 @@ function WorkspacePanelContent({ const normalizedWorkspace = useMemo( () => ({ - features: [...workspace.features], flags: { public: workspace.public, enableAi: workspace.enableAi, @@ -91,14 +86,10 @@ function WorkspacePanelContent({ [workspace] ); - const [featureSelection, setFeatureSelection] = useState( - normalizedWorkspace.features - ); const [flags, setFlags] = useState(normalizedWorkspace.flags); const [baseline, setBaseline] = useState(normalizedWorkspace); useEffect(() => { - setFeatureSelection(normalizedWorkspace.features); setFlags(normalizedWorkspace.flags); setBaseline(normalizedWorkspace); }, [normalizedWorkspace]); @@ -110,20 +101,14 @@ function WorkspacePanelContent({ flags.enableSharing !== baseline.flags.enableSharing || flags.enableUrlPreview !== baseline.flags.enableUrlPreview || flags.enableDocEmbedding !== baseline.flags.enableDocEmbedding || - flags.name !== baseline.flags.name || - featureSelection.length !== baseline.features.length || - featureSelection.some(f => !baseline.features.includes(f)) + flags.name !== baseline.flags.name ); - }, [baseline, featureSelection, flags]); + }, [baseline, flags]); useEffect(() => { setHasDirtyChanges(hasChanges); }, [hasChanges, setHasDirtyChanges]); - const handleFeaturesChange = useCallback((features: FeatureType[]) => { - setFeatureSelection(features); - }, []); - const handleSave = useCallback(() => { const update = async () => { try { @@ -136,7 +121,6 @@ function WorkspacePanelContent({ enableUrlPreview: flags.enableUrlPreview, enableDocEmbedding: flags.enableDocEmbedding, name: flags.name || null, - features: featureSelection, }, }); await Promise.all([ @@ -146,7 +130,6 @@ function WorkspacePanelContent({ toast.success('Workspace updated successfully'); setBaseline({ flags: { ...flags }, - features: [...featureSelection], }); setHasDirtyChanges(false); onClose(); @@ -156,7 +139,6 @@ function WorkspacePanelContent({ }; update().catch(() => {}); }, [ - featureSelection, flags, onClose, revalidate, @@ -239,18 +221,6 @@ function WorkspacePanelContent({ /> -
-
Features
- -
-
= nil, enableSharing: GraphQLNullable = nil, enableUrlPreview: GraphQLNullable = nil, - features: GraphQLNullable<[GraphQLEnum]> = nil, id: String, name: GraphQLNullable = nil, `public`: GraphQLNullable = nil @@ -27,7 +26,6 @@ public struct AdminUpdateWorkspaceInput: InputObject { "enableDocEmbedding": enableDocEmbedding, "enableSharing": enableSharing, "enableUrlPreview": enableUrlPreview, - "features": features, "id": id, "name": name, "public": `public` @@ -59,11 +57,6 @@ public struct AdminUpdateWorkspaceInput: InputObject { set { __data["enableUrlPreview"] = newValue } } - public var features: GraphQLNullable<[GraphQLEnum]> { - get { __data["features"] } - set { __data["features"] = newValue } - } - public var id: String { get { __data["id"] } set { __data["id"] = newValue } diff --git a/packages/frontend/core/src/modules/cloud/entities/user-quota.spec.ts b/packages/frontend/core/src/modules/cloud/entities/user-quota.spec.ts new file mode 100644 index 0000000000..653ba3c75d --- /dev/null +++ b/packages/frontend/core/src/modules/cloud/entities/user-quota.spec.ts @@ -0,0 +1,126 @@ +import type { GetCurrentUserProfileQuery } from '@affine/graphql'; +import type { UserQuotaStateSnapshot } from '@affine/realtime'; +import { Framework } from '@toeverything/infra'; +import { Subject } from 'rxjs'; +import { describe, expect, test, vi } from 'vitest'; + +import { AuthService } from '../services/auth'; +import { UserQuotaStore } from '../stores/user-quota'; +import { UserQuota } from './user-quota'; + +type Quota = NonNullable['quota']; + +const authService = { + session: { + ['account$']: { + value: { id: 'user-1' }, + }, + }, +} as unknown as AuthService; + +function createQuotaState( + overrides: Partial = {} +): UserQuotaStateSnapshot { + return { + userId: 'user-1', + plan: 'pro', + sourceEntitlementId: 'entitlement-1', + blobLimit: 1024, + storageQuota: 2048, + usedStorageQuota: 512, + historyPeriodSeconds: 30 * 24 * 60 * 60, + copilotActionLimit: null, + flags: {}, + known: true, + stale: false, + lastReconciledAt: null, + staleAfter: null, + createdAt: new Date(0), + updatedAt: new Date(0), + ...overrides, + }; +} + +function createQuota(overrides: Partial = {}): Quota { + return { + name: 'Legacy', + blobLimit: 1024, + storageQuota: 2048, + historyPeriod: 30 * 24 * 60 * 60, + memberLimit: 8, + humanReadable: { + name: 'Legacy', + blobLimit: '1 KB', + storageQuota: '2 KB', + historyPeriod: '30 days', + memberLimit: '8', + }, + ...overrides, + }; +} + +function createStore( + overrides: Partial, + eventSubject = new Subject<{ type: 'ready' } | { changed: true }>() +) { + return { + fetchUserQuotaState: vi.fn(), + subscribeUserQuotaState: vi.fn(() => eventSubject), + fetchUserQuota: vi.fn(), + ...overrides, + } as unknown as UserQuotaStore; +} + +function createEntity(store: UserQuotaStore) { + const framework = new Framework(); + framework + .service(AuthService, authService) + .store(UserQuotaStore, store) + .entity(UserQuota, [AuthService, UserQuotaStore]); + + return framework.provider().createEntity(UserQuota); +} + +describe('UserQuota', () => { + test('uses realtime quota state snapshots and refreshes on quota events', async () => { + const events$ = new Subject<{ type: 'ready' } | { changed: true }>(); + const store = createStore( + { + fetchUserQuotaState: vi + .fn() + .mockResolvedValueOnce(createQuotaState({ usedStorageQuota: 512 })) + .mockResolvedValueOnce(createQuotaState({ usedStorageQuota: 768 })), + }, + events$ + ); + const quota = createEntity(store); + + quota.revalidate(); + await vi.waitFor(() => expect(quota.used$.value).toBe(512)); + expect(quota.quota$.value?.humanReadable.historyPeriod).toBe('30 days'); + + events$.next({ changed: true }); + await vi.waitFor(() => expect(quota.used$.value).toBe(768)); + + expect(store.fetchUserQuotaState).toHaveBeenCalledTimes(2); + expect(store.fetchUserQuota).not.toHaveBeenCalled(); + quota.dispose(); + }); + + test('falls back to legacy GraphQL quota when realtime request fails', async () => { + const store = createStore({ + fetchUserQuotaState: vi.fn().mockRejectedValue(new Error('offline')), + fetchUserQuota: vi.fn().mockResolvedValue({ + quota: createQuota(), + used: 256, + }), + }); + const quota = createEntity(store); + + quota.revalidate(); + + await vi.waitFor(() => expect(quota.quota$.value?.name).toBe('Legacy')); + expect(quota.used$.value).toBe(256); + quota.dispose(); + }); +}); diff --git a/packages/frontend/core/src/modules/cloud/entities/user-quota.ts b/packages/frontend/core/src/modules/cloud/entities/user-quota.ts index 51ea9bea2d..c1673adaa2 100644 --- a/packages/frontend/core/src/modules/cloud/entities/user-quota.ts +++ b/packages/frontend/core/src/modules/cloud/entities/user-quota.ts @@ -1,26 +1,74 @@ -import type { QuotaQuery } from '@affine/graphql'; -import { - catchErrorInto, - effect, - Entity, - exhaustMapSwitchUntilChanged, - fromPromise, - LiveData, - onComplete, - onStart, - smartRetry, -} from '@toeverything/infra'; +import type { GetCurrentUserProfileQuery } from '@affine/graphql'; +import type { + RealtimeTopicEventOf, + UserQuotaStateSnapshot, +} from '@affine/realtime'; +import { Entity, LiveData } from '@toeverything/infra'; import { cssVarV2 } from '@toeverything/theme/v2'; import bytes from 'bytes'; -import { map, tap } from 'rxjs'; +import { RealtimeLiveQuery } from '../realtime/live-query'; import type { AuthService } from '../services/auth'; import type { UserQuotaStore } from '../stores/user-quota'; +type QuotaType = NonNullable< + GetCurrentUserProfileQuery['currentUser'] +>['quota']; + +const DAY_SECONDS = 24 * 60 * 60; + +function formatSize(size: number) { + return size === 0 ? '0 B' : (bytes.format(size) ?? '0 B'); +} + +function formatHistoryPeriod(value: number) { + return `${(value / DAY_SECONDS).toFixed(0)} days`; +} + +function userMemberLimit(plan: string) { + return plan === 'pro' || plan === 'lifetime_pro' || plan === 'selfhost_free' + ? 10 + : 3; +} + +function planName(plan: string) { + switch (plan) { + case 'pro': + case 'selfhost_free': + return 'Pro'; + case 'lifetime_pro': + return 'Lifetime Pro'; + case 'ai': + return 'AI'; + default: + return 'Free'; + } +} + +function userQuotaFromState(state: UserQuotaStateSnapshot): QuotaType { + const name = planName(state.plan); + const memberLimit = userMemberLimit(state.plan); + return { + name, + blobLimit: state.blobLimit, + storageQuota: state.storageQuota, + historyPeriod: state.historyPeriodSeconds, + memberLimit, + humanReadable: { + name, + blobLimit: formatSize(state.blobLimit), + storageQuota: formatSize(state.storageQuota), + historyPeriod: formatHistoryPeriod(state.historyPeriodSeconds), + memberLimit: memberLimit.toString(), + }, + }; +} + export class UserQuota extends Entity { - quota$ = new LiveData['quota'] | null>( - null - ); + quota$ = new LiveData(null); + isRevalidating$ = new LiveData(false); + error$ = new LiveData(null); + /** Used storage in bytes */ used$ = new LiveData(null); /** Formatted used storage */ @@ -53,8 +101,17 @@ export class UserQuota extends Entity { : null ); - isRevalidating$ = new LiveData(false); - error$ = new LiveData(null); + private started = false; + private readonly liveQuery = new RealtimeLiveQuery< + { quota: QuotaType; used: number }, + RealtimeTopicEventOf<'user.quota-state.changed'> + >({ + request: signal => this.requestQuota(signal), + subscribe: () => this.store.subscribeUserQuotaState(), + applySnapshot: data => this.applyQuota(data.quota, data.used), + applyEvent: () => 'revalidate', + onError: error => this.error$.next(error), + }); constructor( private readonly authService: AuthService, @@ -63,42 +120,19 @@ export class UserQuota extends Entity { super(); } - revalidate = effect( - map(() => ({ - accountId: this.authService.session.account$.value?.id, - })), - exhaustMapSwitchUntilChanged( - (a, b) => a.accountId === b.accountId, - ({ accountId }) => - fromPromise(async signal => { - if (!accountId) { - return; // no quota if no user - } - const { quota, used } = await this.store.fetchUserQuota(signal); - - return { quota, used }; - }).pipe( - smartRetry(), - tap(data => { - if (data) { - const { quota, used } = data; - this.quota$.next(quota); - this.used$.next(used); - } else { - this.quota$.next(null); - this.used$.next(null); - } - }), - catchErrorInto(this.error$), - onStart(() => this.isRevalidating$.next(true)), - onComplete(() => this.isRevalidating$.next(false)) - ), - () => { - // Reset the state when the user is changed - this.reset(); - } - ) - ); + revalidate = () => { + if (!this.authService.session.account$.value?.id) { + this.liveQuery.stop(); + this.started = false; + this.reset(); + return; + } + if (!this.started) { + this.started = true; + this.liveQuery.start(); + } + this.liveQuery.revalidate(); + }; reset() { this.quota$.next(null); @@ -108,6 +142,28 @@ export class UserQuota extends Entity { } override dispose(): void { - this.revalidate.unsubscribe(); + this.liveQuery.dispose(); + } + + private applyQuota(quota: QuotaType | null, used: number | null) { + this.error$.next(null); + this.quota$.next(quota); + this.used$.next(used); + } + + private async requestQuota(signal: AbortSignal) { + this.isRevalidating$.setValue(true); + try { + const state = await this.store.fetchUserQuotaState(signal); + return { + quota: userQuotaFromState(state), + used: state.usedStorageQuota, + }; + } catch { + const { quota, used } = await this.store.fetchUserQuota(signal); + return { quota, used }; + } finally { + this.isRevalidating$.setValue(false); + } } } diff --git a/packages/frontend/core/src/modules/cloud/index.ts b/packages/frontend/core/src/modules/cloud/index.ts index ec60596ecd..143f812479 100644 --- a/packages/frontend/core/src/modules/cloud/index.ts +++ b/packages/frontend/core/src/modules/cloud/index.ts @@ -165,7 +165,7 @@ export function configureCloudModule(framework: Framework) { .entity(Subscription, [AuthService, ServerService, SubscriptionStore]) .entity(SubscriptionPrices, [ServerService, SubscriptionStore]) .service(UserQuotaService) - .store(UserQuotaStore, [GraphQLService]) + .store(UserQuotaStore, [GraphQLService, NbstoreService]) .entity(UserQuota, [AuthService, UserQuotaStore]) .service(UserCopilotQuotaService) .store(UserCopilotQuotaStore, [GraphQLService]) diff --git a/packages/frontend/core/src/modules/cloud/stores/user-quota.ts b/packages/frontend/core/src/modules/cloud/stores/user-quota.ts index e738a458d5..a3515eb63a 100644 --- a/packages/frontend/core/src/modules/cloud/stores/user-quota.ts +++ b/packages/frontend/core/src/modules/cloud/stores/user-quota.ts @@ -1,13 +1,33 @@ import { getCurrentUserProfileQuery } from '@affine/graphql'; import { Store } from '@toeverything/infra'; +import type { NbstoreService } from '../../storage'; import type { GraphQLService } from '../services/graphql'; export class UserQuotaStore extends Store { - constructor(private readonly graphqlService: GraphQLService) { + constructor( + private readonly graphqlService: GraphQLService, + private readonly nbstoreService: NbstoreService + ) { super(); } + async fetchUserQuotaState(abortSignal?: AbortSignal) { + const response = await this.nbstoreService.realtime.request( + 'user.quota-state.get', + {}, + { signal: abortSignal, timeoutMs: 10000 } + ); + return response.state; + } + + subscribeUserQuotaState() { + return this.nbstoreService.realtime.subscribe( + 'user.quota-state.changed', + {} + ); + } + async fetchUserQuota(abortSignal?: AbortSignal) { const data = await this.graphqlService.gql({ query: getCurrentUserProfileQuery, diff --git a/packages/frontend/core/src/modules/permissions/entities/permission.ts b/packages/frontend/core/src/modules/permissions/entities/permission.ts index b467eadf0e..4f84d22469 100644 --- a/packages/frontend/core/src/modules/permissions/entities/permission.ts +++ b/packages/frontend/core/src/modules/permissions/entities/permission.ts @@ -1,4 +1,3 @@ -import { Permission } from '@affine/graphql'; import { backoffRetry, effect, @@ -46,9 +45,11 @@ export class WorkspacePermission extends Entity { signal ); + const isOwner = info.workspace.permissions.Workspace_Delete; return { - isOwner: info.workspace.role === Permission.Owner, - isAdmin: info.workspace.role === Permission.Admin, + isOwner, + isAdmin: + !isOwner && info.workspace.permissions.Workspace_Settings_Update, isTeam: info.workspace.team, }; } else { diff --git a/packages/frontend/core/src/modules/quota/entities/quota.spec.ts b/packages/frontend/core/src/modules/quota/entities/quota.spec.ts new file mode 100644 index 0000000000..906b142ab2 --- /dev/null +++ b/packages/frontend/core/src/modules/quota/entities/quota.spec.ts @@ -0,0 +1,130 @@ +import { WorkspaceService } from '@affine/core/modules/workspace'; +import type { WorkspaceQuotaQuery } from '@affine/graphql'; +import type { WorkspaceQuotaStateSnapshot } from '@affine/realtime'; +import { Framework } from '@toeverything/infra'; +import { Subject } from 'rxjs'; +import { describe, expect, test, vi } from 'vitest'; + +import { WorkspaceQuotaStore } from '../stores/quota'; +import { WorkspaceQuota } from './quota'; + +type Quota = WorkspaceQuotaQuery['workspace']['quota']; + +const workspaceService = { + workspace: { id: 'workspace-1' }, +} as unknown as WorkspaceService; + +function createQuotaState( + overrides: Partial = {} +): WorkspaceQuotaStateSnapshot { + return { + workspaceId: 'workspace-1', + plan: 'Team', + sourceEntitlementId: 'entitlement-1', + ownerUserId: 'user-1', + usesOwnerQuota: false, + seatLimit: 10, + memberCount: 3, + overcapacityMemberCount: 0, + blobLimit: 1024, + storageQuota: 2048, + usedStorageQuota: 512, + historyPeriodSeconds: 30 * 24 * 60 * 60, + readonly: false, + readonlyReasons: [], + flags: {}, + known: true, + stale: false, + lastReconciledAt: null, + staleAfter: null, + createdAt: new Date(0), + updatedAt: new Date(0), + ...overrides, + }; +} + +function createQuota(overrides: Partial = {}): Quota { + return { + name: 'Legacy', + blobLimit: 1024, + storageQuota: 2048, + usedStorageQuota: 256, + historyPeriod: 30 * 24 * 60 * 60 * 1000, + memberLimit: 8, + memberCount: 2, + overcapacityMemberCount: 0, + humanReadable: { + name: 'Legacy', + blobLimit: '1 KB', + storageQuota: '2 KB', + historyPeriod: '30 days', + memberLimit: '8', + memberCount: '2', + overcapacityMemberCount: '0', + }, + ...overrides, + }; +} + +function createStore( + overrides: Partial, + eventSubject = new Subject<{ type: 'ready' } | { changed: true }>() +) { + return { + fetchWorkspaceQuotaState: vi.fn(), + subscribeWorkspaceQuotaState: vi.fn(() => eventSubject), + fetchWorkspaceQuota: vi.fn(), + ...overrides, + } as unknown as WorkspaceQuotaStore; +} + +function createEntity(store: WorkspaceQuotaStore) { + const framework = new Framework(); + framework + .service(WorkspaceService, workspaceService) + .store(WorkspaceQuotaStore, store) + .entity(WorkspaceQuota, [WorkspaceService, WorkspaceQuotaStore]); + + return framework.provider().createEntity(WorkspaceQuota); +} + +describe('WorkspaceQuota', () => { + test('uses realtime quota state snapshots and refreshes on quota events', async () => { + const events$ = new Subject<{ type: 'ready' } | { changed: true }>(); + const store = createStore( + { + fetchWorkspaceQuotaState: vi + .fn() + .mockResolvedValueOnce(createQuotaState({ memberCount: 3 })) + .mockResolvedValueOnce(createQuotaState({ memberCount: 4 })), + }, + events$ + ); + const quota = createEntity(store); + + quota.revalidate(); + await vi.waitFor(() => expect(quota.quota$.value?.memberCount).toBe(3)); + expect(quota.quota$.value?.humanReadable.historyPeriod).toBe('30 days'); + + events$.next({ changed: true }); + await vi.waitFor(() => expect(quota.quota$.value?.memberCount).toBe(4)); + + expect(store.fetchWorkspaceQuotaState).toHaveBeenCalledTimes(2); + expect(store.fetchWorkspaceQuota).not.toHaveBeenCalled(); + quota.dispose(); + }); + + test('falls back to legacy GraphQL quota when realtime request fails', async () => { + const store = createStore({ + fetchWorkspaceQuotaState: vi.fn().mockRejectedValue(new Error('offline')), + fetchWorkspaceQuota: vi.fn().mockResolvedValue(createQuota()), + }); + const quota = createEntity(store); + + quota.revalidate(); + + await vi.waitFor(() => expect(quota.quota$.value?.name).toBe('Legacy')); + expect(quota.error$.value).toBeNull(); + quota.dispose(); + }); +}); diff --git a/packages/frontend/core/src/modules/quota/entities/quota.ts b/packages/frontend/core/src/modules/quota/entities/quota.ts index 0cac2c2c4f..502e10f0f1 100644 --- a/packages/frontend/core/src/modules/quota/entities/quota.ts +++ b/packages/frontend/core/src/modules/quota/entities/quota.ts @@ -1,31 +1,90 @@ import { DebugLogger } from '@affine/debug'; import type { WorkspaceQuotaQuery } from '@affine/graphql'; -import { - catchErrorInto, - effect, - Entity, - exhaustMapWithTrailing, - fromPromise, - LiveData, - onComplete, - onStart, - smartRetry, -} from '@toeverything/infra'; +import type { + RealtimeTopicEventOf, + WorkspaceQuotaStateSnapshot, +} from '@affine/realtime'; +import { Entity, LiveData } from '@toeverything/infra'; import { cssVarV2 } from '@toeverything/theme/v2'; import bytes from 'bytes'; -import { tap } from 'rxjs'; +import { RealtimeLiveQuery } from '../../cloud/realtime/live-query'; import type { WorkspaceService } from '../../workspace'; import type { WorkspaceQuotaStore } from '../stores/quota'; type QuotaType = WorkspaceQuotaQuery['workspace']['quota']; const logger = new DebugLogger('affine:workspace-permission'); +const DAY_SECONDS = 24 * 60 * 60; + +function formatSize(size: number) { + return size === 0 ? '0 B' : (bytes.format(size) ?? '0 B'); +} + +function formatHistoryPeriod(value: number) { + return `${(value / DAY_SECONDS).toFixed(0)} days`; +} + +function planName(plan: string) { + switch (plan) { + case 'pro': + case 'selfhost_free': + return 'Pro'; + case 'lifetime_pro': + return 'Lifetime Pro'; + case 'ai': + return 'AI'; + case 'team': + case 'selfhost_team': + return 'Team'; + default: + return 'Free'; + } +} + +function workspaceQuotaFromState( + state: WorkspaceQuotaStateSnapshot +): QuotaType { + const name = planName(state.plan); + return { + name, + blobLimit: state.blobLimit, + storageQuota: state.storageQuota, + usedStorageQuota: state.usedStorageQuota, + historyPeriod: state.historyPeriodSeconds, + memberLimit: state.seatLimit, + memberCount: state.memberCount, + overcapacityMemberCount: state.overcapacityMemberCount, + humanReadable: { + name, + blobLimit: formatSize(state.blobLimit), + storageQuota: formatSize(state.storageQuota), + historyPeriod: formatHistoryPeriod(state.historyPeriodSeconds), + memberLimit: state.seatLimit.toString(), + memberCount: state.memberCount.toString(), + overcapacityMemberCount: state.overcapacityMemberCount.toString(), + }, + }; +} export class WorkspaceQuota extends Entity { quota$ = new LiveData(null); isRevalidating$ = new LiveData(false); error$ = new LiveData(null); + private started = false; + private readonly liveQuery = new RealtimeLiveQuery< + QuotaType, + RealtimeTopicEventOf<'workspace.quota-state.changed'> + >({ + request: signal => this.requestQuota(signal), + subscribe: () => + this.store.subscribeWorkspaceQuotaState( + this.workspaceService.workspace.id + ), + applySnapshot: quota => this.applyQuota(quota), + applyEvent: () => 'revalidate', + onError: error => this.handleError(error), + }); /** Used storage in bytes */ used$ = new LiveData(null); @@ -66,34 +125,13 @@ export class WorkspaceQuota extends Entity { super(); } - revalidate = effect( - exhaustMapWithTrailing(() => { - return fromPromise(async signal => { - const data = await this.store.fetchWorkspaceQuota( - this.workspaceService.workspace.id, - signal - ); - return { quota: data, used: data.usedStorageQuota }; - }).pipe( - smartRetry(), - tap(data => { - if (data) { - const { quota, used } = data; - this.quota$.next(quota); - this.used$.next(used); - } else { - this.quota$.next(null); - this.used$.next(null); - } - }), - catchErrorInto(this.error$, error => { - logger.error('Failed to fetch workspace quota', error); - }), - onStart(() => this.isRevalidating$.setValue(true)), - onComplete(() => this.isRevalidating$.setValue(false)) - ); - }) - ); + revalidate = () => { + if (!this.started) { + this.started = true; + this.liveQuery.start(); + } + this.liveQuery.revalidate(); + }; waitForRevalidation(signal?: AbortSignal) { this.revalidate(); @@ -111,6 +149,36 @@ export class WorkspaceQuota extends Entity { } override dispose(): void { - this.revalidate.unsubscribe(); + this.liveQuery.dispose(); + } + + private applyQuota(quota: QuotaType | null) { + this.error$.next(null); + this.quota$.next(quota); + this.used$.next(quota?.usedStorageQuota ?? null); + } + + private handleError(error: unknown) { + logger.error('Failed to fetch workspace quota', error); + this.error$.next(error); + } + + private async requestQuota(signal: AbortSignal) { + this.isRevalidating$.setValue(true); + try { + return workspaceQuotaFromState( + await this.store.fetchWorkspaceQuotaState( + this.workspaceService.workspace.id, + signal + ) + ); + } catch { + return await this.store.fetchWorkspaceQuota( + this.workspaceService.workspace.id, + signal + ); + } finally { + this.isRevalidating$.setValue(false); + } } } diff --git a/packages/frontend/core/src/modules/quota/index.ts b/packages/frontend/core/src/modules/quota/index.ts index db5b5bc2bc..294493c2ab 100644 --- a/packages/frontend/core/src/modules/quota/index.ts +++ b/packages/frontend/core/src/modules/quota/index.ts @@ -4,6 +4,7 @@ export { QuotaCheck } from './views/quota-check'; import { type Framework } from '@toeverything/infra'; import { WorkspaceServerService } from '../cloud'; +import { NbstoreService } from '../storage'; import { WorkspaceScope, WorkspaceService } from '../workspace'; import { WorkspaceQuota } from './entities/quota'; import { WorkspaceQuotaService } from './services/quota'; @@ -13,6 +14,6 @@ export function configureQuotaModule(framework: Framework) { framework .scope(WorkspaceScope) .service(WorkspaceQuotaService) - .store(WorkspaceQuotaStore, [WorkspaceServerService]) + .store(WorkspaceQuotaStore, [WorkspaceServerService, NbstoreService]) .entity(WorkspaceQuota, [WorkspaceService, WorkspaceQuotaStore]); } diff --git a/packages/frontend/core/src/modules/quota/stores/quota.ts b/packages/frontend/core/src/modules/quota/stores/quota.ts index 3de8016fda..595b1aa828 100644 --- a/packages/frontend/core/src/modules/quota/stores/quota.ts +++ b/packages/frontend/core/src/modules/quota/stores/quota.ts @@ -2,11 +2,32 @@ import type { WorkspaceServerService } from '@affine/core/modules/cloud'; import { workspaceQuotaQuery } from '@affine/graphql'; import { Store } from '@toeverything/infra'; +import type { NbstoreService } from '../../storage'; + export class WorkspaceQuotaStore extends Store { - constructor(private readonly workspaceServerService: WorkspaceServerService) { + constructor( + private readonly workspaceServerService: WorkspaceServerService, + private readonly nbstoreService: NbstoreService + ) { super(); } + async fetchWorkspaceQuotaState(workspaceId: string, signal?: AbortSignal) { + const response = await this.nbstoreService.realtime.request( + 'workspace.quota-state.get', + { workspaceId }, + { signal, timeoutMs: 10000 } + ); + return response.state; + } + + subscribeWorkspaceQuotaState(workspaceId: string) { + return this.nbstoreService.realtime.subscribe( + 'workspace.quota-state.changed', + { workspaceId } + ); + } + async fetchWorkspaceQuota(workspaceId: string, signal?: AbortSignal) { if (!this.workspaceServerService.server) { throw new Error('No Server'); diff --git a/packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts b/packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts index e98650c003..1abe526c89 100644 --- a/packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts +++ b/packages/frontend/core/src/modules/workspace-engine/impls/cloud.ts @@ -5,7 +5,6 @@ import { deleteWorkspaceMutation, getWorkspaceInfoQuery, getWorkspacesQuery, - Permission, ServerDeploymentType, ServerFeature, } from '@affine/graphql'; @@ -355,13 +354,12 @@ class CloudWorkspaceFlavourProvider implements WorkspaceFlavourProvider { const info = await this.getWorkspaceInfo(id, signal); + const isOwner = info.workspace.permissions.Workspace_Delete; + const isAdmin = + !isOwner && info.workspace.permissions.Workspace_Settings_Update; + if (!cloudData && !localData) { - return { - isOwner: info.workspace.role === Permission.Owner, - isAdmin: info.workspace.role === Permission.Admin, - isTeam: info.workspace.team, - isEmpty, - }; + return { isOwner, isAdmin, isTeam: info.workspace.team, isEmpty }; } const client = getWorkspaceProfileWorker(); @@ -374,8 +372,8 @@ class CloudWorkspaceFlavourProvider implements WorkspaceFlavourProvider { return { name: result.name, avatar: result.avatar, - isOwner: info.workspace.role === Permission.Owner, - isAdmin: info.workspace.role === Permission.Admin, + isOwner, + isAdmin, isTeam: info.workspace.team, isEmpty, }; diff --git a/tests/kit/src/utils/cloud.ts b/tests/kit/src/utils/cloud.ts index 35bd4cc381..8bfcd2510b 100644 --- a/tests/kit/src/utils/cloud.ts +++ b/tests/kit/src/utils/cloud.ts @@ -179,7 +179,7 @@ export async function createRandomAIUser(): Promise<{ password: '123456', }; const result = await runPrisma(async client => { - await client.user.create({ + const created = await client.user.create({ data: { ...user, emailVerifiedAt: new Date(), @@ -203,11 +203,21 @@ export async function createRandomAIUser(): Promise<{ }, }); - return await client.user.findUnique({ - where: { - email: user.email, + await client.entitlement.create({ + data: { + targetType: 'user', + targetId: created.id, + source: 'cloud_subscription', + plan: 'ai', + status: 'active', + subjectId: `test-ai:${created.id}`, + metadata: { + legacySync: false, + }, }, }); + + return created; }); cloudUserSchema.parse(result); return {