feat(server): entitlement based model (#14996)

#### PR Dependency Tree


* **PR #14996** 👈

This tree was auto-generated by
[Charcoal](https://github.com/danerwilliams/charcoal)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## 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_stack_entry_start -->

[![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)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
This commit is contained in:
DarkSky
2026-05-19 22:48:05 +08:00
committed by GitHub
parent 103ad2a810
commit c53457691d
150 changed files with 13463 additions and 2458 deletions
@@ -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<Context>;
const test = ava.serial as TestFn<Context>;
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<Context>) {
@@ -59,6 +77,73 @@ function workspaceHash(workspaceId: string) {
return createHash('sha256').update(workspaceId).digest('hex').slice(0, 12);
}
async function grantUserPlan(
t: ExecutionContext<Context>,
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<Context>,
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<Context>,
workspaceId: string
) {
await t.context.entitlement.upsertFromCloudSubscription({
targetId: workspaceId,
plan: SubscriptionPlan.Team,
recurring: SubscriptionRecurring.Yearly,
status: SubscriptionStatus.Active,
});
}
async function revokeTeamPlan(
t: ExecutionContext<Context>,
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(
@@ -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(
@@ -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<ReturnType<QuotaStateService['reconcileUserQuotaState']>>);
};
// payment disabled -> allow requested if in optional; pro not blocked
@@ -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, {
@@ -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);
@@ -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',
@@ -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,
/<meta name="apple-itunes-app" content="app-id=6736937980" \/>/
);
t.notRegex(
res.text,
/<meta name="apple-itunes-app" content="app-id=6736937980" \/>/
);
} finally {
// @ts-expect-error restore mutable test env singleton
globalThis.env.DEPLOYMENT_TYPE = previousDeploymentType;
}
}
);
@@ -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 => {
@@ -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');
@@ -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 => {
@@ -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,
@@ -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)
@@ -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: {
@@ -45,6 +45,55 @@ export class MockWorkspace extends Mocker<MockWorkspaceInput, MockedWorkspace> {
: 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) {
@@ -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';
@@ -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,
@@ -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;
@@ -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<void>) => {
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,
});
});
@@ -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);
@@ -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<string, never>;
const test = ava as TestFn<Context>;
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,
},
},
]);
});
@@ -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();
@@ -9,7 +9,7 @@ import type { TestingApp } from './utils';
type TestContext = {
app: TestingApp;
};
const test = ava as TestFn<TestContext>;
const test = ava.serial as TestFn<TestContext>;
let safeFetchStub: Sinon.SinonStub | undefined;
let safeFetchHandler:
@@ -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));
@@ -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;
@@ -62,6 +62,7 @@ export type KnownMetricScopes =
| 'queue'
| 'storage'
| 'process'
| 'permission'
| 'workspace';
const metricCreators: MetricCreators = {
@@ -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
) {}
@@ -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
@@ -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 [];
}
}
@@ -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<RenderOptions | null> {
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<RenderOptions | null> {
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);
@@ -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) => {
@@ -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<Context>;
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);
});
@@ -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<Context>;
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;
}
}
);
@@ -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<Context>;
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');
});
@@ -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 };
@@ -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<string, unknown>) {
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;
}
}
@@ -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<string, Entitlement>();
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';
}
}
File diff suppressed because it is too large Load Diff
@@ -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,
@@ -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;
}
}
@@ -5,24 +5,14 @@ import { Feature, UserFeatureName } from '../../models';
@Injectable()
export class AvailableUserFeatureConfig {
availableUserFeatures(): Set<UserFeatureName> {
return new Set([
Feature.Admin,
Feature.UnlimitedCopilot,
Feature.EarlyAccess,
Feature.AIEarlyAccess,
]);
return new Set([Feature.Admin, Feature.EarlyAccess, Feature.AIEarlyAccess]);
}
configurableUserFeatures(): Set<UserFeatureName> {
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]
);
}
}
@@ -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, {
@@ -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',
@@ -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>(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<void>;
@@ -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']);
@@ -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)
@@ -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<Context>;
const READONLY_FEATURE = 'quota_exceeded_readonly_workspace_v1' as const;
type WorkspaceQuotaSnapshot = Awaited<
ReturnType<QuotaService['getWorkspaceQuotaWithUsage']>
ReturnType<QuotaStateService['reconcileWorkspaceQuotaState']>
> & {
ownerQuota?: string;
readonlyReasons: string[];
};
const readonlyWorkspaceState = (
workspaceId: string,
readonlyReasons: string[],
overrides: Partial<WorkspaceQuotaSnapshot> = {}
) =>
({
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<Array<{ current: boolean }>>`
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,
File diff suppressed because it is too large Load Diff
@@ -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>(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']);
@@ -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<T[]> {
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<WorkspaceAction, boolean>,
};
}
}
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<DocAction, boolean>,
};
}
}
@@ -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'],
},
});
@@ -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<ReturnType<Models['workspaceUser']['get']>>
>;
workspacePolicy: Map<string, Awaited<ReturnType<Models['workspace']['get']>>>;
workspaceRuntime: Map<
string,
Awaited<ReturnType<Models['workspaceRuntimeState']['get']>>
>;
workspaceQuotaRuntime: Map<string, NewWorkspaceRuntimeState>;
docPolicies: Map<
string,
Awaited<ReturnType<Models['doc']['findDefaultRoles']>>
>;
docGrants: Map<string, Awaited<ReturnType<Models['docUser']['findMany']>>>;
};
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<PermissionEvaluationInputV1> {
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<PermissionEvaluationInputV1> {
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<T>(
map: Map<string, Promise<T> | T>,
key: string,
load: () => Promise<T>
) {
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<NewWorkspaceRuntimeState[]>`
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<NewWorkspaceMemberRow[]>`
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<NewWorkspacePolicyRow[]>`
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<NewDocPolicyRow[]>`
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<NewDocGrantRow[]>`
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)]})
`;
}
}
@@ -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<PermissionEvaluationInputV1['docs']>[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, PermissionDocRole>([
[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<PermissionDocRole, DocRole>([
['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,
};
}
@@ -1,53 +0,0 @@
import { Logger, OnModuleInit } from '@nestjs/common';
import type {
Resource,
ResourceAction,
ResourceRole,
ResourceType,
} from './resource';
const ACTION_CHECKER_PROVIDERS = new Map<ResourceType, AccessController<any>>();
function registerAccessController<Type extends ResourceType>(
type: Type,
provider: AccessController<Type>
) {
ACTION_CHECKER_PROVIDERS.set(type, provider);
}
export function getAccessController<Type extends ResourceType>(
type: Type
): AccessController<Type> {
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<Type>,
action: ResourceAction<Type>
): Promise<void>;
abstract can(
resource: Resource<Type>,
action: ResourceAction<Type>
): Promise<boolean>;
abstract role(resource: Resource<Type>): Promise<{
role: ResourceRole<Type> | null;
permissions: Record<ResourceAction<Type>, boolean>;
}>;
}
@@ -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<typeof docLegacyBoundary> & { decisions: unknown }
>,
current: Array<
ReturnType<typeof docLegacyBoundary> & { 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 });
}
}
@@ -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<DocRole | null> {
const workspaceController = getAccessController(
'ws'
) as WorkspaceAccessController;
const docRoles = await workspaceController.getDocRoles(payload, [
payload.docId,
]);
return docRoles[0];
}
}
@@ -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,
@@ -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<QuotaService['getWorkspaceQuotaWithUsage']>
ReturnType<QuotaStateService['reconcileWorkspaceQuotaState']>
> & {
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<WorkspaceRole | null>;
docRoles(
resource: Resource<'ws'>,
docIds: string[]
): Promise<Array<{ role: unknown; permissions: Record<DocAction, boolean> }>>;
};
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<WorkspaceState> {
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')
@@ -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<ProjectionDecisionSample[]>`
(
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;
}
}
@@ -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<PermissionWorkspaceAction>(
[
'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<PermissionDocAction>([
'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<T extends { docId: string }>(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<PermissionContextLoader['load']>[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<PermissionContextLoader['load']>[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;
@@ -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<string, string[]> };
workspace?: { roles?: Record<string, string[]> };
};
private readonly legacyDocRoleValues = new Map<string, DocRole>([
['external', DocRole.External],
['reader', DocRole.Reader],
['commenter', DocRole.Commenter],
['editor', DocRole.Editor],
['manager', DocRole.Manager],
['owner', DocRole.Owner],
]);
private readonly legacyWorkspaceRoleValues = new Map<string, number>([
['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[]`}))
)
)
`;
}
}
@@ -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,
];
@@ -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;
}
}
@@ -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<Context>;
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<Context>) {
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<Context>,
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<Context>,
workspace: Workspace,
key: string,
size: number
) {
await t.context.models.blob.upsert({
workspaceId: workspace.id,
key,
mime: 'application/octet-stream',
size,
});
}
@@ -19,4 +19,5 @@ export class QuotaModule {}
export { QuotaService };
export { QuotaServiceModule };
export { QuotaStateService } from './state';
export { WorkspaceQuotaHumanReadableType, WorkspaceQuotaType } from './types';
@@ -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<string, never>;
output: { state: UserQuotaStateSnapshot };
};
'workspace.quota-state.get': {
input: { workspaceId: string };
output: { state: WorkspaceQuotaStateSnapshot };
};
}
interface RealtimeTopicMap {
'user.quota-state.changed': {
input: Record<string, never>;
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<T extends Record<string, unknown>>(state: T) {
return Object.fromEntries(
Object.entries(state).map(([key, value]) => [
key,
typeof value === 'bigint' ? Number(value) : value,
])
);
}
}
@@ -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 {}
+88 -126
View File
@@ -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<UserQuota> {
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<UserQuotaWithUsage> {
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<WorkspaceQuota> {
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<WorkspaceQuotaWithUsage> {
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<ReturnType<QuotaStateService['reconcileUserQuotaState']>>
): 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<QuotaStateService['reconcileWorkspaceQuotaState']>
>
): 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';
}
}
}
@@ -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<EntitlementService['resolveUserEntitlement']>
>['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<unknown>>) {
await this.mapMany(tasks, task => task());
}
private async mapMany<T, U>(items: T[], mapper: (item: T) => Promise<U>) {
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<PrismaClient['effectiveUserQuotaState']['findUnique']>
>,
current: Awaited<
ReturnType<PrismaClient['effectiveUserQuotaState']['upsert']>
>
) {
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<PrismaClient['effectiveWorkspaceQuotaState']['findUnique']>
>,
current: Awaited<
ReturnType<PrismaClient['effectiveWorkspaceQuotaState']['upsert']>
>
) {
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);
}
}
@@ -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`;
}
@@ -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,
@@ -18,9 +18,11 @@ export {
realtimeCommentRoom,
realtimeNotificationRoom,
realtimeTranscriptTaskRoom,
realtimeUserQuotaStateRoom,
realtimeUserRoom,
realtimeWorkspaceDocRoom,
realtimeWorkspaceEmbeddingProgressRoom,
realtimeWorkspaceQuotaStateRoom,
realtimeWorkspaceRoom,
} from './rooms';
export type { RealtimeRequestHandler, RealtimeTopicHandler } from './types';
@@ -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');
}
@@ -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
) {
@@ -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')
@@ -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,
@@ -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);
@@ -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<PaginatedDocType> {
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
) {}
@@ -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<SnapshotHistory> {
export class DocHistoryResolver {
constructor(
private readonly workspace: PgWorkspaceDocStorageAdapter,
private readonly ac: AccessController
private readonly ac: PermissionAccess
) {}
@ResolveField(() => [DocHistoryType])
@@ -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 });
}
}
}
@@ -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,
@@ -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<InviteInfo> {
@@ -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) {
@@ -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) {}
}
@@ -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) {}
}
@@ -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';
@@ -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: {
@@ -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<T>(operation: Promise<T>) {
try {
return await operation;
} catch (err) {
this.models.permissionProjection.recordTriggerErrorMetric(err);
throw err;
}
}
}
+55 -87
View File
@@ -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<TransactionalAdapterPrisma>({ 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,
};
}
}
+41 -19
View File
@@ -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
@@ -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';
@@ -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<Prisma.TransactionClient, '$executeRaw'>
) => Promise<void>
) => Promise<void>;
};
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<string, number>;
};
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<CountRow[]>) {
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<PermissionProjectionCheckReport> {
const [
oldWorkspacePolicyMismatch,
oldAcceptedMemberMismatch,
extraProjectedMember,
oldInvitationMismatch,
extraProjectedInvitation,
oldDocGrantMismatch,
extraProjectedDocGrant,
oldDocPolicyMismatch,
extraProjectedDocPolicy,
ownerConflict,
invalidLegacyRows,
] = await Promise.all([
count(this.db.$queryRaw<CountRow[]>`
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<CountRow[]>`
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<CountRow[]>`
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<CountRow[]>`
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<CountRow[]>`
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<CountRow[]>`
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<CountRow[]>`
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<CountRow[]>`
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<CountRow[]>`
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<CountRow[]>`
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<ProjectionIssueRow[]>`
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)
`;
}
}
@@ -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<boolean>;
@Transactional()
async set(
workspaceId: string,
userId: string,
role: WorkspaceRole,
status: WorkspaceMemberStatus,
data: {
source?: WorkspaceMemberSource;
inviterId?: string;
} = {}
): Promise<void> {
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<Array<{ exists: boolean }>>`
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,
},
});
}
}
@@ -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<boolean>;
async get(workspaceId: string): Promise<WorkspaceRuntimeState> {
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<WorkspaceRuntimeStateRow[]>`
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<Array<{ exists: boolean }>>`
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<LegacyWorkspaceRuntimeStateRow[]>`
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()
`;
}
}
@@ -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<User, keyof typeof workspaceUserSelect>;
};
export type WorkspaceMemberWithUser = WorkspaceMember & {
user?: Pick<User, keyof typeof workspaceUserSelect>;
};
export type WorkspaceInvitationWithUser = WorkspaceInvitation & {
inviteeUser?: Pick<User, keyof typeof workspaceUserSelect> | 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<WorkspaceUserCompatRow[]>`
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<WorkspaceUserCompatRow[]>`
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<void> };
workspaceMember: {
setActive(
workspaceId: string,
userId: string,
role: WorkspaceRole
): Promise<unknown>;
};
},
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<WorkspaceMemberSource, WorkspaceInvitation[]>;
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;
}
}
@@ -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<WorkspaceMemberSource, WorkspaceUserRole[]>;
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
);
}
}
+49 -15
View File
@@ -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
@@ -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)
@@ -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;
}
}
}
@@ -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
) {}
@@ -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,
@@ -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,

Some files were not shown because too many files have changed in this diff Show More