feat(server): synthetic root doc (#14794)

packages/backend/server/src/__tests__/copilot/copilot.e2e.ts

@@ -11,6 +11,7 @@ import { JobQueue } from '../../base';
 import { ConfigModule } from '../../base/config';
 import { AuthService } from '../../core/auth';
 import { DocReader } from '../../core/doc';
+import { ContextCategories, DocRole, WorkspaceRole } from '../../models';
 import { CopilotContextService } from '../../plugins/copilot/context';
 import {
   CopilotEmbeddingJob,
@@ -25,7 +26,7 @@ import {
   OpenAIProvider,
 } from '../../plugins/copilot/providers';
 import { CopilotStorage } from '../../plugins/copilot/storage';
-import { MockCopilotProvider } from '../mocks';
+import { MockCopilotProvider, Mockers } from '../mocks';
 import {
   acceptInviteById,
   createTestingApp,
@@ -36,6 +37,7 @@ import {
   TestUser,
 } from '../utils';
 import {
+  addContextCategory,
   addContextDoc,
   addContextFile,
   array2sse,
@@ -60,6 +62,7 @@ import {
   getPinnedSessions,
   getWorkspaceSessions,
   listContext,
+  listContextCategories,
   listContextDocAndFiles,
   matchFiles,
   matchWorkspaceDocs,
@@ -1040,6 +1043,114 @@ test('should be able to manage context', async t => {
   }
 });
 
+test('should reject context reads from another user', async t => {
+  const { app, context, jobs, u1 } = t.context;
+
+  const u2 = await app.signupV1();
+  await app.switchUser(u1);
+
+  const { id: workspaceId } = await createWorkspace(app);
+  const sessionId = await createCopilotSession(
+    app,
+    workspaceId,
+    randomUUID(),
+    textPromptName
+  );
+
+  Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient());
+  Sinon.stub(jobs, 'embeddingClient').get(() => new MockEmbeddingClient());
+
+  const contextId = await createCopilotContext(app, workspaceId, sessionId);
+  await addContextFile(app, contextId, 'sample.txt', Buffer.from('test file'));
+
+  await app.switchUser(u2);
+
+  await t.throwsAsync(
+    app.gql(`
+      query {
+        currentUser {
+          copilot {
+            contexts(contextId: "${contextId}") {
+              id
+            }
+          }
+        }
+      }
+    `)
+  );
+  await t.throwsAsync(matchFiles(app, contextId, 'test', 1));
+});
+
+test('should skip unauthorized docs when adding context category', async t => {
+  const { app, context, jobs, u1 } = t.context;
+
+  const member = await app.signupV1();
+  await app.switchUser(u1);
+
+  const { id: workspaceId } = await createWorkspace(app);
+  await app.create(Mockers.WorkspaceUser, {
+    workspaceId,
+    userId: member.id,
+    type: WorkspaceRole.Collaborator,
+  });
+
+  const readableSnapshot = await app.create(Mockers.DocSnapshot, {
+    workspaceId,
+    user: u1,
+  });
+  const hiddenSnapshot = await app.create(Mockers.DocSnapshot, {
+    workspaceId,
+    user: u1,
+  });
+
+  await app.create(Mockers.DocMeta, {
+    workspaceId,
+    docId: readableSnapshot.id,
+    title: 'readable-doc',
+  });
+  await app.create(Mockers.DocMeta, {
+    workspaceId,
+    docId: hiddenSnapshot.id,
+    title: 'hidden-doc',
+    defaultRole: DocRole.None,
+  });
+
+  Sinon.stub(context, 'embeddingClient').get(() => new MockEmbeddingClient());
+  Sinon.stub(jobs, 'embeddingClient').get(() => new MockEmbeddingClient());
+
+  await app.switchUser(member);
+  const sessionId = await createCopilotSession(
+    app,
+    workspaceId,
+    randomUUID(),
+    textPromptName
+  );
+  const contextId = await createCopilotContext(app, workspaceId, sessionId);
+  const category = await addContextCategory(
+    app,
+    contextId,
+    ContextCategories.Collection,
+    'fav',
+    [readableSnapshot.id, hiddenSnapshot.id]
+  );
+
+  t.deepEqual(
+    category.docs.map(doc => doc.id),
+    [readableSnapshot.id]
+  );
+
+  const ret = await listContextCategories(
+    app,
+    workspaceId,
+    sessionId,
+    contextId
+  );
+  t.deepEqual(
+    ret?.collections?.[0]?.docs.map(doc => doc.id),
+    [readableSnapshot.id]
+  );
+});
+
 test('should be able to transcript', async t => {
   const { app } = t.context;
 

packages/backend/server/src/__tests__/e2e/doc/resolver.spec.ts

@@ -3,9 +3,11 @@ import { randomUUID } from 'node:crypto';
 import {
   getRecentlyUpdatedDocsQuery,
   getWorkspacePageByIdQuery,
+  type GraphQLQuery,
   publishPageMutation,
 } from '@affine/graphql';
 
+import { DocRole, WorkspaceRole } from '../../../models';
 import { Mockers } from '../../mocks';
 import { app, e2e } from '../test';
 
@@ -152,3 +154,97 @@ e2e('should get doc with title and null summary', async t => {
   t.is(result.workspace.doc.title, doc.title);
   t.is(result.workspace.doc.summary, null);
 });
+
+e2e('should require owner or admin to query workspace docs', 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 docSnapshot = await app.create(Mockers.DocSnapshot, {
+    workspaceId: workspace.id,
+    user: owner,
+  });
+  await app.create(Mockers.DocMeta, {
+    workspaceId: workspace.id,
+    docId: docSnapshot.id,
+    title: 'private-doc',
+    defaultRole: DocRole.None,
+  });
+
+  await app.switchUser(member);
+  await t.throwsAsync(
+    app.gql({
+      query: {
+        id: 'workspaceDocsPermissionCheck',
+        op: 'workspaceDocsPermissionCheck',
+        query: `
+          query {
+            workspace(id: "${workspace.id}") {
+              docs(pagination: { first: 10 }) {
+                totalCount
+              }
+            }
+          }
+        `,
+      } satisfies GraphQLQuery,
+      variables: undefined,
+    })
+  );
+});
+
+e2e('should require Doc.Read to query workspace page meta', 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 docSnapshot = await app.create(Mockers.DocSnapshot, {
+    workspaceId: workspace.id,
+    user: owner,
+  });
+  const doc = await app.create(Mockers.DocMeta, {
+    workspaceId: workspace.id,
+    docId: docSnapshot.id,
+    title: 'private-doc',
+    defaultRole: DocRole.None,
+  });
+
+  await app.switchUser(member);
+  await t.throwsAsync(
+    app.gql({
+      query: {
+        id: 'workspacePageMetaPermissionCheck',
+        op: 'workspacePageMetaPermissionCheck',
+        query: `
+          query {
+            workspace(id: "${workspace.id}") {
+              pageMeta(pageId: "${doc.docId}") {
+                createdAt
+              }
+            }
+          }
+        `,
+      } satisfies GraphQLQuery,
+      variables: undefined,
+    })
+  );
+});

packages/backend/server/src/__tests__/sync/gateway.spec.ts

@@ -4,6 +4,12 @@ import { io, type Socket as SocketIOClient } from 'socket.io-client';
 import { Doc, encodeStateAsUpdate } from 'yjs';
 
 import { CANARY_CLIENT_VERSION_MAX_AGE_DAYS } from '../../base';
+import {
+  DocRole,
+  Models,
+  WorkspaceMemberStatus,
+  WorkspaceRole,
+} from '../../models';
 import { createTestingApp, TestingApp } from '../utils';
 
 type WebsocketResponse<T> =
@@ -25,6 +31,16 @@ function unwrapResponse<T>(t: ExecutionContext, res: WebsocketResponse<T>): T {
   throw new Error(`Websocket error: ${res.error.name}: ${res.error.message}`);
 }
 
+function getErrorResponse<T>(
+  t: ExecutionContext,
+  res: WebsocketResponse<T>
+): { name: string; message: string } {
+  if ('error' in res) return res.error;
+
+  t.log(res);
+  throw new Error(`Expected websocket error response, got data instead`);
+}
+
 async function withTimeout<T>(
   promise: Promise<T>,
   timeoutMs: number,
@@ -129,7 +145,7 @@ function expectNoEvent(
 }
 
 async function login(app: TestingApp) {
-  const user = await app.createUser('u1@affine.pro');
+  const user = await app.createUser();
   const res = await app
     .POST('/api/auth/sign-in')
     .send({ email: user.email, password: user.password })
@@ -519,3 +535,66 @@ test('active users metric should dedupe multiple sockets for one user', async t
     await Promise.all([waitForDisconnect(first), waitForDisconnect(second)]);
   }
 });
+
+test('workspace sync delete-doc should enforce doc permissions', async t => {
+  const db = app.get(PrismaClient);
+  const models = app.get(Models);
+  const { user: owner } = await login(app);
+  const { user: collaborator, cookieHeader } = await login(app);
+  const workspace = await models.workspace.create(owner.id);
+  const docId = 'private-doc';
+
+  await models.workspaceUser.set(
+    workspace.id,
+    collaborator.id,
+    WorkspaceRole.Collaborator,
+    {
+      status: WorkspaceMemberStatus.Accepted,
+    }
+  );
+  await models.doc.setDefaultRole(workspace.id, docId, DocRole.None);
+  await db.snapshot.create({
+    data: {
+      id: docId,
+      workspaceId: workspace.id,
+      blob: Buffer.from([1, 1]),
+      state: Buffer.from([1, 1]),
+      createdAt: new Date(),
+      updatedAt: new Date(),
+      createdBy: owner.id,
+      updatedBy: owner.id,
+    },
+  });
+
+  const socket = createClient(url, cookieHeader);
+
+  try {
+    await waitForConnect(socket);
+
+    const join = unwrapResponse(
+      t,
+      await emitWithAck<{ clientId: string; success: boolean }>(
+        socket,
+        'space:join',
+        {
+          spaceType: 'workspace',
+          spaceId: workspace.id,
+          clientVersion: '0.26.0',
+        }
+      )
+    );
+    t.true(join.success);
+
+    const error = getErrorResponse(
+      t,
+      await emitWithAck(socket, 'space:delete-doc', {
+        spaceType: 'workspace',
+        spaceId: workspace.id,
+        docId,
+      })
+    );
+    t.true(error.message.includes('Doc.Delete'));
+  } finally {
+    socket.disconnect();
+  }
+});

packages/backend/server/src/__tests__/utils/copilot.ts

@@ -1,3 +1,4 @@
+import { ContextCategories } from '../../models';
 import { PromptConfig, PromptMessage } from '../../plugins/copilot/providers';
 import { NodeExecutorType } from '../../plugins/copilot/workflow/executor';
 import {
@@ -318,6 +319,33 @@ export async function addContextDoc(
   return res.addContextDoc;
 }
 
+export async function addContextCategory(
+  app: TestingApp,
+  contextId: string,
+  type: ContextCategories,
+  categoryId: string,
+  docs: string[]
+): Promise<{ type: string; id: string; docs: { id: string }[] }> {
+  const graphqlType =
+    type === ContextCategories.Collection ? 'Collection' : 'Tag';
+  const res = await app.gql(
+    `
+      mutation addContextCategory($options: AddContextCategoryInput!) {
+        addContextCategory(options: $options) {
+          type
+          id
+          docs {
+            id
+          }
+        }
+      }
+    `,
+    { options: { contextId, type: graphqlType, categoryId, docs } }
+  );
+
+  return res.addContextCategory;
+}
+
 export async function removeContextDoc(
   app: TestingApp,
   contextId: string,
@@ -390,6 +418,50 @@ export async function listContextDocAndFiles(
   return { docs, files };
 }
 
+export async function listContextCategories(
+  app: TestingApp,
+  workspaceId: string,
+  sessionId: string,
+  contextId: string
+): Promise<
+  | {
+      collections: {
+        type: string;
+        id: string;
+        docs: {
+          id: string;
+          status: string;
+          createdAt: number;
+        }[];
+      }[];
+    }
+  | undefined
+> {
+  const res = await app.gql(`
+        query {
+          currentUser {
+            copilot(workspaceId: "${workspaceId}") {
+              contexts(sessionId: "${sessionId}", contextId: "${contextId}") {
+                collections {
+                  type
+                  id
+                  docs {
+                    id
+                    status
+                    createdAt
+                  }
+                }
+              }
+            }
+          }
+        }
+      `);
+
+  const { collections } = res.currentUser?.copilot?.contexts?.[0] || {};
+
+  return { collections };
+}
+
 export async function submitAudioTranscription(
   app: TestingApp,
   workspaceId: string,

packages/backend/server/src/__tests__/workspace.e2e.ts

@@ -2,6 +2,7 @@ import { PrismaClient } from '@prisma/client';
 import type { TestFn } from 'ava';
 import ava from 'ava';
 
+import { addDocToRootDoc, readAllDocIdsFromRootDoc } from '../native';
 import {
   acceptInviteById,
   createTestingApp,
@@ -58,10 +59,15 @@ test('should be able to publish workspace', async t => {
 });
 
 test('should visit public page', async t => {
-  const { app } = t.context;
-  await app.signupV1('u1@affine.pro');
+  const { app, client } = t.context;
+  const user = await app.signupV1('u1@affine.pro');
 
   const workspace = await createWorkspace(app);
+  const rootDoc = addDocToRootDoc(Buffer.from([0, 0]), 'doc1', 'doc1');
+  await client.snapshot.update({
+    where: { workspaceId_id: { workspaceId: workspace.id, id: workspace.id } },
+    data: { blob: rootDoc },
+  });
   const share = await publishDoc(app, workspace.id, 'doc1');
 
   t.is(share.id, 'doc1', 'failed to share doc');
@@ -78,18 +84,32 @@ test('should visit public page', async t => {
     `/api/workspaces/${workspace.id}/docs/${workspace.id}`
   );
   t.is(resp1.statusCode, 200, 'failed to get root doc with u1 token');
+
+  await app.logout();
   const resp2 = await app.GET(
     `/api/workspaces/${workspace.id}/docs/${workspace.id}`
   );
-  t.is(resp2.statusCode, 200, 'failed to get root doc with public pages');
+  t.is(resp2.statusCode, 403, 'legacy root doc should not be public');
+
+  const respPublicRoot = await app.GET(
+    `/api/workspaces/${workspace.id}/public-docs/doc1/root-doc`
+  );
+  t.is(
+    respPublicRoot.statusCode,
+    200,
+    'failed to get filtered public root doc'
+  );
+  t.deepEqual(readAllDocIdsFromRootDoc(respPublicRoot.body, false), ['doc1']);
 
   const resp3 = await app.GET(`/api/workspaces/${workspace.id}/docs/doc1`);
   // 404 because we don't put the page doc to server
-  t.is(resp3.statusCode, 404, 'failed to get shared doc with u1 token');
+  t.is(resp3.statusCode, 404, 'failed to get shared doc without snapshot');
   const resp4 = await app.GET(`/api/workspaces/${workspace.id}/docs/doc1`);
   // 404 because we don't put the page doc to server
   t.is(resp4.statusCode, 404, 'should not get shared doc without token');
 
+  await app.login(user);
+
   const revoke = await revokePublicDoc(app, workspace.id, 'doc1');
   t.false(revoke.public, 'failed to revoke doc');
   const docs2 = await getWorkspacePublicDocs(app, workspace.id);

packages/backend/server/src/__tests__/workspace/controller.spec.ts

@@ -4,10 +4,17 @@ import { HttpStatus } from '@nestjs/common';
 import { PrismaClient, WorkspaceMemberStatus } from '@prisma/client';
 import ava, { TestFn } from 'ava';
 import Sinon from 'sinon';
+import supertest from 'supertest';
+import { applyUpdate, Doc as YDoc, Map as YMap } from 'yjs';
 
 import { PgWorkspaceDocStorageAdapter } from '../../core/doc';
 import { WorkspaceBlobStorage } from '../../core/storage';
 import { Models, PublicDocMode, WorkspaceRole } from '../../models';
+import {
+  addDocToRootDoc,
+  mergeUpdatesInApplyWay,
+  readAllDocIdsFromRootDoc,
+} from '../../native';
 import { createTestingApp, TestingApp, TestUser } from '../utils';
 
 const test = ava as TestFn<{
@@ -217,6 +224,81 @@ test('should be able to get doc', async t => {
   t.deepEqual(res.body, Buffer.from([0, 0]));
 });
 
+test('should not expose legacy root doc for private workspace with public pages', async t => {
+  const { app } = t.context;
+
+  const res = await app.GET('/api/workspaces/private/docs/private');
+
+  t.is(res.status, HttpStatus.FORBIDDEN);
+});
+
+test('should expose filtered public root doc for shared page', async t => {
+  const { app, workspace: doc } = t.context;
+
+  let root = addDocToRootDoc(Buffer.from([0, 0]), 'public', 'Public Doc');
+  const privateUpdate = addDocToRootDoc(root, 'private-page', 'Private Doc');
+  root = mergeUpdatesInApplyWay([root, privateUpdate]);
+
+  doc.getDoc.resolves({
+    spaceId: 'private',
+    docId: 'private',
+    bin: root,
+    timestamp: Date.now(),
+  });
+
+  const res = await app.GET(
+    '/api/workspaces/private/public-docs/public/root-doc'
+  );
+
+  t.is(res.status, HttpStatus.OK);
+  const body = Buffer.isBuffer(res.body) ? res.body : Buffer.from(res.body);
+  t.deepEqual(readAllDocIdsFromRootDoc(body, false), ['public']);
+
+  const ydoc = new YDoc({ guid: 'private' });
+  t.notThrows(() =>
+    applyUpdate(
+      ydoc,
+      new Uint8Array(body.buffer, body.byteOffset, body.byteLength)
+    )
+  );
+  const pages = (ydoc.getMap('meta') as YMap<unknown>).get('pages') as
+    | { toArray: () => Array<{ get: (key: string) => unknown }> }
+    | undefined;
+  t.deepEqual(
+    pages?.toArray().map(page => page.get('id')),
+    ['public']
+  );
+});
+
+test('should expose public doc publish mode through HEAD route', async t => {
+  const { app } = t.context;
+
+  const res = await supertest(app.getHttpServer()).head(
+    '/api/workspaces/private/public-docs/public'
+  );
+
+  t.is(res.status, HttpStatus.OK);
+  t.is(res.get('publish-mode'), 'page');
+});
+
+test('should expose public doc binary through public route', async t => {
+  const { app, workspace: doc } = t.context;
+
+  doc.getDoc.resolves({
+    spaceId: 'private',
+    docId: 'public',
+    bin: Buffer.from([1, 2, 3]),
+    timestamp: Date.now(),
+  });
+
+  const res = await app.GET('/api/workspaces/private/public-docs/public');
+
+  t.is(res.status, HttpStatus.OK);
+  t.is(res.get('content-type'), 'application/octet-stream');
+  t.is(res.get('publish-mode'), 'page');
+  t.deepEqual(res.body, Buffer.from([1, 2, 3]));
+});
+
 test('should record doc view when reading doc', async t => {
   const { app, workspace: doc, models } = t.context;