wip: frame-worked tracking state and making authentication RPC call

MatrixAI · Nov 28, 2024 · 3e7b842 · 3e7b842
1 parent 18e0e48
commit 3e7b842
Show file tree

Hide file tree

Showing 7 changed files with 165 additions and 9 deletions.
diff --git a/src/nodes/NodeConnectionManager.ts b/src/nodes/NodeConnectionManager.ts
@@ -52,13 +52,13 @@ type ConnectionAndTimer = {
   connection: NodeConnection;
   timer: Timer | null;
   usageCount: number;
-  authenticatedForward: boolean;
-  authenticatedReverse: boolean;
 };
 
 type ConnectionsEntry = {
   activeConnection: string;
   connections: Record<string, ConnectionAndTimer>;
+  authenticatedForward: boolean;
+  authenticatedReverse: boolean;
 };
 
 type ConnectionInfo = {
@@ -349,6 +349,68 @@ class NodeConnectionManager {
     }
   };
 
+  public forwardAuthenticate(
+    nodeId: NodeId,
+    ctx?: Partial<ContextTimedInput>,
+  ): PromiseCancellable<void>;
+  @timedCancellable(
+    true,
+    (nodeConnectionManager: NodeConnectionManager) =>
+      nodeConnectionManager.connectionConnectTimeoutTime,
+  )
+  public async forwardAuthenticate(
+    nodeId: NodeId,
+    @context ctx: ContextTimed,
+  ): Promise<void> {
+    const targetNodeIdString = nodeId.toString() as NodeIdString;
+    const connectionsEntry = this.connections.get(targetNodeIdString);
+    if (connectionsEntry == null) {
+      throw new nodesErrors.ErrorNodeConnectionManagerConnectionNotFound();
+    }
+    // Need to make an authenticate request here. Get the connection and RPC.
+    await this.withConnF(nodeId, async (conn) => {
+      await conn.rpcClient.methods.nodesAuthenticateConnection({}, ctx);
+    });
+    connectionsEntry.authenticatedForward = true;
+    this.logger.warn(
+      `Node ${nodesUtils.encodeNodeId(nodeId)} has been forward authenticated`,
+    );
+    if (connectionsEntry.authenticatedReverse) {
+      this.logger.warn(
+        `Node ${nodesUtils.encodeNodeId(nodeId)} has been fully authenticated`,
+      );
+    }
+  }
+
+  public handleReverseAuthenticate(
+    nodeId: NodeId,
+    ctx?: Partial<ContextTimedInput>,
+  ): PromiseCancellable<void>;
+  @timedCancellable(
+    true,
+    (nodeConnectionManager: NodeConnectionManager) =>
+      nodeConnectionManager.connectionConnectTimeoutTime,
+  )
+  public async handleReverseAuthenticate(
+    nodeId: NodeId,
+    @context ctx: ContextTimed,
+  ): Promise<void> {
+    const targetNodeIdString = nodeId.toString() as NodeIdString;
+    const connectionsEntry = this.connections.get(targetNodeIdString);
+    if (connectionsEntry == null) {
+      throw new nodesErrors.ErrorNodeConnectionManagerConnectionNotFound();
+    }
+    connectionsEntry.authenticatedReverse = true;
+    this.logger.warn(
+      `Node ${nodesUtils.encodeNodeId(nodeId)} has been reverse authenticated`,
+    );
+    if (connectionsEntry.authenticatedForward) {
+      this.logger.warn(
+        `Node ${nodesUtils.encodeNodeId(nodeId)} has been fully authenticated`,
+      );
+    }
+  }
+
   /**
    * Constructs the `NodeConnectionManager`.
    */
@@ -745,7 +807,6 @@ class NodeConnectionManager {
       },
       ctx,
     );
-    // FIXME: How do we even handle cancellation here if we stop the NCM?
     this.addConnection(nodeConnection.validatedNodeId, nodeConnection);
     // Dispatch the connection event
     const connectionData: ConnectionData = {
@@ -917,8 +978,6 @@ class NodeConnectionManager {
       connection: nodeConnection,
       timer: null,
       usageCount: 0,
-      authenticatedForward: false,
-      authenticatedReverse: false,
     };
 
     // Adding the new connection into the connection map
@@ -931,12 +990,13 @@ class NodeConnectionManager {
           await this.destroyConnection(nodeId, false, connectionId),
         delay: this.getStickyTimeoutValue(nodeId, true),
       });
