Allow agent to claim custom networks

jest.config.mjs

@@ -1,6 +1,6 @@
 import path from 'node:path';
 import url from 'node:url';
-import tsconfigJSON from './tsconfig.json' assert { type: "json" };
+import tsconfigJSON from './tsconfig.json' with { type: "json" };
 
 const projectPath = path.dirname(url.fileURLToPath(import.meta.url));
 

src/PolykeyAgent.ts

@@ -774,6 +774,7 @@ class PolykeyAgent {
         });
         const initialNodesShortlist = initialNodes.slice(0, 3);
         await this.nodeManager.syncNodeGraph(
+          options.network,
           initialNodesShortlist,
           undefined,
           false,

src/client/callers/index.ts

@@ -46,6 +46,7 @@ import nodesFind from './nodesFind.js';
 import nodesGetAll from './nodesGetAll.js';
 import nodesListConnections from './nodesListConnections.js';
 import nodesPing from './nodesPing.js';
+import nodesSyncGraph from './nodesSyncGraph.js';
 import notificationsInboxClear from './notificationsInboxClear.js';
 import notificationsInboxRead from './notificationsInboxRead.js';
 import notificationsInboxRemove from './notificationsInboxRemove.js';
@@ -128,6 +129,7 @@ const clientManifest = {
   nodesGetAll,
   nodesListConnections,
   nodesPing,
+  nodesSyncGraph,
   notificationsInboxClear,
   notificationsInboxRead,
   notificationsInboxRemove,
@@ -209,6 +211,7 @@ export {
   nodesGetAll,
   nodesListConnections,
   nodesPing,
+  nodesSyncGraph,
   notificationsInboxClear,
   notificationsInboxRead,
   notificationsInboxRemove,

src/client/callers/nodesSyncGraph.ts

@@ -0,0 +1,12 @@
+import type { HandlerTypes } from '@matrixai/rpc';
+import type NodesSyncGraph from '../handlers/NodesSyncGraph.js';
+import { UnaryCaller } from '@matrixai/rpc';
+
+type CallerTypes = HandlerTypes<NodesSyncGraph>;
+
+const nodesSyncGraph = new UnaryCaller<
+  CallerTypes['input'],
+  CallerTypes['output']
+>();
+
+export default nodesSyncGraph;

src/client/handlers/NodesSyncGraph.ts

@@ -0,0 +1,48 @@
+import type { ContextTimed } from '@matrixai/contexts';
+import type { JSONValue } from '@matrixai/rpc';
+import type NodeManager from '../../nodes/NodeManager.js';
+import type {
+  ClientRPCRequestParams,
+  ClientRPCResponseResult,
+  NodesSyncGraphMessage,
+} from '../types.js';
+import type { AgentClientManifest } from '../../nodes/agent/callers/index.js';
+import type { NodeId, NodeAddress } from '../../nodes/types.js';
+import { UnaryHandler } from '@matrixai/rpc';
+import * as nodesUtils from '../../nodes/utils.js';
+
+class NodesSyncGraph extends UnaryHandler<
+  {
+    nodeManager: NodeManager<AgentClientManifest>;
+  },
+  ClientRPCRequestParams<NodesSyncGraphMessage>,
+  ClientRPCResponseResult
+> {
+  public handle = async (
+    input: ClientRPCRequestParams<NodesSyncGraphMessage>,
+    _cancel: (reason?: any) => void,
+    _meta: Record<string, JSONValue> | undefined,
+    ctx: ContextTimed,
+  ): Promise<ClientRPCResponseResult> => {
+    const {
+      nodeManager,
+    }: {
+      nodeManager: NodeManager<AgentClientManifest>;
+    } = this.container;
+    // Convert the encoded node id to the binary one we expect
+    const parsedInitialNodes = input.initialNodes.map(
+      (value) =>
+        [nodesUtils.decodeNodeId(value[0]), value[1]] as [NodeId, NodeAddress],
+    );
+    await nodeManager.syncNodeGraph(
+      input.network,
+      parsedInitialNodes,
+      input.connectionTimeout,
+      true,
+      ctx,
+    );
+    return {};
+  };
+}
+
+export default NodesSyncGraph;

src/client/handlers/index.ts

@@ -64,6 +64,7 @@ import NodesFind from './NodesFind.js';
 import NodesGetAll from './NodesGetAll.js';
 import NodesListConnections from './NodesListConnections.js';
 import NodesPing from './NodesPing.js';
+import NodesSyncGraph from './NodesSyncGraph.js';
 import NotificationsInboxClear from './NotificationsInboxClear.js';
 import NotificationsInboxRead from './NotificationsInboxRead.js';
 import NotificationsInboxRemove from './NotificationsInboxRemove.js';
@@ -169,6 +170,7 @@ const serverManifest = (container: {
     nodesGetAll: new NodesGetAll(container),
     nodesListConnections: new NodesListConnections(container),
     nodesPing: new NodesPing(container),
+    nodesSyncGraph: new NodesSyncGraph(container),
     notificationsInboxClear: new NotificationsInboxClear(container),
     notificationsInboxRead: new NotificationsInboxRead(container),
     notificationsInboxRemove: new NotificationsInboxRemove(container),
@@ -252,6 +254,7 @@ export {
   NodesGetAll,
   NodesListConnections,
   NodesPing,
+  NodesSyncGraph,
   NotificationsInboxClear,
   NotificationsInboxRead,
   NotificationsInboxRemove,

src/client/types.ts

@@ -151,6 +151,12 @@ type NodeConnectionMessage = NodeAddressMessage & {
   authenticated: boolean;
 };
 
+type NodesSyncGraphMessage = {
+  network: string;
+  initialNodes: Array<[NodeIdEncoded, NodeAddress]>;
+  connectionTimeout?: number;
+};
+
 // Gestalts messages
 
 type ActionsListMessage = {
@@ -428,6 +434,7 @@ export type {
   NodeAddressMessage,
   NodesFindMessage,
   NodeConnectionMessage,
+  NodesSyncGraphMessage,
   ActionsListMessage,
   SetIdentityActionMessage,
   SetNodeActionMessage,

src/nodes/NodeManager.ts

@@ -253,6 +253,7 @@ class NodeManager<Manifest extends AgentClientManifestNodeManager> {
   protected syncNodeGraphHandler = async (
     ctx: ContextTimed,
     _taskInfo: TaskInfo | undefined,
+    network: string | undefined,
     initialNodes: Array<[NodeIdEncoded, NodeAddress]>,
     connectionConnectTimeoutTime: number | undefined,
   ) => {
@@ -299,6 +300,14 @@ class NodeManager<Manifest extends AgentClientManifestNodeManager> {
     }
     if (ctx.signal.aborted) return;
 
+    if (network != null) {
+      if ((await this.getClaimNetworkAccess(network)) == null) {
+        await this.claimNetwork(successfulConnections[0].value.nodeId, network);
+      } else {
+        await this.switchNetwork(network);
+      }
+    }
+
     // Attempt a findNode operation looking for ourselves
     await this.findNode(
       {
@@ -1748,7 +1757,7 @@ class NodeManager<Manifest extends AgentClientManifestNodeManager> {
   /**
    * This returns the `ClaimNetworkAccess` for the given network.
    */
-  protected async getClaimNetworkAccess(
+  public async getClaimNetworkAccess(
     network: string,
     tran?: DBTransaction,
   ): Promise<Token<ClaimNetworkAccess> | undefined> {
@@ -1875,6 +1884,12 @@ class NodeManager<Manifest extends AgentClientManifestNodeManager> {
       network,
       targetNodeId,
     );
+
+    // Error out if a network access claim already exists
+    if ((await this.getClaimNetworkAccess(network, tran)) != null) {
+      throw new Error('TMP network access claim already exists');
+    }
+
     const encodedNetworkAuthority = claimsUtils.generateSignedClaim(
       claimNetworkAuthority.toSigned(),
     );
@@ -2787,6 +2802,7 @@ class NodeManager<Manifest extends AgentClientManifestNodeManager> {
    *
    */
   public syncNodeGraph(
+    network: string | undefined,
     initialNodes: Array<[NodeId, NodeAddress]>,
     connectionConnectTimeoutTime?: number,
     blocking?: boolean,
@@ -2795,6 +2811,7 @@ class NodeManager<Manifest extends AgentClientManifestNodeManager> {
   @startStop.ready(new nodesErrors.ErrorNodeManagerNotRunning())
   @decorators.timedCancellable(true)
   public async syncNodeGraph(
+    network: string,
     initialNodes: Array<[NodeId, NodeAddress]>,
     connectionConnectTimeoutTime: number = this.connectionConnectTimeoutTime,
     blocking: boolean = false,
@@ -2817,6 +2834,7 @@ class NodeManager<Manifest extends AgentClientManifestNodeManager> {
       await this.syncNodeGraphHandler(
         ctx,
         undefined,
+        network,
         initialNodesParameter,
         connectionConnectTimeoutTime,
       );
@@ -2826,7 +2844,11 @@ class NodeManager<Manifest extends AgentClientManifestNodeManager> {
         delay: 0,
         handlerId: this.syncNodeGraphHandlerId,
         lazy: true,
-        parameters: [initialNodesParameter, connectionConnectTimeoutTime],
+        parameters: [
+          network,
+          initialNodesParameter,
+          connectionConnectTimeoutTime,
+        ],
         path: [this.tasksPath, this.syncNodeGraphHandlerId],
         priority: 0,
       });