Add optional exec websocket keepalive with lifecycle safeguards

src/exec.ts

@@ -7,6 +7,15 @@ import { KubeConfig } from './config.js';
 import { isResizable, ResizableStream, TerminalSizeQueue } from './terminal-size-queue.js';
 import { WebSocketHandler, WebSocketInterface } from './web-socket-handler.js';
 
+export interface ExecOptions {
+    /**
+     * Optional websocket keepalive interval in milliseconds (1 to 2147483647).
+     * Only one ping is outstanding until a pong is received; this is not a pong timeout.
+     * Ignored when the socket does not support ping and Node-style event listeners.
+     */
+    pingIntervalMs?: number;
+}
+
 export class Exec {
     public 'handler': WebSocketInterface;
 
@@ -27,6 +36,7 @@ export class Exec {
      * @param {boolean} tty - Should the command execute in a TTY enabled session.
      * @param {(V1Status) => void} statusCallback -
      *       A callback to received the status (e.g. exit code) from the command, optional.
+     * @param {ExecOptions} options - Optional websocket keepalive settings.
      * @return {Promise<WebSocket>} A promise that will return the web socket created for this command.
      */
     public async exec(
@@ -39,7 +49,16 @@ export class Exec {
         stdin: stream.Readable | null,
         tty: boolean,
         statusCallback?: (status: V1Status) => void,
+        options?: ExecOptions,
     ): Promise<WebSocket.WebSocket> {
+        const pingIntervalMs = options?.pingIntervalMs;
+        if (
+            pingIntervalMs !== undefined &&
+            (!Number.isInteger(pingIntervalMs) || pingIntervalMs <= 0 || pingIntervalMs > 2147483647)
+        ) {
+            // Validate before connecting so invalid options cannot start a remote command.
+            throw new Error('pingIntervalMs must be an integer between 1 and 2147483647');
+        }
         const query = {
             stdout: stdout != null,
             stderr: stderr != null,
@@ -68,6 +87,45 @@ export class Exec {
             WebSocketHandler.handleStandardInput(conn, this.terminalSizeQueue, WebSocketHandler.ResizeStream);
             this.terminalSizeQueue.handleResizes(stdout as any as ResizableStream);
         }
+        if (pingIntervalMs !== undefined) {
+            this.setupPing(conn, pingIntervalMs);
+        }
         return conn;
     }
+
+    private setupPing(conn: WebSocket.WebSocket, pingIntervalMs: number): void {
+        if (
+            typeof conn.ping !== 'function' ||
+            typeof conn.on !== 'function' ||
+            typeof conn.removeListener !== 'function' ||
+            conn.readyState !== WebSocket.OPEN
+        ) {
+            return;
+        }
+
+        let awaitingPong = false;
+        const onPong = () => {
+            awaitingPong = false;
+        };
+        const clearKeepAlive = () => {
+            clearInterval(timer);
+            conn.removeListener('pong', onPong);
+            conn.removeListener('close', clearKeepAlive);
+            conn.removeListener('error', clearKeepAlive);
+        };
+        const timer = setInterval(() => {
+            if (conn.readyState !== WebSocket.OPEN) {
+                clearKeepAlive();
+                return;
+            }
+            if (!awaitingPong) {
+                awaitingPong = true;
+                conn.ping();
+            }
+        }, pingIntervalMs);
+        conn.on('pong', onPong);
+        conn.on('close', clearKeepAlive);
+        conn.on('error', clearKeepAlive);
+        timer.unref();
+    }
 }

src/exec_test.ts

@@ -1,17 +1,21 @@
-import { describe, it } from 'node:test';
-import { deepStrictEqual, strictEqual } from 'node:assert';
+import { beforeEach, describe, it } from 'node:test';
+import { deepStrictEqual, ok, rejects, strictEqual } from 'node:assert';
+import { EventEmitter, once } from 'node:events';
 import WebSocket from 'isomorphic-ws';
 import { ReadableStreamBuffer, WritableStreamBuffer } from 'stream-buffers';
 import { anyFunction, anything, capture, instance, mock, verify, when } from 'ts-mockito';
 
 import { CallAwaiter, matchBuffer, ResizableWriteableStreamBuffer } from './test/index.js';
 import { V1Status } from './api.js';
 import { KubeConfig } from './config.js';
-import { Exec } from './exec.js';
+import { Exec, ExecOptions } from './exec.js';
 import { TerminalSize } from './terminal-size-queue.js';
 import { WebSocketHandler, WebSocketInterface } from './web-socket-handler.js';
 
 describe('Exec', () => {
+    const startExec = (exec: Exec, options?: ExecOptions) =>
+        exec.exec('ns', 'pod', 'container', 'command', null, null, null, false, undefined, options);
+
     describe('basic', () => {
         it('should correctly exec to a url', async () => {
             const kc = new KubeConfig();
@@ -157,4 +161,189 @@ describe('Exec', () => {
             verify(fakeWebSocket.close()).called();
         });
     });
+
+    describe('keepalive', () => {
+        let conn: EventEmitter & { readyState: number; ping: () => void };
+        let pingCount: number;
+        let connectCount: number;
+        let exec: Exec;
+        const start = (options?: ExecOptions) => startExec(exec, options);
+
+        beforeEach((t) => {
+            ok('mock' in t);
+            t.mock.timers.enable({ apis: ['setInterval'] });
+            pingCount = 0;
+            connectCount = 0;
+            conn = Object.assign(new EventEmitter(), {
+                readyState: WebSocket.OPEN as number,
+                ping: () => {
+                    pingCount++;
+                },
+            });
+            exec = new Exec(new KubeConfig(), {
+                connect: async () => {
+                    connectCount++;
+                    return conn as unknown as WebSocket.WebSocket;
+                },
+            });
+        });
+
+        it('sends one ping at a time and resumes after pong', async (t) => {
+            await start({ pingIntervalMs: 10 });
+            t.mock.timers.tick(9);
+            strictEqual(pingCount, 0);
+            t.mock.timers.tick(1);
+            strictEqual(pingCount, 1);
+            t.mock.timers.tick(100);
+            strictEqual(pingCount, 1);
+            conn.emit('pong');
+            t.mock.timers.tick(10);
+            strictEqual(pingCount, 2);
+            t.mock.timers.tick(100);
+            strictEqual(pingCount, 2);
+        });
+
+        const disabledCases: { name: string; options?: ExecOptions; socket?: object }[] = [
+            { name: 'omitted options' },
+            { name: 'empty options', options: {} },
+            { name: 'undefined interval', options: { pingIntervalMs: undefined } },
+            { name: 'missing ping', socket: { ping: undefined } },
+            { name: 'non-function on', socket: { on: true } },
+            { name: 'missing removeListener', socket: { removeListener: undefined } },
+            { name: 'connecting socket', socket: { readyState: WebSocket.CONNECTING } },
+            { name: 'closing socket', socket: { readyState: WebSocket.CLOSING } },
+            { name: 'closed socket', socket: { readyState: WebSocket.CLOSED } },
+        ];
+        for (const { name, options, socket } of disabledCases) {
+            it(`keeps pings disabled for ${name}`, async (t) => {
+                Object.assign(conn, socket);
+                const interval = t.mock.method(globalThis, 'setInterval');
+                await start(socket ? { pingIntervalMs: 10 } : options);
+                t.mock.timers.tick(100);
+                strictEqual(pingCount, 0);
+                strictEqual(interval.mock.callCount(), 0);
+                deepStrictEqual(conn.eventNames(), []);
+            });
+        }
+
+        for (const value of [0, -1, 1.5, NaN, Infinity, 2147483648, '10', null]) {
+            it(`rejects invalid interval ${String(value)} before connecting`, async () => {
+                await rejects(
+                    start({ pingIntervalMs: value as number }),
+                    /pingIntervalMs must be an integer/,
+                );
+                strictEqual(connectCount, 0);
+            });
+        }
+
+        for (const pingIntervalMs of [1, 2147483647]) {
+            it(`accepts interval boundary ${pingIntervalMs}`, async (t) => {
+                await start({ pingIntervalMs });
+                t.mock.timers.tick(pingIntervalMs - 1);
+                strictEqual(pingCount, 0);
+                t.mock.timers.tick(1);
+                strictEqual(pingCount, 1);
+            });
+        }
+
+        for (const event of ['close', 'error']) {
+            it(`cleans up its timer and listeners on ${event}`, async (t) => {
+                const externalListener = () => {};
+                for (const name of ['pong', 'close', 'error']) {
+                    conn.on(name, externalListener);
+                }
+                const clear = t.mock.method(globalThis, 'clearInterval');
+                await start({ pingIntervalMs: 10 });
+                t.mock.timers.tick(10);
+                strictEqual(pingCount, 1);
+                conn.emit(event);
+                strictEqual(clear.mock.callCount(), 1);
+                for (const name of ['pong', 'close', 'error']) {
+                    deepStrictEqual(conn.listeners(name), [externalListener]);
+                }
+                conn.emit('pong');
+                t.mock.timers.tick(100);
+                strictEqual(pingCount, 1);
+                conn.emit('close');
+                strictEqual(clear.mock.callCount(), 1);
+            });
+        }
+
+        it('cleans up if the socket stops being open without a close event', async (t) => {
+            await start({ pingIntervalMs: 10 });
+            conn.readyState = WebSocket.CLOSING;
+            const clear = t.mock.method(globalThis, 'clearInterval');
+            t.mock.timers.tick(10);
+            strictEqual(clear.mock.callCount(), 1);
+            strictEqual(pingCount, 0);
+            deepStrictEqual(conn.eventNames(), []);
+        });
+
+        it('handles a synchronous pong without getting stuck', async (t) => {
+            conn.ping = () => {
+                pingCount++;
+                conn.emit('pong');
+            };
+            await start({ pingIntervalMs: 10 });
+            t.mock.timers.tick(10);
+            t.mock.timers.tick(10);
+            strictEqual(pingCount, 2);
+        });
+
+        it('keeps concurrent sessions independent', async (t) => {
+            await start({ pingIntervalMs: 10 });
+            const firstConn = conn;
+            conn = Object.assign(new EventEmitter(), {
+                readyState: WebSocket.OPEN as number,
+                ping: conn.ping,
+            });
+            await start({ pingIntervalMs: 10 });
+            t.mock.timers.tick(10);
+            strictEqual(pingCount, 2);
+            firstConn.emit('close');
+            conn.emit('pong');
+            t.mock.timers.tick(10);
+            strictEqual(pingCount, 3);
+            strictEqual(conn.listenerCount('pong'), 1);
+            deepStrictEqual(firstConn.eventNames(), []);
+        });
+    });
+
+    it('exchanges keepalive ping/pong frames with a real websocket', { timeout: 5000 }, async (t) => {
+        const server = new WebSocket.Server({ port: 0 });
+        let client: WebSocket.WebSocket | undefined;
+        t.after(() => {
+            client?.terminate();
+            for (const socket of server.clients) {
+                socket.terminate();
+            }
+            server.close();
+        });
+        await once(server, 'listening');
+        const address = server.address();
+        ok(typeof address === 'object' && address !== null);
+        let pingCount = 0;
+        server.on('connection', (socket) => {
+            socket.on('ping', () => {
+                pingCount++;
+            });
+        });
+        const exec = new Exec(new KubeConfig(), {
+            connect: async () => {
+                client = new WebSocket(`ws://127.0.0.1:${address.port}`);
+                await once(client, 'open');
+                return client;
+            },
+        });
+        const conn = await startExec(exec, { pingIntervalMs: 10 });
+        await once(conn, 'pong');
+        await once(conn, 'pong');
+        strictEqual(pingCount, 2);
+        const closed = once(conn, 'close');
+        conn.close();
+        await closed;
+        strictEqual(conn.listenerCount('pong'), 0);
+        strictEqual(conn.listenerCount('close'), 0);
+        strictEqual(conn.listenerCount('error'), 0);
+    });
 });