-      // TODO: only update the active connection once the authentication has been completed.
       entry = {
         activeConnection: connectionId,
         connections: {
           [connectionId]: newConnAndTimer,
         },
+        authenticatedForward: false,
+        authenticatedReverse: false,
       };
       this.connections.set(nodeIdString, entry);
     } else {
@@ -951,11 +1011,12 @@ class NodeConnectionManager {
       // Updating existing entry
       entry.connections[connectionId] = newConnAndTimer;
       // If the new connection ID is less than the old then replace it
-      // TODO: only update the active connection once the authentication has been completed.
       if (entry.activeConnection > connectionId) {
         entry.activeConnection = connectionId;
       }
     }
+    // TODO: this is backgrounded so we need to track this in a map and clean up.
+    this.forwardAuthenticate(nodeId);
     return newConnAndTimer;
   }
 

diff --git a/src/nodes/agent/callers/index.ts b/src/nodes/agent/callers/index.ts
@@ -1,3 +1,4 @@
+import nodesAuthenticateConnection from './nodesAuthenticateConnection';
 import nodesClaimsGet from './nodesClaimsGet';
 import nodesClosestActiveConnectionsGet from './nodesClosestActiveConnectionsGet';
 import nodesClosestLocalNodesGet from './nodesClosestLocalNodesGet';
