@dxos/echo-pipeline 0.7.3-staging.cc8dd3e → 0.7.3

package/src/edge/echo-edge-replicator.test.ts

@@ -2,7 +2,8 @@
 // Copyright 2024 DXOS.org
 //
 
-import { describe, test } from 'vitest';
+import { getRandomPort } from 'get-port-please';
+import { describe, expect, onTestFinished, test } from 'vitest';
 
 import { Event } from '@dxos/async';
 import { cbor } from '@dxos/automerge/automerge-repo';
@@ -16,69 +17,111 @@ import type { Peer } from '@dxos/protocols/proto/dxos/edge/messenger';
 import { openAndClose } from '@dxos/test-utils';
 
 import { EchoEdgeReplicator } from './echo-edge-replicator';
-import type { EchoReplicatorContext } from '../automerge';
+import type { EchoReplicatorContext, ReplicatorConnection } from '../automerge';
 
 describe('EchoEdgeReplicator', () => {
-  test('reconnects', async ({ onTestFinished }) => {
-    const { endpoint, cleanup, sendMessage } = await createTestEdgeWsServer(8001);
-    onTestFinished(cleanup);
-    const client = new EdgeClient(await createEphemeralEdgeIdentity(), { socketEndpoint: endpoint });
-    await openAndClose(client);
+  test('reconnects', async () => {
+    const { client, server } = await createClientServer();
 
     const spaceId = SpaceId.random();
 
-    const replicator = new EchoEdgeReplicator({ edgeConnection: client });
     const { context, connectionOpen } = createMockContext();
-    await replicator.connect(context);
+    const replicator = await connectReplicator(client, context);
     await replicator.connectToSpace(spaceId);
 
     client.setIdentity(await createEphemeralEdgeIdentity());
     await connectionOpen.waitForCount(1);
 
-    sendMessage(
-      createForbiddenMessage(
-        {
-          identityKey: client.identityKey,
-          peerKey: client.peerKey,
-        },
-        spaceId,
-      ),
-    );
+    const forbidden = createForbiddenMessage({ identityKey: client.identityKey, peerKey: client.peerKey }, spaceId);
+    server.sendMessage(forbidden);
     await connectionOpen.waitForCount(1);
 
     // Double restart to check for race conditions.
     client.setIdentity(await createEphemeralEdgeIdentity());
-    sendMessage(
-      createForbiddenMessage(
-        {
-          identityKey: client.identityKey,
-          peerKey: client.peerKey,
-        },
-        spaceId,
-      ),
-    );
+    server.sendMessage(forbidden);
     await connectionOpen.waitForCount(1);
 
     await replicator.disconnect();
   });
+
+  describe('shouldAdvertise', () => {
+    test('true if space document belongs to connection space', async () => {
+      const { client } = await createClientServer();
+
+      const spaceId = SpaceId.random();
+      const documentId = PublicKey.random().toHex();
+      const { context, openConnections } = createMockContext({
+        documentSpaceId: { [documentId]: spaceId },
+      });
+      const replicator = await connectReplicator(client, context);
+      await replicator.connectToSpace(spaceId);
+
+      await expect.poll(() => openConnections.length === 1).toBeTruthy();
+      expect(openConnections[0].shouldAdvertise({ documentId })).toBeTruthy();
+    });
+
+    test('checks remote collection if space id can not be resolved', async () => {
+      const { client } = await createClientServer();
+
+      const spaceId = SpaceId.random();
+      const documentId = PublicKey.random().toHex();
+      const remoteCollections: { [peerId: string]: { [documentId: string]: boolean } } = {};
+      const { context, openConnections } = createMockContext({ remoteCollections });
+      const replicator = await connectReplicator(client, context);
+      await replicator.connectToSpace(spaceId);
+
+      await expect.poll(() => openConnections.length === 1).toBeTruthy();
+      const connection = openConnections[0];
+      expect(await connection.shouldAdvertise({ documentId })).toBeFalsy();
+      remoteCollections[connection.peerId] = { [documentId]: true };
+      expect(await connection.shouldAdvertise({ documentId })).toBeTruthy();
+    });
+  });
+
+  const connectReplicator = async (client: EdgeClient, context: EchoReplicatorContext) => {
+    const replicator = new EchoEdgeReplicator({ edgeConnection: client });
+    await replicator.connect(context);
+    onTestFinished(() => replicator.disconnect());
+    return replicator;
+  };
+
+  const createClientServer = async () => {
+    const server = await createTestEdgeWsServer(await getRandomPort());
+    onTestFinished(server.cleanup);
+    const client = new EdgeClient(await createEphemeralEdgeIdentity(), { socketEndpoint: server.endpoint });
+    await openAndClose(client);
+    return { client, server };
+  };
 });
 
-const createMockContext = () => {
+const createMockContext = (args?: {
+  remoteCollections?: { [peerId: string]: { [documentId: string]: boolean } };
+  documentSpaceId?: { [documentId: string]: SpaceId };
+}) => {
   const connectionOpen = new Event();
+  const openConnections: ReplicatorConnection[] = [];
   return {
     context: {
-      getContainingSpaceIdForDocument: async (documentId) => null,
+      getContainingSpaceIdForDocument: async (documentId) => args?.documentSpaceId?.[documentId] ?? null,
       getContainingSpaceForDocument: async (documentId) => null,
       onConnectionAuthScopeChanged: (connection) => {},
-      isDocumentInRemoteCollection: async (params) => false,
-      onConnectionClosed: (connection) => {},
+      isDocumentInRemoteCollection: async (params) =>
+        args?.remoteCollections?.[params.peerId]?.[params.documentId] ?? false,
+      onConnectionClosed: (connection) => {
+        const idx = openConnections.indexOf(connection);
+        if (idx >= 0) {
+          openConnections.splice(idx, 1);
+        }
+      },
       onConnectionOpen: (connection) => {
+        openConnections.push(connection);
         connectionOpen.emit();
       },
 
       peerId: PublicKey.random().toHex(),
     } satisfies EchoReplicatorContext,
 
+    openConnections,
     connectionOpen,
   };
 };

package/src/edge/echo-edge-replicator.ts

@@ -267,9 +267,20 @@ class EdgeReplicatorConnection extends Resource implements ReplicatorConnection
     }
     const spaceId = await this._context.getContainingSpaceIdForDocument(params.documentId);
     if (!spaceId) {
-      // There's no spaceId if the document is not present locally. This means the sharePolicy check is being
-      // performed on message reception, so spaceId check was already performed in _onMessage.
-      return true;
+      const remoteDocumentExists = await this._context.isDocumentInRemoteCollection({
+        documentId: params.documentId,
+        peerId: this._remotePeerId as PeerId,
+      });
+
+      log.info('document not found locally for share policy check, accepting the remote document', {
+        documentId: params.documentId,
+        remoteDocumentExists,
+      });
+
+      // If a document is not present locally return true only if it already exists on edge.
+      // Simply returning true will add edge to "generous peers list" for this document which will
+      // start replication of the document after we receive it potentially pushing it to replicator of the wrong space.
+      return remoteDocumentExists;
     }
     return spaceId === this._spaceId;
   }