Fixes #342 - added connectionInfoGetter utility

MatrixAI · Mar 9, 2022 · 202c444 · 202c444
1 parent 204d3b2
commit 202c444
Show file tree

Hide file tree

Showing 10 changed files with 188 additions and 14 deletions.
diff --git a/src/PolykeyAgent.ts b/src/PolykeyAgent.ts
@@ -564,6 +564,7 @@ class PolykeyAgent {
         notificationsManager: this.notificationsManager,
         acl: this.acl,
         gestaltGraph: this.gestaltGraph,
+        revProxy: this.revProxy,
       });
       const clientService = createClientService({
         pkAgent: this,

diff --git a/src/agent/errors.ts b/src/agent/errors.ts
@@ -1,4 +1,4 @@
-import { ErrorPolykey } from '../errors';
+import { ErrorPolykey, sysexits } from '../errors';
 
 class ErrorAgent extends ErrorPolykey {}
 
@@ -8,9 +8,15 @@ class ErrorAgentClientNotStarted extends ErrorAgent {}
 
 class ErrorAgentClientDestroyed extends ErrorAgent {}
 
+class ErrorConnectionInfoMissing extends ErrorAgent {
+  description = 'Vault already exists';
+  exitCode = sysexits.UNAVAILABLE;
+}
+
 export {
   ErrorAgent,
   ErrorAgentClientNotStarted,
   ErrorAgentRunning,
   ErrorAgentClientDestroyed,
+  ErrorConnectionInfoMissing,
 };
diff --git a/src/agent/index.ts b/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';
diff --git a/src/agent/service/echo.ts b/src/agent/service/echo.ts
@@ -1,11 +1,13 @@
 import type * as grpc from '@grpc/grpc-js';
+import type { ConnectionInfoGet } from 'agent/types';
 import * as utilsPB from '../../proto/js/polykey/v1/utils/utils_pb';
 
-function echo(_) {
+function echo({ connectionInfoGet }: { connectionInfoGet: ConnectionInfoGet }) {
   return async (
     call: grpc.ServerUnaryCall<utilsPB.EchoMessage, utilsPB.EchoMessage>,
     callback: grpc.sendUnaryData<utilsPB.EchoMessage>,
   ): Promise<void> => {
+    connectionInfoGet(call);
     const response = new utilsPB.EchoMessage();
     response.setChallenge(call.request.getChallenge());
     callback(null, response);

diff --git a/src/agent/service/index.ts b/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,12 @@ function createService(container: {
   sigchain: Sigchain;
   acl: ACL;
   gestaltGraph: GestaltGraph;
+  revProxy: ReverseProxy;
 }): IAgentServiceServer {
+  const connectionInfoGet = agentUtils.connectionInfoGetter(container.revProxy);
   const container_ = {
     ...container,
+    connectionInfoGet: connectionInfoGet,
   };
   const service: IAgentServiceServer = {
     echo: echo(container_),

diff --git a/src/agent/types.ts b/src/agent/types.ts
@@ -0,0 +1,8 @@
+import type { ConnectionInfo } from 'network/types';
+import type { ServerSurfaceCall } from '@grpc/grpc-js/build/src/server-call';
+
+type ConnectionInfoGet = (
+  call: ServerSurfaceCall,
+) => ConnectionInfo | undefined;
+
+export type { ConnectionInfoGet };
diff --git a/src/agent/utils.ts b/src/agent/utils.ts
@@ -0,0 +1,18 @@
+import type { Host, Port } from 'network/types';
+import type ReverseProxy from 'network/ReverseProxy';
+import type { ConnectionInfoGet } from './types';
+import type { ServerSurfaceCall } from '@grpc/grpc-js/build/src/server-call';
+
+function connectionInfoGetter(revProxy: ReverseProxy): ConnectionInfoGet {
+  return (call: ServerSurfaceCall) => {
+    let urlString = call.getPeer();
+    if (!/^.*:\/\//.test(urlString)) urlString = 'pk://' + urlString;
+    const url = new URL(urlString);
+    return revProxy.getConnectionInfoByProxy(
+      url.hostname as Host,
+      parseInt(url.port) as Port,
+    );
+  };
+}
+
+export { connectionInfoGetter };
diff --git a/tests/agent/GRPCClientAgent.test.ts b/tests/agent/GRPCClientAgent.test.ts
@@ -1,5 +1,6 @@
-import type { TLSConfig } from '@/network/types';
+import type { Host, Port, TLSConfig } from '@/network/types';
 import type * as grpc from '@grpc/grpc-js';
+import type { NodeId } from '@/nodes/types';
 import fs from 'fs';
 import path from 'path';
 import os from 'os';
@@ -21,26 +22,20 @@ import * as utilsPB from '@/proto/js/polykey/v1/utils/utils_pb';
 import * as agentErrors from '@/agent/errors';
 import * as keysUtils from '@/keys/utils';
 import * as testAgentUtils from './utils';
-import * as testUtils from '../utils';
 
 describe(GRPCClientAgent.name, () => {
+  const host = '127.0.0.1' as Host;
   const password = 'password';
   const logger = new Logger(`${GRPCClientAgent.name} test`, LogLevel.WARN, [
     new StreamHandler(),
   ]);
-  let mockedGenerateKeyPair: jest.SpyInstance;
   let mockedGenerateDeterministicKeyPair: jest.SpyInstance;
   beforeAll(async () => {
-    const globalKeyPair = await testUtils.setupGlobalKeypair();
-    mockedGenerateKeyPair = jest
-      .spyOn(keysUtils, 'generateKeyPair')
-      .mockResolvedValue(globalKeyPair);
     mockedGenerateDeterministicKeyPair = jest
       .spyOn(keysUtils, 'generateDeterministicKeyPair')
-      .mockResolvedValue(globalKeyPair);
+      .mockImplementation((bits, _) => keysUtils.generateKeyPair(bits));
   });
   afterAll(async () => {
-    mockedGenerateKeyPair.mockRestore();
     mockedGenerateDeterministicKeyPair.mockRestore();
   });
   let client: GRPCClientAgent;
@@ -85,6 +80,8 @@ describe(GRPCClientAgent.name, () => {
     });
     await fwdProxy.start({
       tlsConfig,
+      egressHost: host,
+      proxyHost: host,
     });
     revProxy = new ReverseProxy({
       logger: logger,
@@ -168,12 +165,20 @@ describe(GRPCClientAgent.name, () => {
       notificationsManager,
       acl,
       gestaltGraph,
+      revProxy,
+    });
+    await revProxy.start({
+      ingressHost: host,
+      serverHost: host,
+      serverPort: port as Port,
+      tlsConfig: tlsConfig,
     });
     client = await testAgentUtils.openTestAgentClient(port);
   }, global.defaultTimeout);
   afterEach(async () => {
     await testAgentUtils.closeTestAgentClient(client);
     await testAgentUtils.closeTestAgentServer(server);
+    await revProxy.stop();
     await vaultManager.stop();
     await notificationsManager.stop();
     await sigchain.stop();
@@ -210,4 +215,120 @@ describe(GRPCClientAgent.name, () => {
     expect(response.getChallenge()).toBe('yes');
     expect(client.secured).toBeFalsy();
   });
+  describe('With connection through proxies', () => {
+    const logger = new Logger(`${GRPCClientAgent.name} test`, LogLevel.WARN, [
+      new StreamHandler(),
+    ]);
+    const localHost = '127.0.0.1' as Host;
+
+    let clientWithProxies1: GRPCClientAgent;
+    let clientFwdProxy1: ForwardProxy;
+    let clientKeyManager1: KeyManager;
+    let nodeId1: NodeId;
+
+    let clientWithProxies2: GRPCClientAgent;
+    let clientFwdProxy2: ForwardProxy;
+    let clientKeyManager2: KeyManager;
+    let nodeId2: NodeId;
+
+    beforeEach(async () => {
+      dataDir = await fs.promises.mkdtemp(
+        path.join(os.tmpdir(), 'polykey-test-'),
+      );
+      // Setting up clients
+      clientFwdProxy1 = new ForwardProxy({
+        authToken: 'auth',
+        logger,
+      });
+      clientKeyManager1 = await KeyManager.createKeyManager({
+        keysPath: path.join(dataDir, 'clientKeys1'),
+        password: 'password',
+        logger,
+      });
+      nodeId1 = clientKeyManager1.getNodeId();
+      await clientFwdProxy1.start({
+        tlsConfig: {
+          keyPrivatePem: clientKeyManager1.getRootKeyPairPem().privateKey,
+          certChainPem: await clientKeyManager1.getRootCertChainPem(),
+        },
+        egressHost: localHost,
+        proxyHost: localHost,
+      });
+      clientWithProxies1 = await GRPCClientAgent.createGRPCClientAgent({
+        host: localHost,
+        nodeId: keyManager.getNodeId(),
+        port: revProxy.getIngressPort(),
+        proxyConfig: {
+          host: clientFwdProxy1.getProxyHost(),
+          port: clientFwdProxy1.getProxyPort(),
+          authToken: clientFwdProxy1.authToken,
+        },
+        timeout: 5000,
+        logger,
+      });
+
+      clientFwdProxy2 = new ForwardProxy({
+        authToken: 'auth',
+        logger,
+      });
+      clientKeyManager2 = await KeyManager.createKeyManager({
+        keysPath: path.join(dataDir, 'clientKeys2'),
+        password: 'password',
+        logger,
+      });
+      nodeId2 = clientKeyManager2.getNodeId();
+      await clientFwdProxy2.start({
+        tlsConfig: {
+          keyPrivatePem: clientKeyManager2.getRootKeyPairPem().privateKey,
+          certChainPem: await clientKeyManager2.getRootCertChainPem(),
+        },
+        egressHost: localHost,
+        proxyHost: localHost,
+      });
+      clientWithProxies2 = await GRPCClientAgent.createGRPCClientAgent({
+        host: '127.0.0.1' as Host,
+        logger,
+        nodeId: keyManager.getNodeId(),
+        port: revProxy.getIngressPort(),
+        proxyConfig: {
+          host: clientFwdProxy2.getProxyHost(),
+          port: clientFwdProxy2.getProxyPort(),
+          authToken: clientFwdProxy2.authToken,
+        },
+        timeout: 5000,
+      });
+    }, 26000);
+    afterEach(async () => {
+      await testAgentUtils.closeTestAgentClient(clientWithProxies1);
+      await clientFwdProxy1.stop();
+      await clientKeyManager1.stop();
+      await testAgentUtils.closeTestAgentClient(clientWithProxies2);
+      await clientFwdProxy2.stop();
+      await clientKeyManager2.stop();
+    }, 25000);
+    test('connectionInfoGetter returns correct information for each connection', 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 clientWithProxies1.echo(new utilsPB.EchoMessage());
+      await clientWithProxies2.echo(new utilsPB.EchoMessage());
+      // It should've returned the expected information
+      const returnedInfo1 = getConnectionInfoByProxySpy.mock.results[0].value;
+      expect(returnedInfo1.ingressPort).toEqual(revProxy.getIngressPort());
+      expect(returnedInfo1.ingressHost).toEqual(localHost);
+      expect(returnedInfo1.egressPort).toEqual(clientFwdProxy1.getEgressPort());
+      expect(returnedInfo1.egressHost).toEqual(localHost);
+      expect(returnedInfo1.nodeId).toStrictEqual(nodeId1);
+      // Checking second call
+      const returnedInfo2 = getConnectionInfoByProxySpy.mock.results[1].value;
+      expect(returnedInfo2.ingressPort).toEqual(revProxy.getIngressPort());
+      expect(returnedInfo2.ingressHost).toEqual(localHost);
+      expect(returnedInfo2.egressPort).toEqual(clientFwdProxy2.getEgressPort());
+      expect(returnedInfo2.egressHost).toEqual(localHost);
+      expect(returnedInfo2.nodeId).toStrictEqual(nodeId2);
+    });
+  });
 });
diff --git a/tests/agent/utils.ts b/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;

diff --git a/tests/nodes/NodeConnection.test.ts b/tests/nodes/NodeConnection.test.ts
@@ -282,6 +282,7 @@ describe(`${NodeConnection.name} test`, () => {
       notificationsManager: serverNotificationsManager,
       acl: serverACL,
       gestaltGraph: serverGestaltGraph,
+      revProxy: serverRevProxy,
     });
     agentServer = new GRPCServer({
       logger: logger,