From a36cd3a83dcc11fcea1023907ec9ec1641bd1c7d Mon Sep 17 00:00:00 2001
From: Brian Botha <brian.botha@matrix.ai>
Date: Thu, 24 Mar 2022 18:10:28 +1100
Subject: [PATCH] feat: `Proxy` trigger adding nodes to `Nodegraph`

Added an event to the `EventBus` that is triggered by the `Proxy`'s `connectionEstablishedCallback`. this adds the node to the `NodeGraph`.

Related #344
---
 src/PolykeyAgent.ts             | 28 +++++++++++++++++++++-
 tests/nodes/NodeManager.test.ts | 42 +++++++++++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/src/PolykeyAgent.ts b/src/PolykeyAgent.ts
index 68e3b5571..a3f5241f3 100644
--- a/src/PolykeyAgent.ts
+++ b/src/PolykeyAgent.ts
@@ -1,6 +1,6 @@
 import type { FileSystem } from './types';
 import type { PolykeyWorkerManagerInterface } from './workers/types';
-import type { Host, Port } from './network/types';
+import type { ConnectionData, Host, Port } from './network/types';
 import type { SeedNodes } from './nodes/types';
 import type { KeyManagerChangeData } from './keys/types';
 import path from 'path';
@@ -8,6 +8,7 @@ import process from 'process';
 import Logger from '@matrixai/logger';
 import { DB } from '@matrixai/db';
 import { CreateDestroyStartStop } from '@matrixai/async-init/dist/CreateDestroyStartStop';
+import * as networkUtils from '@/network/utils';
 import KeyManager from './keys/KeyManager';
 import Status from './status/Status';
 import Schema from './schema/Schema';
@@ -59,8 +60,10 @@ class PolykeyAgent {
    */
   public static readonly eventSymbols = {
     [KeyManager.name]: Symbol(KeyManager.name),
+    [Proxy.name]: Symbol(Proxy.name),
   } as {
     readonly KeyManager: unique symbol;
+    readonly Proxy: unique symbol;
   };
 
   public static async createPolykeyAgent({
@@ -266,6 +269,8 @@ class PolykeyAgent {
         proxy ??
         new Proxy({
           ...proxyConfig_,
+          connectionEstablishedCallback: (data) =>
+            events.emitAsync(PolykeyAgent.eventSymbols.Proxy, data),
           logger: logger.getChild(Proxy.name),
         });
       nodeGraph =
@@ -542,6 +547,27 @@ class PolykeyAgent {
           this.logger.info(`${KeyManager.name} change propagated`);
         },
       );
+      this.events.on(
+        PolykeyAgent.eventSymbols.Proxy,
+        async (data: ConnectionData) => {
+          if (data.type === 'reverse') {
+            const address = networkUtils.buildAddress(
+              data.remoteHost,
+              data.remotePort,
+            );
+            const nodeIdEncoded = nodesUtils.encodeNodeId(data.remoteNodeId);
+            this.logger.info(
+              `Reverse connection adding ${nodeIdEncoded}:${address} to ${NodeGraph.name}`,
+            );
+            // Reverse connection was established and authenticated,
+            //  add it to the node graph
+            await this.nodeManager.setNode(data.remoteNodeId, {
+              host: data.remoteHost,
+              port: data.remotePort,
+            });
+          }
+        },
+      );
       const networkConfig_ = {
         ...config.defaults.networkConfig,
         ...utils.filterEmptyObject(networkConfig),
diff --git a/tests/nodes/NodeManager.test.ts b/tests/nodes/NodeManager.test.ts
index d74df5c6a..8d2c249b8 100644
--- a/tests/nodes/NodeManager.test.ts
+++ b/tests/nodes/NodeManager.test.ts
@@ -18,6 +18,7 @@ import Sigchain from '@/sigchain/Sigchain';
 import * as claimsUtils from '@/claims/utils';
 import { promisify, sleep } from '@/utils';
 import * as nodesUtils from '@/nodes/utils';
+import * as utilsPB from '@/proto/js/polykey/v1/utils/utils_pb';
 import * as nodesTestUtils from './utils';
 
 describe(`${NodeManager.name} test`, () => {
@@ -604,4 +605,45 @@ describe(`${NodeManager.name} test`, () => {
     expect(oldestNodeNew).toBeUndefined();
     nodeManagerPingMock.mockRestore();
   });
+  test('should add node when an incoming connection is established', async () => {
+    let server: PolykeyAgent | undefined;
+    try {
+      server = await PolykeyAgent.createPolykeyAgent({
+        password: 'password',
+        nodePath: path.join(dataDir, 'server'),
+        keysConfig: {
+          rootKeyPairBits: 2048,
+        },
+        logger: logger,
+      });
+      const serverNodeId = server.keyManager.getNodeId();
+      const serverNodeAddress: NodeAddress = {
+        host: server.proxy.getProxyHost(),
+        port: server.proxy.getProxyPort(),
+      };
+      await nodeGraph.setNode(serverNodeId, serverNodeAddress);
+
+      const expectedHost = proxy.getProxyHost();
+      const expectedPort = proxy.getProxyPort();
+      const expectedNodeId = keyManager.getNodeId();
+
+      const nodeData = await server.nodeGraph.getNode(expectedNodeId);
+      expect(nodeData).toBeUndefined();
+
+      // Now we want to connect to the server by making an echo request.
+      await nodeConnectionManager.withConnF(serverNodeId, async (conn) => {
+        const client = conn.getClient();
+        await client.echo(new utilsPB.EchoMessage().setChallenge('hello'));
+      });
+
+      const nodeData2 = await server.nodeGraph.getNode(expectedNodeId);
+      expect(nodeData2).toBeDefined();
+      expect(nodeData2?.address.host).toEqual(expectedHost);
+      expect(nodeData2?.address.port).toEqual(expectedPort);
+    } finally {
+      // Clean up
+      await server?.stop();
+      await server?.destroy();
+    }
+  });
 });