@dxos/echo-pipeline 0.3.11-main.ec7a2b3 → 0.3.11-main.ee2b64c

package/src/automerge/automerge-host.ts

@@ -2,6 +2,8 @@
 // Copyright 2023 DXOS.org
 //
 
+import { Trigger } from '@dxos/async';
+import { next as automerge } from '@dxos/automerge/automerge';
 import {
   Repo,
   NetworkAdapter,
@@ -15,19 +17,27 @@ import {
 import { IndexedDBStorageAdapter } from '@dxos/automerge/automerge-repo-storage-indexeddb';
 import { Stream } from '@dxos/codec-protobuf';
 import { invariant } from '@dxos/invariant';
+import { PublicKey } from '@dxos/keys';
 import { log } from '@dxos/log';
 import { type HostInfo, type SyncRepoRequest, type SyncRepoResponse } from '@dxos/protocols/proto/dxos/echo/service';
 import { type PeerInfo } from '@dxos/protocols/proto/dxos/mesh/teleport/automerge';
 import { StorageType, type Directory } from '@dxos/random-access-storage';
 import { AutomergeReplicator } from '@dxos/teleport-extension-automerge-replicator';
-import { arrayToBuffer, bufferToArray } from '@dxos/util';
+import { trace } from '@dxos/tracing';
+import { ComplexMap, ComplexSet, arrayToBuffer, bufferToArray, defaultMap, mapValues } from '@dxos/util';
 
+@trace.resource()
 export class AutomergeHost {
   private readonly _repo: Repo;
   private readonly _meshNetwork: MeshNetworkAdapter;
   private readonly _clientNetwork: LocalHostNetworkAdapter;
   private readonly _storage: StorageAdapter;
 
+  /**
+   * spaceKey -> deviceKey[]
+   */
+  private readonly _authorizedDevices = new ComplexMap<PublicKey, ComplexSet<PublicKey>>(PublicKey.hash);
+
   constructor(storageDirectory: Directory) {
     this._meshNetwork = new MeshNetworkAdapter();
     this._clientNetwork = new LocalHostNetworkAdapter();
@@ -38,11 +48,52 @@ export class AutomergeHost {
         ? new IndexedDBStorageAdapter(storageDirectory.path, 'data')
         : new AutomergeStorageAdapter(storageDirectory);
     this._repo = new Repo({
+      peerId: `host-${PublicKey.random().toHex()}` as PeerId,
       network: [this._clientNetwork, this._meshNetwork],
       storage: this._storage,
 
       // TODO(dmaretskyi): Share based on HALO permissions and space affinity.
-      sharePolicy: async (peerId, documentId) => true, // Share everything.
+      // Hosts, running in the worker, don't share documents unless requested by other peers.
+      sharePolicy: async (peerId /* device key */, documentId /* space key */) => {
+        if (peerId.startsWith('client-')) {
+          return true;
+        }
+
+        if (!documentId) {
+          return false;
+        }
+
+        const doc = this._repo.handles[documentId]?.docSync();
+        if (!doc) {
+          log('doc not found for share policy check', { peerId, documentId });
+          return false;
+        }
+
+        try {
+          if (!doc.experimental_spaceKey) {
+            log('space key not found for share policy check', { peerId, documentId });
+            return false;
+          }
+
+          const spaceKey = PublicKey.from(doc.experimental_spaceKey);
+          const authorizedDevices = this._authorizedDevices.get(spaceKey);
+
+          // TODO(mykola): Hack, stop abusing `peerMetadata` field.
+          const deviceKeyHex = (this.repo.peerMetadataByPeerId[peerId] as any)?.dxos_deviceKey;
+          if (!deviceKeyHex) {
+            log('device key not found for share policy check', { peerId, documentId });
+            return false;
+          }
+          const deviceKey = PublicKey.from(deviceKeyHex);
+
+          const isAuthorized = authorizedDevices?.has(deviceKey) ?? false;
+          log('share policy check', { peerId, documentId, deviceKey, spaceKey, isAuthorized });
+          return isAuthorized;
+        } catch (err) {
+          log.catch(err);
+          return false;
+        }
+      }, // Share everything.
     });
     this._clientNetwork.ready();
     this._meshNetwork.ready();
@@ -52,7 +103,22 @@ export class AutomergeHost {
     return this._repo;
   }
 
+  @trace.info({ depth: null })
+  private _automergeDocs() {
+    return mapValues(this._repo.handles, (handle) => ({
+      state: handle.state,
+      hasDoc: !!handle.docSync(),
+      heads: handle.docSync() ? automerge.getHeads(handle.docSync()) : null,
+    }));
+  }
+
+  @trace.info({ depth: null })
+  private _automergePeers() {
+    return this._repo.peers;
+  }
+
   async close() {
+    this._storage instanceof AutomergeStorageAdapter && (await this._storage.close());
     await this._clientNetwork.close();
   }
 
@@ -68,7 +134,7 @@ export class AutomergeHost {
     return this._clientNetwork.sendSyncMessage(request);
   }
 
-  getHostInfo(): HostInfo {
+  async getHostInfo(): Promise<HostInfo> {
     return this._clientNetwork.getHostInfo();
   }
 
@@ -79,6 +145,10 @@ export class AutomergeHost {
   createExtension(): AutomergeReplicator {
     return this._meshNetwork.createExtension();
   }
+
+  authorizeDevice(spaceKey: PublicKey, deviceKey: PublicKey) {
+    defaultMap(this._authorizedDevices, spaceKey, () => new ComplexSet(PublicKey.hash)).add(deviceKey);
+  }
 }
 
 type ClientSyncState = {
@@ -104,8 +174,11 @@ class LocalHostNetworkAdapter extends NetworkAdapter {
     });
   }
 
+  private _connected = new Trigger();
+
   override connect(peerId: PeerId): void {
     this.peerId = peerId;
+    this._connected.wake();
     // No-op. Client always connects first
   }
 
@@ -146,18 +219,26 @@ class LocalHostNetworkAdapter extends NetworkAdapter {
         },
       });
 
-      this.emit('peer-candidate', {
-        peerId,
-      });
+      this._connected
+        .wait({ timeout: 1_000 })
+        .then(() => {
+          this.emit('peer-candidate', {
+            peerMetadata: {},
+            peerId,
+          });
+        })
+        .catch((err) => log.catch(err));
     });
   }
 
   async sendSyncMessage({ id, syncMessage }: SyncRepoRequest): Promise<void> {
+    await this._connected.wait({ timeout: 1_000 });
     const message = cbor.decode(syncMessage!) as Message;
     this.emit('message', message);
   }
 
