From 03df31f3595ae928f1c5de58d2dd5854eecda0a8 Mon Sep 17 00:00:00 2001 From: Roy Razon Date: Thu, 9 Nov 2023 09:25:53 +0200 Subject: [PATCH] tunnel server: ssh connections gauge --- tunnel-server/index.ts | 3 +- tunnel-server/src/metrics.ts | 6 ++ tunnel-server/src/multimap.test.ts | 125 +++++++++++++++-------------- tunnel-server/src/ssh/index.ts | 39 +++++---- 4 files changed, 95 insertions(+), 78 deletions(-) diff --git a/tunnel-server/index.ts b/tunnel-server/index.ts index e5aa2a78..5fb50d8c 100644 --- a/tunnel-server/index.ts +++ b/tunnel-server/index.ts @@ -9,7 +9,7 @@ import { activeTunnelStoreKey, inMemoryActiveTunnelStore } from './src/tunnel-st import { getSSHKeys } from './src/ssh-keys' import { proxy } from './src/proxy' import { appLoggerFromEnv } from './src/logging' -import { tunnelsGauge, runMetricsServer } from './src/metrics' +import { tunnelsGauge, runMetricsServer, sshConnectionsGauge } from './src/metrics' import { numberFromEnv, requiredEnv } from './src/env' import { editUrl } from './src/url' import { cookieSessionStore } from './src/session' @@ -83,6 +83,7 @@ const sshServer = createSshServer({ rootUrl: BASE_URL.toString(), }, tunnelsGauge, + sshConnectionsGauge, tunnelUrl: (clientId, remotePath) => tunnelUrl(BASE_URL, clientId, remotePath), }) .listen(SSH_PORT, LISTEN_HOST, () => { diff --git a/tunnel-server/src/metrics.ts b/tunnel-server/src/metrics.ts index 1dc47fce..0a318e4f 100644 --- a/tunnel-server/src/metrics.ts +++ b/tunnel-server/src/metrics.ts @@ -1,6 +1,12 @@ import fastify from 'fastify' import { Gauge, Counter, register } from 'prom-client' +export const sshConnectionsGauge = new Gauge({ + name: 'sshConnections', + help: 'Current number of open SSH connections', + labelNames: ['envId'], +}) + export const tunnelsGauge = new Gauge({ name: 'tunnels', help: 'Current number of open tunnels', diff --git a/tunnel-server/src/multimap.test.ts b/tunnel-server/src/multimap.test.ts index 8e1a2f84..382d7339 100644 --- a/tunnel-server/src/multimap.test.ts +++ b/tunnel-server/src/multimap.test.ts @@ -4,91 +4,94 @@ import { MultiMap, multimap } from './multimap' describe('multimap', () => { type ObjType = { x: number } let a: MultiMap - const expectedValues = [{ x: 12 }, { x: 13 }] as const - beforeEach(() => { - a = multimap() - a.add('foo', expectedValues[0]) - a.add('foo', expectedValues[1]) - }) - - describe('when the key does not exist', () => { - it('returns undefined', () => { - expect(multimap().get('bar')).toBeUndefined() - }) - }) - describe('when the key exists', () => { - let values: readonly ObjType[] | undefined + describe('after adding a value with length 2', () => { + const expectedValues = [{ x: 12 }, { x: 13 }] as const beforeEach(() => { - values = a.get('foo') - }) - it('returns the values', () => { - expect(values).toBeDefined() - expect(values).toHaveLength(2) - expect(values).toContain(expectedValues[0]) - expect(values).toContain(expectedValues[1]) + a = multimap() + a.add('foo', expectedValues[0]) + a.add('foo', expectedValues[1]) }) - describe('when the returned array is mutated', () => { - beforeEach(() => { - (values as ObjType[]).push({ x: 14 }) - }) - it('does not affect the multimap', () => { - expect(a.get('foo')).toHaveLength(2) + describe('when getting a key which does not exist', () => { + it('returns undefined', () => { + expect(multimap().get('bar')).toBeUndefined() }) }) - describe('when delete is called with a predicate that returns false for everything', () => { - let deleteReturn: boolean + describe('when getting a key which exists', () => { + let values: readonly ObjType[] | undefined beforeEach(() => { - deleteReturn = a.delete('foo', () => false) values = a.get('foo') }) - - it('returns false', () => { - expect(deleteReturn).toBe(false) - }) - - it('does not delete the values', () => { + it('returns the values', () => { expect(values).toBeDefined() expect(values).toHaveLength(2) expect(values).toContain(expectedValues[0]) expect(values).toContain(expectedValues[1]) }) - }) - describe('when delete is called with a predicate that returns true for everything', () => { - let deleteReturn: boolean - beforeEach(() => { - deleteReturn = a.delete('foo', () => true) - values = a.get('foo') + describe('when the returned array is mutated', () => { + beforeEach(() => { + (values as ObjType[]).push({ x: 14 }) + }) + it('does not affect the multimap', () => { + expect(a.get('foo')).toHaveLength(2) + }) }) - it('returns true', () => { - expect(deleteReturn).toBe(true) - }) + describe('when delete is called with a predicate that returns false for everything', () => { + let deleteReturn: boolean + beforeEach(() => { + deleteReturn = a.delete('foo', () => false) + values = a.get('foo') + }) - it('deletes the values', () => { - expect(values).toBeUndefined() - }) - }) + it('returns false', () => { + expect(deleteReturn).toBe(false) + }) - describe('when delete is called with a predicate that returns true for a specific value', () => { - let deleteReturn: boolean - beforeEach(() => { - deleteReturn = a.delete('foo', ({ x }) => x === expectedValues[0].x) - values = a.get('foo') + it('does not delete the values', () => { + expect(values).toBeDefined() + expect(values).toHaveLength(2) + expect(values).toContain(expectedValues[0]) + expect(values).toContain(expectedValues[1]) + }) }) - it('returns true', () => { - expect(deleteReturn).toBe(true) + describe('when delete is called with a predicate that returns true for everything', () => { + let deleteReturn: boolean + beforeEach(() => { + deleteReturn = a.delete('foo', () => true) + values = a.get('foo') + }) + + it('returns true', () => { + expect(deleteReturn).toBe(true) + }) + + it('deletes the values', () => { + expect(values).toBeUndefined() + }) }) - it('deletes the specific value', () => { - expect(values).toBeDefined() - expect(values).toHaveLength(1) - expect(values).not.toContain(expectedValues[0]) - expect(values).toContain(expectedValues[1]) + describe('when delete is called with a predicate that returns true for a specific value', () => { + let deleteReturn: boolean + beforeEach(() => { + deleteReturn = a.delete('foo', ({ x }) => x === expectedValues[0].x) + values = a.get('foo') + }) + + it('returns true', () => { + expect(deleteReturn).toBe(true) + }) + + it('deletes the specific value', () => { + expect(values).toBeDefined() + expect(values).toHaveLength(1) + expect(values).not.toContain(expectedValues[0]) + expect(values).toContain(expectedValues[1]) + }) }) }) }) diff --git a/tunnel-server/src/ssh/index.ts b/tunnel-server/src/ssh/index.ts index f89f2356..ce444046 100644 --- a/tunnel-server/src/ssh/index.ts +++ b/tunnel-server/src/ssh/index.ts @@ -1,6 +1,7 @@ import { Logger } from 'pino' import { inspect } from 'util' import { Gauge } from 'prom-client' +import lodash from 'lodash' import { BaseSshClient, baseSshServer } from './base-server' import { ActiveTunnelStore, activeTunnelStoreKey } from '../tunnel-store' import { KeyAlreadyExistsError } from '../memory-store' @@ -14,6 +15,7 @@ export const createSshServer = ({ tunnelUrl, helloBaseResponse, tunnelsGauge, + sshConnectionsGauge, }: { log: Logger sshPrivateKey: string @@ -22,10 +24,12 @@ export const createSshServer = ({ tunnelUrl: (clientId: string, remotePath: string) => string helloBaseResponse: Record tunnelsGauge: Pick + sshConnectionsGauge: Pick }) => { - const storeKeyToClient = new Map() const onClient = (client: BaseSshClient) => { const { clientId, publicKey, envId, connectionId, publicKeyThumbprint, log } = client + sshConnectionsGauge.inc({ envId }) + const cleanupClient = lodash.once(() => { sshConnectionsGauge.inc({ envId }) }) const tunnels = new Map() client .on('forward', async (requestId, { path: tunnelPath, access, meta, inject }, localSocketPath, accept, reject) => { @@ -67,7 +71,7 @@ export const createSshServer = ({ const setResult = await set().catch(err => { reject(err) }) if (!setResult) { - return undefined + return } const { tx: setTx } = setResult const forward = await accept().catch(async e => { @@ -75,27 +79,30 @@ export const createSshServer = ({ await activeTunnelStore.delete(key, setTx) }) if (!forward) { - return undefined + return } - storeKeyToClient.set(key, client) tunnels.set(requestId, tunnelUrl(clientId, tunnelPath)) - const onForwardClose = (event: 'close' | 'error') => (err?: Error) => { - if (err) { - log.info('%s: deleting tunnel %s due to forward server error: %j', event, key, inspect(err)) - } else { - log.info('%s: deleting tunnel %s', event, key) - } + const cleanupTunnel = lodash.once(() => { tunnels.delete(requestId) - storeKeyToClient.delete(key) void activeTunnelStore.delete(key, setTx) tunnelsGauge.dec({ clientId }) - } - forward.on('close', onForwardClose('close')) - forward.on('error', onForwardClose('error')) + }) + forward + .on('close', () => { + log.info('forward close: deleting tunnel %s', key) + cleanupTunnel() + }) + .on('error', err => { + log.info('forward error: deleting tunnel %s due to forward server error: %j', key, inspect(err)) + cleanupTunnel() + }) tunnelsGauge.inc({ clientId }) - return undefined }) - .on('error', err => { log.warn('client error %j: %j', clientId, inspect(err)) }) + .on('end', () => { cleanupClient() }) + .on('error', err => { + log.warn('client error %j: %j', clientId, inspect(err)) + cleanupClient() + }) .on('exec', (command, respondWithJson, reject) => { if (command === 'hello') { respondWithJson({