Fixes #342 - added connectionInfoGetter utility

src/PolykeyAgent.ts

@@ -563,6 +563,7 @@ class PolykeyAgent {
         notificationsManager: this.notificationsManager,
         acl: this.acl,
         gestaltGraph: this.gestaltGraph,
+        revProxy: this.revProxy,
       });
       const clientService = createClientService({
         pkAgent: this,

src/agent/index.ts

@@ -1,3 +1,5 @@
 export { default as createAgentService, AgentServiceService } from './service';
 export { default as GRPCClientAgent } from './GRPCClientAgent';
 export * as errors from './errors';
+export * as types from './types';
+export * as utils from './utils';

src/agent/service/echo.ts

@@ -1,11 +1,17 @@
 import type * as grpc from '@grpc/grpc-js';
+import type { ConnectionInfoGetter } from 'agent/types';
 import * as utilsPB from '../../proto/js/polykey/v1/utils/utils_pb';
 
-function echo(_) {
+function echo({
+  connectionInfoGetter,
+}: {
+  connectionInfoGetter: ConnectionInfoGetter;
+}) {
   return async (
     call: grpc.ServerUnaryCall<utilsPB.EchoMessage, utilsPB.EchoMessage>,
     callback: grpc.sendUnaryData<utilsPB.EchoMessage>,
   ): Promise<void> => {
+    await connectionInfoGetter(call.getPeer());
     const response = new utilsPB.EchoMessage();
     response.setChallenge(call.request.getChallenge());
     callback(null, response);

src/agent/service/index.ts

@@ -10,6 +10,7 @@ import type { Sigchain } from '../../sigchain';
 import type { ACL } from '../../acl';
 import type { GestaltGraph } from '../../gestalts';
 import type { IAgentServiceServer } from '../../proto/js/polykey/v1/agent_service_grpc_pb';
+import type ReverseProxy from '../../network/ReverseProxy';
 import echo from './echo';
 import nodesChainDataGet from './nodesChainDataGet';
 import nodesClaimsGet from './nodesClaimsGet';
@@ -21,6 +22,7 @@ import vaultsGitInfoGet from './vaultsGitInfoGet';
 import vaultsGitPackGet from './vaultsGitPackGet';
 import vaultsScan from './vaultsScan';
 import { AgentServiceService } from '../../proto/js/polykey/v1/agent_service_grpc_pb';
+import * as agentUtils from '../utils';
 
 function createService(container: {
   keyManager: KeyManager;
@@ -32,9 +34,14 @@ function createService(container: {
   sigchain: Sigchain;
   acl: ACL;
   gestaltGraph: GestaltGraph;
+  revProxy: ReverseProxy;
 }): IAgentServiceServer {
+  const connectionInfoGetter = agentUtils.connectionInfoGetter(
+    container.revProxy,
+  );
   const container_ = {
     ...container,
+    connectionInfoGetter,
   };
   const service: IAgentServiceServer = {
     echo: echo(container_),

src/agent/types.ts

@@ -0,0 +1,7 @@
+import type { ConnectionInfo } from 'network/types';
+
+type ConnectionInfoGetter = (
+  peerInfo: string,
+) => Promise<ConnectionInfo | undefined>;
+
+export type { ConnectionInfoGetter };

src/agent/utils.ts

@@ -0,0 +1,16 @@
+import type { Host, Port } from 'network/types';
+import type ReverseProxy from 'network/ReverseProxy';
+import type { ConnectionInfoGetter } from 'agent/types';
+
+function connectionInfoGetter(revProxy: ReverseProxy): ConnectionInfoGetter {
+  return async (peerInfo: string) => {
+    const address = peerInfo.split(':');
+    const host = address[0] as Host;
+    const port = parseInt(address[1]) as Port;
+    // Return undefined in input was invalid
+    if (host == null || isNaN(port)) return;
+    return revProxy.getConnectionInfoByProxy(host, port);
+  };
+}
+
+export { connectionInfoGetter };

tests/agent/GRPCClientAgent.test.ts

@@ -7,21 +7,24 @@ import os from 'os';
 import path from 'path';
 import Logger, { LogLevel, StreamHandler } from '@matrixai/logger';
 import { DB } from '@matrixai/db';
-import { GRPCClientAgent } from '@/agent';
+import { errors as agentErrors, GRPCClientAgent } from '@/agent';
 import { KeyManager } from '@/keys';
-import { NodeConnectionManager, NodeGraph, NodeManager } from '@/nodes';
+import {
+  NodeConnectionManager,
+  NodeGraph,
+  NodeManager,
+  utils as nodesUtils,
+} from '@/nodes';
 import { VaultManager } from '@/vaults';
 import { Sigchain } from '@/sigchain';
 import { ACL } from '@/acl';
 import { GestaltGraph } from '@/gestalts';
-import { errors as agentErrors } from '@/agent';
 import { ForwardProxy, ReverseProxy } from '@/network';
 import { NotificationsManager } from '@/notifications';
-import { utils as claimsUtils, errors as claimsErrors } from '@/claims';
+import { errors as claimsErrors, utils as claimsUtils } from '@/claims';
 import * as keysUtils from '@/keys/utils';
 import * as utilsPB from '@/proto/js/polykey/v1/utils/utils_pb';
 import * as nodesPB from '@/proto/js/polykey/v1/nodes/nodes_pb';
-import { utils as nodesUtils } from '@/nodes';
 import { RWLock } from '@/utils';
 import * as testAgentUtils from './utils';
 import * as testUtils from '../utils';
@@ -171,12 +174,19 @@ describe(GRPCClientAgent.name, () => {
       notificationsManager,
       acl,
       gestaltGraph,
+      revProxy,
+    });
+    await revProxy.start({
+      serverHost: '127.0.0.1' as Host,
+      serverPort: port as Port,
+      tlsConfig: tlsConfig,
     });
     client = await testAgentUtils.openTestAgentClient(port);
   }, global.polykeyStartupTimeout);
   afterEach(async () => {
     await testAgentUtils.closeTestAgentClient(client);
     await testAgentUtils.closeTestAgentServer(server);
+    await revProxy.stop();
     await vaultManager.stop();
     await notificationsManager.stop();
     await sigchain.stop();
@@ -408,4 +418,59 @@ describe(GRPCClientAgent.name, () => {
       global.defaultTimeout * 4,
     );
   });
+  describe('With connection through proxies', () => {
+    let clientWithProxies: GRPCClientAgent;
+    let clientFwdProxy: ForwardProxy;
+    let clientKeyManager: KeyManager;
+
+    beforeEach(async () => {
+      clientFwdProxy = new ForwardProxy({
+        authToken: 'auth',
+        logger,
+      });
+
+      clientKeyManager = await KeyManager.createKeyManager({
+        keysPath: path.join(dataDir, 'clientKeys'),
+        password: 'password',
+        logger,
+      });
+      const clientTlsConfig: TLSConfig = {
+        keyPrivatePem: clientKeyManager.getRootKeyPairPem().privateKey,
+        certChainPem: await clientKeyManager.getRootCertChainPem(),
+      };
+      await clientFwdProxy.start({
+        tlsConfig: clientTlsConfig,
+      });
+      clientWithProxies = await testAgentUtils.openTestAgentClient(
+        revProxy.getIngressPort(),
+        clientKeyManager.getNodeId(),
+        {
+          host: clientFwdProxy.getProxyHost(),
+          port: clientFwdProxy.getProxyPort(),
+          authToken: clientFwdProxy.authToken,
+        },
+      );
+    });
+
+    afterEach(async () => {
+      await testAgentUtils.closeTestAgentClient(clientWithProxies);
+      await clientFwdProxy.stop();
+      await clientKeyManager.stop();
+    });
+
+    test('connectionInfoGetter is called and returns the expected information', async () => {
+      // We can't directly spy on the connectionInfoGetter result
+      // but we can check that it called `getConnectionInfoByProxy` properly
+      const getConnectionInfoByProxySpy = jest.spyOn(
+        ReverseProxy.prototype,
+        'getConnectionInfoByProxy',
+      );
+      await clientWithProxies.echo(new utilsPB.EchoMessage());
+      // It should've returned the expected information
+      const returnedInfo = getConnectionInfoByProxySpy.mock.results[0].value;
+      expect(returnedInfo.ingressPort).toEqual(revProxy.getIngressPort());
+      expect(returnedInfo.egressPort).toEqual(clientFwdProxy.getEgressPort());
+      expect(returnedInfo.nodeId).toStrictEqual(clientKeyManager.getNodeId());
+    });
+  });
 });

tests/agent/utils.ts

@@ -1,4 +1,4 @@
-import type { Host, Port } from '@/network/types';
+import type { Host, Port, ProxyConfig } from '@/network/types';
 
 import type { IAgentServiceServer } from '@/proto/js/polykey/v1/agent_service_grpc_pb';
 import type { KeyManager } from '@/keys';
@@ -8,6 +8,8 @@ import type { Sigchain } from '@/sigchain';
 import type { NotificationsManager } from '@/notifications';
 import type { ACL } from '@/acl';
 import type { GestaltGraph } from '@/gestalts';
+import type { NodeId } from 'nodes/types';
+import type { ReverseProxy } from 'network/index';
 import Logger, { LogLevel, StreamHandler } from '@matrixai/logger';
 import * as grpc from '@grpc/grpc-js';
 import { promisify } from '@/utils';
@@ -28,6 +30,7 @@ async function openTestAgentServer({
   notificationsManager,
   acl,
   gestaltGraph,
+  revProxy,
 }: {
   keyManager: KeyManager;
   vaultManager: VaultManager;
@@ -38,6 +41,7 @@ async function openTestAgentServer({
   notificationsManager: NotificationsManager;
   acl: ACL;
   gestaltGraph: GestaltGraph;
+  revProxy: ReverseProxy;
 }) {
   const agentService: IAgentServiceServer = createAgentService({
     keyManager,
@@ -49,6 +53,7 @@ async function openTestAgentServer({
     nodeConnectionManager,
     acl,
     gestaltGraph,
+    revProxy,
   });
 
   const server = new grpc.Server();
@@ -67,16 +72,21 @@ async function closeTestAgentServer(server) {
   await tryShutdown();
 }
 
-async function openTestAgentClient(port: number): Promise<GRPCClientAgent> {
+async function openTestAgentClient(
+  port: number,
+  nodeId?: NodeId,
+  proxyConfig?: ProxyConfig,
+): Promise<GRPCClientAgent> {
   const logger = new Logger('AgentClientTest', LogLevel.WARN, [
     new StreamHandler(),
   ]);
   const agentClient = await GRPCClientAgent.createGRPCClientAgent({
-    nodeId: testUtils.generateRandomNodeId(),
+    nodeId: nodeId ?? testUtils.generateRandomNodeId(),
     host: '127.0.0.1' as Host,
     port: port as Port,
     logger: logger,
     destroyCallback: async () => {},
+    proxyConfig,
     timeout: 30000,
   });
   return agentClient;

tests/nodes/NodeConnection.test.ts

@@ -281,6 +281,7 @@ describe(`${NodeConnection.name} test`, () => {
       notificationsManager: serverNotificationsManager,
       acl: serverACL,
       gestaltGraph: serverGestaltGraph,
+      revProxy: serverRevProxy,
     });
     agentServer = new GRPCServer({
       logger: logger,

tests/notifications/NotificationsManager.test.ts

@@ -210,6 +210,7 @@ describe('NotificationsManager', () => {
       notificationsManager: receiverNotificationsManager,
       acl: receiverACL,
       gestaltGraph: receiverGestaltGraph,
+      revProxy: receiverRevProxy,
     });
     agentServer = new GRPCServer({
       logger: logger,