-  getHostInfo(): HostInfo {
+  async getHostInfo(): Promise<HostInfo> {
+    await this._connected.wait({ timeout: 1_000 });
     invariant(this.peerId, 'Peer id not set.');
     return {
       peerId: this.peerId,
@@ -172,8 +253,9 @@ class LocalHostNetworkAdapter extends NetworkAdapter {
 /**
  * Used to replicate with other peers over the network.
  */
-class MeshNetworkAdapter extends NetworkAdapter {
+export class MeshNetworkAdapter extends NetworkAdapter {
   private readonly _extensions: Map<string, AutomergeReplicator> = new Map();
+  private _connected = new Trigger();
 
   /**
    * Emits `ready` event. That signals to `Repo` that it can start using the adapter.
@@ -188,6 +270,7 @@ class MeshNetworkAdapter extends NetworkAdapter {
 
   override connect(peerId: PeerId): void {
     this.peerId = peerId;
+    this._connected.wake();
   }
 
   override send(message: Message): void {
@@ -210,20 +293,28 @@ class MeshNetworkAdapter extends NetworkAdapter {
         peerId: this.peerId,
       },
       {
-        onStartReplication: async (info) => {
+        onStartReplication: async (info, remotePeerId /** Teleport ID */) => {
+          await this._connected.wait();
+
           // Note: We store only one extension per peer.
           //       There can be a case where two connected peers have more than one teleport connection between them
           //       and each of them uses different teleport connections to send messages.
           //       It works because we receive messages from all teleport connections and Automerge Repo dedup them.
           // TODO(mykola): Use only one teleport connection per peer.
-          if (this._extensions.has(info.id)) {
-            return;
+          if (!this._extensions.has(info.id)) {
+            peerInfo = info;
+            // TODO(mykola): Fix race condition?
+            this._extensions.set(info.id, extension);
+          } else {
+            // TODO(mykola): retry hack.
+            this.emit('peer-disconnected', { peerId: info.id as PeerId });
           }
 
-          peerInfo = info;
-          // TODO(mykola): Fix race condition?
-          this._extensions.set(info.id, extension);
           this.emit('peer-candidate', {
+            // TODO(mykola): Hack, stop abusing `peerMetadata` field.
+            peerMetadata: {
+              dxos_deviceKey: remotePeerId.toHex(),
+            } as any,
             peerId: info.id as PeerId,
           });
         },
@@ -233,10 +324,13 @@ class MeshNetworkAdapter extends NetworkAdapter {
           this.emit('message', message);
         },
         onClose: async () => {
-          peerInfo &&
-            this.emit('peer-disconnected', {
-              peerId: peerInfo.id as PeerId,
-            });
+          if (!peerInfo) {
+            return;
+          }
+          this.emit('peer-disconnected', {
+            peerId: peerInfo.id as PeerId,
+          });
+          this._extensions.delete(peerInfo.id);
         },
       },
     );
@@ -245,11 +339,18 @@ class MeshNetworkAdapter extends NetworkAdapter {
 }
 
 export class AutomergeStorageAdapter extends StorageAdapter {
+  // TODO(mykola): Hack for restricting automerge Repo to access storage if Host is `closed`.
+  //               Automerge Repo do not have any lifetime management.
+  private _state: 'opened' | 'closed' = 'opened';
+
   constructor(private readonly _directory: Directory) {
     super();
   }
 
   override async load(key: StorageKey): Promise<Uint8Array | undefined> {
+    if (this._state !== 'opened') {
+      return undefined;
+    }
     const filename = this._getFilename(key);
     const file = this._directory.getOrCreateFile(filename);
     const { size } = await file.stat();
@@ -261,6 +362,9 @@ export class AutomergeStorageAdapter extends StorageAdapter {
   }
 
   override async save(key: StorageKey, data: Uint8Array): Promise<void> {
+    if (this._state !== 'opened') {
+      return undefined;
+    }
     const filename = this._getFilename(key);
     const file = this._directory.getOrCreateFile(filename);
     await file.write(0, arrayToBuffer(data));
@@ -270,13 +374,19 @@ export class AutomergeStorageAdapter extends StorageAdapter {
   }
 
   override async remove(key: StorageKey): Promise<void> {
+    if (this._state !== 'opened') {
+      return undefined;
+    }
     // TODO(dmaretskyi): Better deletion.
     const filename = this._getFilename(key);
     const file = this._directory.getOrCreateFile(filename);
-    await file.truncate?.(0);
+    await file.destroy();
   }
 
   override async loadRange(keyPrefix: StorageKey): Promise<Chunk[]> {
+    if (this._state !== 'opened') {
+      return [];
+    }
     const filename = this._getFilename(keyPrefix);
     const entries = await this._directory.list();
     return Promise.all(
@@ -295,18 +405,25 @@ export class AutomergeStorageAdapter extends StorageAdapter {
   }
 
   override async removeRange(keyPrefix: StorageKey): Promise<void> {
+    if (this._state !== 'opened') {
+      return undefined;
+    }
     const filename = this._getFilename(keyPrefix);
     const entries = await this._directory.list();
     await Promise.all(
       entries
         .filter((entry) => entry.startsWith(filename))
         .map(async (entry): Promise<void> => {
-          const file = this._directory.getOrCreateFile(filename);
-          await file.truncate?.(0);
+          const file = this._directory.getOrCreateFile(entry);
+          await file.destroy();
         }),
     );
   }
 
+  async close(): Promise<void> {
+    this._state = 'closed';
+  }
+
   private _getFilename(key: StorageKey): string {
     return key.map((k) => k.replaceAll('%', '%25').replaceAll('-', '%2D')).join('-');
   }

package/src/automerge/index.ts

@@ -2,4 +2,4 @@
 // Copyright 2023 DXOS.org
 //
 
-export * from './automerge-host';
+export { AutomergeHost } from './automerge-host';

package/src/pipeline/pipeline-stress.test.ts

@@ -80,7 +80,11 @@ class Agent {
   public messages: FeedMessageBlock[] = [];
   public writePromise: Promise<any> = Promise.resolve();
 
-  constructor(private readonly builder: TestFeedBuilder, public feedStore: FeedStore<FeedMessage>, public id: string) {}
+  constructor(
+    private readonly builder: TestFeedBuilder,
+    public feedStore: FeedStore<FeedMessage>,
+    public id: string,
+  ) {}
 
   async open() {
     const key = await this.builder.keyring.createKey();
@@ -131,7 +135,10 @@ type Real = {
 };
 
 class WriteCommand implements fc.AsyncCommand<Model, Real> {
-  constructor(public agent: string, public count: number) {}
+  constructor(
+    public agent: string,
+    public count: number,
+  ) {}
 
   check = () => true;
 

package/src/space/data-pipeline.ts

@@ -216,6 +216,7 @@ export class DataPipeline implements CredentialProcessor {
   }
 
   private async _consumePipeline() {
+    const pipeline = this._pipeline;
     if (this.currentEpoch) {
       const waitForOneEpoch = this.onNewEpoch.waitForCount(1);
       await this._processEpochInSeparateTask(this.currentEpoch);
@@ -225,8 +226,8 @@ export class DataPipeline implements CredentialProcessor {
     // CPU bottleneck control.
     let messageCounter = 0;
 
-    invariant(this._pipeline, 'Pipeline is not initialized.');
-    for await (const msg of this._pipeline.consume()) {
+    invariant(pipeline, 'Pipeline is not initialized.');
+    for await (const msg of pipeline.consume()) {
       const span = this._usage.beginRecording();
       this._mutations.inc();
 
@@ -261,7 +262,7 @@ export class DataPipeline implements CredentialProcessor {
           } satisfies DataPipelineProcessed);
 
           // Timeframe clock is not updated yet.
-          await this._noteTargetStateIfNeeded(this._pipeline.state.pendingTimeframe);
+          await this._noteTargetStateIfNeeded(pipeline.state.pendingTimeframe);
         }
       } catch (err: any) {
         log.catch(err);

package/src/space/space-manager.ts

@@ -42,7 +42,7 @@ export type ConstructSpaceParams = {
   /**
    * Called when connection auth passed successful.
    */
-  onNetworkConnection: (session: Teleport) => void;
+  onAuthorizedConnection: (session: Teleport) => void;
   onAuthFailure?: (session: Teleport) => void;
 };
 
@@ -93,7 +93,7 @@ export class SpaceManager {
   async constructSpace({
     metadata,
     swarmIdentity,
-    onNetworkConnection,
+    onAuthorizedConnection,
     onAuthFailure,
     memberKey,
   }: ConstructSpaceParams) {
@@ -108,7 +108,7 @@ export class SpaceManager {
       topic: spaceKey,
       swarmIdentity,
       networkManager: this._networkManager,
-      onSessionAuth: onNetworkConnection,
+      onSessionAuth: onAuthorizedConnection,
       onAuthFailure,
       blobStore: this._blobStore,
     });

package/src/space/space-protocol.ts

@@ -19,6 +19,7 @@ import type { FeedMessage } from '@dxos/protocols/proto/dxos/echo/feed';
 import { type MuxerStats, Teleport } from '@dxos/teleport';
 import { type BlobStore, BlobSync } from '@dxos/teleport-extension-object-sync';
 import { ReplicatorExtension } from '@dxos/teleport-extension-replicator';
+import { trace } from '@dxos/tracing';
 import { ComplexMap } from '@dxos/util';
 
 import { AuthExtension, type AuthProvider, type AuthVerifier } from './auth';
@@ -51,6 +52,7 @@ export type SpaceProtocolOptions = {
 /**
  * Manages Teleport protocol stream creation and joining swarms with replication and presence extensions.
  */
+@trace.resource()
 export class SpaceProtocol {
   private readonly _networkManager: NetworkManager;
   private readonly _swarmIdentity: SwarmIdentity;
@@ -60,8 +62,10 @@ export class SpaceProtocol {
   public readonly blobSync: BlobSync;
 
   @logInfo
+  @trace.info()
   private readonly _topic: Promise<PublicKey>;
 
+  @trace.info()
   private readonly _spaceKey: PublicKey;
 
   private readonly _feeds = new Set<FeedWrapper<FeedMessage>>();

package/src/space/space.ts

@@ -2,7 +2,7 @@
 // Copyright 2022 DXOS.org
 //
 
-import { Event, synchronized, trackLeaks, Lock } from '@dxos/async';
+import { Event, synchronized, trackLeaks, Mutex } from '@dxos/async';
 import { type Context } from '@dxos/context';
 import { type FeedInfo } from '@dxos/credentials';
 import { type FeedOptions, type FeedWrapper } from '@dxos/feed-store';
@@ -52,17 +52,22 @@ export type CreatePipelineParams = {
 @trackLeaks('open', 'close')
 @trace.resource()
 export class Space {
-  private readonly _addFeedLock = new Lock();
+  private readonly _addFeedMutex = new Mutex();
 
   public readonly onCredentialProcessed = new Callback<AsyncCallback<Credential>>();
   public readonly stateUpdate = new Event();
+  @trace.info()
   public readonly protocol: SpaceProtocol;
 
   private readonly _key: PublicKey;
   private readonly _genesisFeedKey: PublicKey;
   private readonly _feedProvider: FeedProvider;
+  @trace.info()
   private readonly _controlPipeline: ControlPipeline;
+
+  @trace.info()
   private readonly _dataPipeline: DataPipeline;
+
   private readonly _snapshotManager: SnapshotManager;
 
   private _isOpen = false;
@@ -130,7 +135,7 @@ export class Space {
         }
 
         // Add existing feeds.
-        await this._addFeedLock.executeSynchronized(async () => {
+        await this._addFeedMutex.executeSynchronized(async () => {
           for (const feed of this._controlPipeline.spaceState.feeds.values()) {
             if (feed.assertion.designation === AdmittedFeed.Designation.DATA && !pipeline.hasFeed(feed.key)) {
               await pipeline.addFeed(await this._feedProvider(feed.key, { sparse: true }));

package/src/testing/database-test-rig.ts

@@ -74,7 +74,10 @@ export class DatabaseTestPeer {
 
   private readonly _writes = new Set<WriteRequest>();
 
-  constructor(public readonly rig: DatabaseTestBuilder, public readonly spaceKey: PublicKey) {}
+  constructor(
+    public readonly rig: DatabaseTestBuilder,
+    public readonly spaceKey: PublicKey,
+  ) {}
 
   async open() {
     this.hostItems = new ItemManager(this.modelFactory);

package/src/testing/test-agent-builder.ts

@@ -189,7 +189,7 @@ export class TestAgent {
         credentialAuthenticator: MOCK_AUTH_VERIFIER,
       },
       memberKey: identityKey,
-      onNetworkConnection: (session) => {
+      onAuthorizedConnection: (session) => {
         session.addExtension(
           'dxos.mesh.teleport.gossip',
           this.createGossip().createExtension({ remotePeerId: session.remotePeerId }),