From 2de4e81c352ccca009f617411bc7a5782125c96f 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 c01027a795..b57d56dd5f 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';
@@ -29,6 +29,7 @@ import { createClientService, ClientServiceService } from './client';
 import config from './config';
 import * as utils from './utils';
 import * as errors from './errors';
+import * as networkUtils from '@/network/utils';
 
 type NetworkConfig = {
   forwardHost?: Host;
@@ -55,8 +56,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({
@@ -260,6 +263,10 @@ class PolykeyAgent {
         proxy ??
         new Proxy({
           ...proxyConfig_,
+          connectionEstablishedCallback: (data) => events.emitAsync(
+            PolykeyAgent.eventSymbols.Proxy,
+            data
+          ),
           logger: logger.getChild(Proxy.name),
         });
       nodeGraph =
@@ -537,6 +544,25 @@ 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 c2f0231c2d..5313dd9abf 100644
--- a/tests/nodes/NodeManager.test.ts
+++ b/tests/nodes/NodeManager.test.ts
@@ -19,6 +19,7 @@ import * as claimsUtils from '@/claims/utils';
 import { promisify, sleep } from '@/utils';
 import * as nodesUtils from '@/nodes/utils';
 import * as nodesTestUtils from './utils';
+import * as utilsPB from '@/proto/js/polykey/v1/utils/utils_pb'
 
 describe(`${NodeManager.name} test`, () => {
   const password = 'password';
@@ -563,4 +564,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();
+      let 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();
+    }
+  })
 });