@@ -15,6 +16,7 @@ import vaultsScan from './vaultsScan';
  * Client manifest
  */
 const manifestClient = {
+  nodesAuthenticateConnection,
   nodesClaimsGet,
   nodesClosestActiveConnectionsGet,
   nodesClosestLocalNodesGet,
@@ -34,6 +36,7 @@ type AgentClientManifest = typeof manifestClient;
 export default manifestClient;
 
 export {
+  nodesAuthenticateConnection,
   nodesClaimsGet,
   nodesClosestActiveConnectionsGet,
   nodesClosestLocalNodesGet,

diff --git a/src/nodes/agent/callers/nodesAuthenticateConnection.ts b/src/nodes/agent/callers/nodesAuthenticateConnection.ts
@@ -0,0 +1,12 @@
+import type { HandlerTypes } from '@matrixai/rpc';
+import type NodesAuthenticateConnection from '../handlers/NodesAuthenticateConnection';
+import { UnaryCaller } from '@matrixai/rpc';
+
+type CallerTypes = HandlerTypes<NodesAuthenticateConnection>;
+
+const nodesAuthenticateConnection = new UnaryCaller<
+  CallerTypes['input'],
+  CallerTypes['output']
+>();
+
+export default nodesAuthenticateConnection;
diff --git a/src/nodes/agent/handlers/NodesAuthenticateConnection.ts b/src/nodes/agent/handlers/NodesAuthenticateConnection.ts
@@ -0,0 +1,43 @@
+import type {
+  AgentRPCRequestParams,
+  AgentRPCResponseResult,
+  SuccessMessage,
+} from '../types';
+import type NodeConnectionManager from '../../../nodes/NodeConnectionManager';
+import type { JSONValue } from '../../../types';
+import type { ContextTimed } from '@matrixai/contexts';
+import { UnaryHandler } from '@matrixai/rpc';
+import * as agentErrors from '../errors';
+import * as agentUtils from '../utils';
+
+class NodesAuthenticateConnection extends UnaryHandler<
+  {
+    nodeConnectionManager: NodeConnectionManager;
+  },
+  AgentRPCRequestParams,
+  AgentRPCResponseResult<SuccessMessage>
+> {
+  public handle = async (
+    _input,
+    _cancel,
+    meta: Record<string, JSONValue> | undefined,
+    ctx: ContextTimed,
+  ): Promise<AgentRPCResponseResult<SuccessMessage>> => {
+    const { nodeConnectionManager } = this.container;
+    // Connections should always be validated
+    const requestingNodeId = agentUtils.nodeIdFromMeta(meta);
+    if (requestingNodeId == null) {
+      throw new agentErrors.ErrorAgentNodeIdMissing();
+    }
+    await nodeConnectionManager.handleReverseAuthenticate(
+      requestingNodeId,
+      ctx,
+    );
+    return {
+      type: 'success',
+      success: true,
+    };
+  };
+}
+
+export default NodesAuthenticateConnection;
diff --git a/src/nodes/agent/handlers/index.ts b/src/nodes/agent/handlers/index.ts
@@ -8,6 +8,7 @@ import type NodeManager from '../../../nodes/NodeManager';
 import type NodeConnectionManager from '../../../nodes/NodeConnectionManager';
 import type NotificationsManager from '../../../notifications/NotificationsManager';
 import type VaultManager from '../../../vaults/VaultManager';
+import NodesAuthenticateConnection from './NodesAuthenticateConnection';
 import NodesClaimsGet from './NodesClaimsGet';
 import NodesClosestActiveConnectionsGet from './NodesClosestActiveConnectionsGet';
 import NodesClosestLocalNodesGet from './NodesClosestLocalNodesGet';
@@ -36,6 +37,7 @@ const manifestServer = (container: {
   vaultManager: VaultManager;
 }) => {
   return {
+    nodesAuthenticateConnection: new NodesAuthenticateConnection(container),
     nodesClaimsGet: new NodesClaimsGet(container),
     nodesClosestActiveConnectionsGet: new NodesClosestActiveConnectionsGet(
       container,
@@ -57,6 +59,7 @@ type AgentServerManifest = ReturnType<typeof manifestServer>;
 export default manifestServer;
 
 export {
+  NodesAuthenticateConnection,
   NodesClaimsGet,
   NodesClosestActiveConnectionsGet,
   NodesClosestLocalNodesGet,

diff --git a/src/nodes/agent/types.ts b/src/nodes/agent/types.ts
@@ -77,6 +77,11 @@ type VaultsScanMessage = VaultInfo & {
   vaultPermissions: Array<VaultAction>;
 };
 
+type SuccessMessage = {
+  type: 'success';
+  success: boolean;
+};
+
 export type {
   AgentRPCRequestParams,
   AgentRPCResponseResult,
@@ -91,4 +96,5 @@ export type {
   SignedNotificationEncoded,
   VaultInfo,
   VaultsScanMessage,
+  SuccessMessage,
 };
diff --git a/tests/nodes/NodeConnectionManager.test.ts b/tests/nodes/NodeConnectionManager.test.ts
@@ -13,6 +13,7 @@ import NodeConnectionManager from '@/nodes/NodeConnectionManager';
 import NodesConnectionSignalFinal from '@/nodes/agent/handlers/NodesConnectionSignalFinal';
 import NodesConnectionSignalInitial from '@/nodes/agent/handlers/NodesConnectionSignalInitial';
 import * as utils from '@/utils';
+import NodesAuthenticateConnection from '@/nodes/agent/handlers/NodesAuthenticateConnection';
 import * as nodesTestUtils from './utils';
 import * as keysTestUtils from '../keys/utils';
 import * as testsUtils from '../utils';
@@ -91,7 +92,13 @@ describe(`${NodeConnectionManager.name}`, () => {
         },
         startOptions: {
           host: localHost,
-          agentService: () => dummyManifest,
+          agentService: (nodeConnectionManager) => {
+            return {
+              nodesAuthenticateConnection: new NodesAuthenticateConnection({
+                nodeConnectionManager,
+              }),
+            } as AgentServerManifest;
+          },
         },
         logger: logger.getChild(`${NodeConnectionManager.name}Local`),
       });
@@ -103,7 +110,13 @@ describe(`${NodeConnectionManager.name}`, () => {
         },
         startOptions: {
           host: localHost,
-          agentService: () => dummyManifest,
+          agentService: (nodeConnectionManager) => {
+            return {
+              nodesAuthenticateConnection: new NodesAuthenticateConnection({
+                nodeConnectionManager,
+              }),
+            } as AgentServerManifest;
+          },
         },
         logger: logger.getChild(`${NodeConnectionManager.name}Peer1`),
       });
@@ -645,6 +658,21 @@ describe(`${NodeConnectionManager.name}`, () => {
         ncmLocal.nodeConnectionManager.hasConnection(ncmPeer1.nodeId),
       ).toBeFalse();
     });
+
+    // TODO: This is a temp test
+    test('can authenticate a connection', async () => {
+      await ncmLocal.nodeConnectionManager.createConnection(
+        [ncmPeer1.nodeId],
+        localHost,
+        ncmPeer1.port,
+      );
+      // Should exist in the map now.
+      expect(
+        ncmLocal.nodeConnectionManager.hasConnection(ncmPeer1.nodeId),
+      ).toBeTrue();
+      await utils.sleep(2000);
+    });
+    test.todo('Cant make most RPC calls while unauthenticated');
   });
   describe('With 2 peers', () => {
     let ncmLocal: NCMState;