tunnel server: memoize ping

tunnel-server/src/events.test.ts

@@ -54,7 +54,7 @@ describe('onceWithTimeout', () => {
     })
 
     describe('when a fallback is specified', () => {
-      let p: Promise<12>
+      let p: Promise<12 | void>
       beforeEach(() => {
         p = onceWithTimeout(emitter, 'foo', { milliseconds: 10, fallback: async () => 12 as const })
         jest.advanceTimersByTime(10)

tunnel-server/src/events.ts

@@ -14,7 +14,7 @@ export async function onceWithTimeout <T = unknown>(
   target: NodeEventTarget,
   event: string | symbol,
   opts: { milliseconds: number; fallback: () => T | Promise<T> },
-): Promise<T>
+): Promise<T | void>
 export async function onceWithTimeout <T = unknown>(
   target: NodeEventTarget,
   event: string | symbol,

tunnel-server/src/memoize.test.ts

@@ -0,0 +1,64 @@
+import { afterAll, beforeAll, beforeEach, describe, it, expect, jest } from '@jest/globals'
+import { memoizeForDuration } from './memoize'
+
+describe('memoizeForDuration', () => {
+  beforeAll(() => {
+    jest.useFakeTimers()
+  })
+  afterAll(() => {
+    jest.useRealTimers()
+  })
+
+  let fn: jest.Mock<() => number>
+  let memoized: () => number
+
+  beforeEach(() => {
+    fn = jest.fn(() => 12)
+    memoized = memoizeForDuration(fn, 1000)
+  })
+
+  describe('before the first call', () => {
+    it('does not call the specified function', () => {
+      expect(fn).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('on the first call', () => {
+    let v: number
+    beforeEach(() => {
+      v = memoized()
+    })
+    it('calls the specified function', () => {
+      expect(fn).toHaveBeenCalledTimes(1)
+    })
+    it('returns the memoized value', () => {
+      expect(v).toBe(12)
+    })
+
+    describe('on the second call, when the expiry duration has not passed', () => {
+      beforeEach(() => {
+        jest.advanceTimersByTime(999)
+        v = memoized()
+      })
+      it('does not call the specified function again', () => {
+        expect(fn).toHaveBeenCalledTimes(1)
+      })
+      it('returns the memoized value', () => {
+        expect(v).toBe(12)
+      })
+    })
+
+    describe('on the second call, when the expiry duration has passed', () => {
+      beforeEach(() => {
+        jest.advanceTimersByTime(1000)
+        v = memoized()
+      })
+      it('calls the specified function again', () => {
+        expect(fn).toHaveBeenCalledTimes(2)
+      })
+      it('returns the memoized value', () => {
+        expect(v).toBe(12)
+      })
+    })
+  })
+})

tunnel-server/src/memoize.ts

@@ -0,0 +1,9 @@
+export const memoizeForDuration = <T>(f: () => T, milliseconds: number) => {
+  let cache: { value: T; expiry: number } | undefined
+  return () => {
+    if (!cache || cache.expiry <= Date.now()) {
+      cache = { value: f(), expiry: Date.now() + milliseconds }
+    }
+    return cache.value
+  }
+}

tunnel-server/src/ssh/base-server.ts

@@ -2,14 +2,14 @@ import crypto, { createPublicKey, randomBytes } from 'crypto'
 import { FastifyBaseLogger } from 'fastify/types/logger'
 import net from 'net'
 import path from 'path'
-import events from 'events'
 import ssh2, { SocketBindInfo } from 'ssh2'
 import { inspect } from 'util'
 import { EventEmitter, IEventEmitter } from 'tseep'
 import { calculateJwkThumbprintUri, exportJWK } from 'jose'
 import { ForwardRequest, parseForwardRequest } from '../forward-request'
 import { createDestroy } from '../destroy-server'
 import { onceWithTimeout } from '../events'
+import { memoizeForDuration } from '../memoize'
 
 const clientIdFromPublicSsh = (key: Buffer) =>
   crypto.createHash('sha1').update(key).digest('base64url').replace(/[_-]/g, '')
@@ -59,7 +59,7 @@ export interface BaseSshClient extends IEventEmitter<BaseSshClientEvents> {
   publicKey: crypto.KeyObject
   publicKeyThumbprint: string
   end: () => Promise<void>
-  ping: (timeoutMs: number) => Promise<boolean>
+  ping: () => Promise<boolean>
   connectionId: string
   log: FastifyBaseLogger
 }
@@ -101,19 +101,23 @@ export const baseSshServer = (
       let ended = false
       const end = async () => {
         if (!ended) {
-          client.end()
-          await events.once(client, 'end')
+          await new Promise(resolve => {
+            client.once('end', resolve)
+            client.end()
+          })
         }
       }
 
       let authContext: ssh2.AuthContext
       let key: ssh2.ParsedKey
 
-      const ping = async (milliseconds: number) => {
-        const result = onceWithTimeout(client, 'rekey', { milliseconds, fallback: () => 'timeout' as const })
+      const REKEY_TIMEOUT = 5000
+      const ping = memoizeForDuration(async () => {
+        const result = onceWithTimeout(client, 'rekey', { milliseconds: REKEY_TIMEOUT })
+          .then(() => 'pong' as const, () => 'error' as const)
         client.rekey()
-        return await result !== 'timeout'
-      }
+        return (await result) === 'pong'
+      }, REKEY_TIMEOUT)
 
       client
         .on('authentication', async ctx => {

tunnel-server/src/ssh/index.ts

@@ -56,7 +56,7 @@ export const createSshServer = ({
           if (otherClient.connectionId === connectionId) {
             throw new Error(`duplicate path: ${key}, from same connection ${connectionId}`)
           }
-          if (!await otherClient.ping(5000)) {
+          if (!await otherClient.ping()) {
             const existingDelete = onceWithTimeout(existingEntry.watcher, 'delete', { milliseconds: 2000 })
             void otherClient.end()
             await existingDelete