diff --git a/src/PolykeyAgent.ts b/src/PolykeyAgent.ts index 4ef549c4ee..14fa1a342c 100644 --- a/src/PolykeyAgent.ts +++ b/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, diff --git a/src/agent/index.ts b/src/agent/index.ts index f45d230fee..4e55eb824b 100644 --- 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 index 45b8f0279b..64b98dd55b 100644 --- a/src/agent/service/echo.ts +++ b/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, callback: grpc.sendUnaryData, ): Promise => { + connectionInfoGetter(call.getPeer()); 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 index 75bb6ee58e..2ab0cbb91a 100644 --- 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,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_), diff --git a/src/agent/types.ts b/src/agent/types.ts new file mode 100644 index 0000000000..ed6023f05f --- /dev/null +++ b/src/agent/types.ts @@ -0,0 +1,7 @@ +import type { ConnectionInfo } from 'network/types'; + +type ConnectionInfoGetter = ( + peerInfo: string, +) => ConnectionInfo | undefined; + +export type { ConnectionInfoGetter }; diff --git a/src/agent/utils.ts b/src/agent/utils.ts new file mode 100644 index 0000000000..0a690159e4 --- /dev/null +++ b/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 './types'; + +function connectionInfoGetter(revProxy: ReverseProxy): ConnectionInfoGetter { + return (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 }; diff --git a/tests/agent/GRPCClientAgent.test.ts b/tests/agent/GRPCClientAgent.test.ts index 4cc7ecb898..3ddf166089 100644 --- a/tests/agent/GRPCClientAgent.test.ts +++ b/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()); + }); + }); }); diff --git a/tests/agent/utils.ts b/tests/agent/utils.ts index 7dd702d91c..6bfda2465c 100644 --- 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 { +async function openTestAgentClient( + port: number, + nodeId?: NodeId, + proxyConfig?: ProxyConfig, +): Promise { 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 index 22e596f8e6..f00d8237eb 100644 --- a/tests/nodes/NodeConnection.test.ts +++ b/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, diff --git a/tests/notifications/NotificationsManager.test.ts b/tests/notifications/NotificationsManager.test.ts index 24b024a18e..485b3dc541 100644 --- a/tests/notifications/NotificationsManager.test.ts +++ b/tests/notifications/NotificationsManager.test.ts @@ -210,6 +210,7 @@ describe('NotificationsManager', () => { notificationsManager: receiverNotificationsManager, acl: receiverACL, gestaltGraph: receiverGestaltGraph, + revProxy: receiverRevProxy, }); agentServer = new GRPCServer({ logger: logger,