diff --git a/src/attach.ts b/src/attach.ts index 3c2f747bb6c..2bb94b95279 100644 --- a/src/attach.ts +++ b/src/attach.ts @@ -23,6 +23,7 @@ export class Attach { stderr: stream.Writable | any, stdin: stream.Readable | any, tty: boolean, + done?: (err: any) => void, ): Promise { const query = { container: containerName, @@ -33,18 +34,21 @@ export class Attach { }; const queryStr = querystring.stringify(query); const path = `/api/v1/namespaces/${namespace}/pods/${podName}/attach?${queryStr}`; - const conn = await this.handler.connect(path, null, (streamNum: number, buff: Buffer): boolean => { - WebSocketHandler.handleStandardStreams(streamNum, buff, stdout, stderr); - return true; - }); - if (stdin != null) { - WebSocketHandler.handleStandardInput(conn, stdin, WebSocketHandler.StdinStream); - } + let resizeStream: stream.Readable | null = null; if (isResizable(stdout)) { this.terminalSizeQueue = new TerminalSizeQueue(); - WebSocketHandler.handleStandardInput(conn, this.terminalSizeQueue, WebSocketHandler.ResizeStream); + resizeStream = this.terminalSizeQueue; this.terminalSizeQueue.handleResizes(stdout as any as ResizableStream); } - return conn; + return WebSocketHandler.connectStandardStreams( + this.handler, + path, + stdout, + stderr, + stdin, + resizeStream, + undefined, + done, + ); } } diff --git a/src/cp.ts b/src/cp.ts index f8920e908d8..6db9e1f5c8d 100644 --- a/src/cp.ts +++ b/src/cp.ts @@ -1,4 +1,5 @@ import { WritableStreamBuffer } from 'stream-buffers'; +import { finished } from 'node:stream/promises'; import tar from 'tar-fs'; import { KubeConfig } from './config.js'; @@ -33,21 +34,37 @@ export class Cp { command.push(srcPath); const writerStream = tar.extract(tgtPath); const errStream = new WritableStreamBuffer(); - this.execInstance.exec( - namespace, - podName, - containerName, - command, - writerStream, - errStream, - null, - false, - async () => { - if (errStream.size()) { - throw new Error(`Error from cpFromPod - details: \n ${errStream.getContentsAsString()}`); - } - }, - ); + const remoteDone = new Promise((resolve, reject) => { + this.execInstance + .exec( + namespace, + podName, + containerName, + command, + writerStream, + errStream, + null, + false, + undefined, + (err: any) => { + if (err) { + reject(err); + return; + } + if (errStream.size()) { + reject( + new Error( + `Error from cpFromPod - details: \n ${errStream.getContentsAsString()}`, + ), + ); + return; + } + resolve(); + }, + ) + .catch(reject); + }); + await Promise.all([remoteDone, finished(writerStream)]); } /** @@ -67,20 +84,36 @@ export class Cp { const command = ['tar', 'xf', '-', '-C', tgtPath]; const readStream = tar.pack(srcPath); const errStream = new WritableStreamBuffer(); - this.execInstance.exec( - namespace, - podName, - containerName, - command, - null, - errStream, - readStream, - false, - async () => { - if (errStream.size()) { - throw new Error(`Error from cpToPod - details: \n ${errStream.getContentsAsString()}`); - } - }, - ); + const remoteDone = new Promise((resolve, reject) => { + this.execInstance + .exec( + namespace, + podName, + containerName, + command, + null, + errStream, + readStream, + false, + undefined, + (err: any) => { + if (err) { + reject(err); + return; + } + if (errStream.size()) { + reject( + new Error( + `Error from cpToPod - details: \n ${errStream.getContentsAsString()}`, + ), + ); + return; + } + resolve(); + }, + ) + .catch(reject); + }); + await Promise.all([remoteDone, finished(readStream)]); } } diff --git a/src/cp_test.ts b/src/cp_test.ts index fb38247df34..4c4ef0632da 100644 --- a/src/cp_test.ts +++ b/src/cp_test.ts @@ -1,9 +1,14 @@ import { describe, it } from 'node:test'; -import { anything, anyFunction, instance, mock, verify, when } from 'ts-mockito'; +import { rejects } from 'node:assert'; +import { anyFunction, anything, capture, instance, mock, verify, when } from 'ts-mockito'; import querystring from 'node:querystring'; import WebSocket from 'isomorphic-ws'; +import { mkdtemp, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { setImmediate as setImmediatePromise } from 'node:timers/promises'; +import tar from 'tar-fs'; -import { CallAwaiter } from './test/index.js'; import { KubeConfig } from './config.js'; import { Exec } from './exec.js'; import { Cp } from './cp.js'; @@ -14,14 +19,16 @@ describe('Cp', () => { it('should run create tar command to a url', async () => { const kc = new KubeConfig(); const fakeWebSocket: WebSocketInterface = mock(WebSocketHandler); + const fakeConn = mock(WebSocket); const exec = new Exec(kc, instance(fakeWebSocket)); const cp = new Cp(kc, exec); const namespace = 'somenamespace'; const pod = 'somepod'; const container = 'container'; - const srcPath = '/'; - const tgtPath = '/'; + const srcPath = await mkdtemp(join(tmpdir(), 'cp-src-')); + const tgtPath = await mkdtemp(join(tmpdir(), 'cp-tgt-')); + await writeFile(join(srcPath, 'test.txt'), 'test'); const cmdArray = ['tar', 'cf', '-', srcPath]; const path = `/api/v1/namespaces/${namespace}/pods/${pod}/exec`; @@ -35,21 +42,36 @@ describe('Cp', () => { }; const queryStr = querystring.stringify(query); - await cp.cpFromPod(namespace, pod, container, srcPath, tgtPath); - verify(fakeWebSocket.connect(`${path}?${queryStr}`, null, anyFunction())).called(); + when( + fakeWebSocket.connect(`${path}?${queryStr}`, null, anyFunction(), anyFunction()), + ).thenResolve(instance(fakeConn)); + const cpPromise = cp.cpFromPod(namespace, pod, container, srcPath, tgtPath); + await setImmediatePromise(); + const chunks = await tarChunks(srcPath); + verify(fakeWebSocket.connect(`${path}?${queryStr}`, null, anyFunction(), anyFunction())).called(); + const captured = requireCapture(fakeWebSocket); + for (const chunk of chunks) { + captured(WebSocketHandler.StdoutStream, chunk); + } + captured(WebSocketHandler.StatusStream, Buffer.from(JSON.stringify({ status: 'Success' }))); + await cpPromise; + await rm(srcPath, { recursive: true, force: true }); + await rm(tgtPath, { recursive: true, force: true }); }); it('should run create tar command to a url with cwd', async () => { const kc = new KubeConfig(); const fakeWebSocket: WebSocketInterface = mock(WebSocketHandler); + const fakeConn = mock(WebSocket); const exec = new Exec(kc, instance(fakeWebSocket)); const cp = new Cp(kc, exec); const namespace = 'somenamespace'; const pod = 'somepod'; const container = 'container'; - const srcPath = '/'; - const tgtPath = '/'; + const srcPath = await mkdtemp(join(tmpdir(), 'cp-src-')); + const tgtPath = await mkdtemp(join(tmpdir(), 'cp-tgt-')); + await writeFile(join(srcPath, 'test.txt'), 'test'); const cwd = '/abc'; const cmdArray = ['tar', 'cf', '-', '-C', cwd, srcPath]; const path = `/api/v1/namespaces/${namespace}/pods/${pod}/exec`; @@ -64,8 +86,21 @@ describe('Cp', () => { }; const queryStr = querystring.stringify(query); - await cp.cpFromPod(namespace, pod, container, srcPath, tgtPath, cwd); - verify(fakeWebSocket.connect(`${path}?${queryStr}`, null, anyFunction())).called(); + when( + fakeWebSocket.connect(`${path}?${queryStr}`, null, anyFunction(), anyFunction()), + ).thenResolve(instance(fakeConn)); + const cpPromise = cp.cpFromPod(namespace, pod, container, srcPath, tgtPath, cwd); + await setImmediatePromise(); + const chunks = await tarChunks(srcPath); + const captured = requireCapture(fakeWebSocket); + for (const chunk of chunks) { + captured(WebSocketHandler.StdoutStream, chunk); + } + captured(WebSocketHandler.StatusStream, Buffer.from(JSON.stringify({ status: 'Success' }))); + verify(fakeWebSocket.connect(`${path}?${queryStr}`, null, anyFunction(), anyFunction())).called(); + await cpPromise; + await rm(srcPath, { recursive: true, force: true }); + await rm(tgtPath, { recursive: true, force: true }); }); }); @@ -74,7 +109,6 @@ describe('Cp', () => { const kc = new KubeConfig(); const fakeWebSocketInterface: WebSocketInterface = mock(WebSocketHandler); const fakeWebSocket: WebSocket.WebSocket = mock(WebSocket) as WebSocket.WebSocket; - const callAwaiter: CallAwaiter = new CallAwaiter(); const exec = new Exec(kc, instance(fakeWebSocketInterface)); const cp = new Cp(kc, exec); @@ -97,14 +131,72 @@ describe('Cp', () => { const queryStr = querystring.stringify(query); const fakeConn: WebSocket.WebSocket = instance(fakeWebSocket); - when(fakeWebSocketInterface.connect(`${path}?${queryStr}`, null, anyFunction())).thenResolve( - fakeConn, + when( + fakeWebSocketInterface.connect(`${path}?${queryStr}`, null, anyFunction(), anyFunction()), + ).thenResolve(fakeConn); + when(fakeWebSocket.send(anything())).thenCall(() => {}); + when(fakeWebSocket.close()).thenCall(() => {}); + + const cpPromise = cp.cpToPod(namespace, pod, container, srcPath, tgtPath); + await setImmediatePromise(); + const captured = requireCapture(fakeWebSocketInterface); + captured(WebSocketHandler.StatusStream, Buffer.from(JSON.stringify({ status: 'Success' }))); + await cpPromise; + verify( + fakeWebSocketInterface.connect(`${path}?${queryStr}`, null, anyFunction(), anyFunction()), + ).called(); + }); + + it('should reject when the remote tar command fails', async () => { + const kc = new KubeConfig(); + const fakeWebSocketInterface: WebSocketInterface = mock(WebSocketHandler); + const fakeWebSocket: WebSocket.WebSocket = mock(WebSocket) as WebSocket.WebSocket; + const exec = new Exec(kc, instance(fakeWebSocketInterface)); + const cp = new Cp(kc, exec); + + const path = `/api/v1/namespaces/ns/pods/pod/exec`; + const query = { + stdout: false, + stderr: true, + stdin: true, + tty: false, + command: ['tar', 'xf', '-', '-C', '/'], + container: 'container', + }; + const queryStr = querystring.stringify(query); + + when( + fakeWebSocketInterface.connect(`${path}?${queryStr}`, null, anyFunction(), anyFunction()), + ).thenResolve(instance(fakeWebSocket)); + when(fakeWebSocket.send(anything())).thenCall(() => {}); + when(fakeWebSocket.close()).thenCall(() => {}); + + const cpPromise = cp.cpToPod('ns', 'pod', 'container', 'testdata/archive.txt', '/'); + await setImmediatePromise(); + const captured = requireCapture(fakeWebSocketInterface); + captured( + WebSocketHandler.StatusStream, + Buffer.from(JSON.stringify({ status: 'Failure', message: 'tar failed' })), ); - when(fakeWebSocket.send(anything())).thenCall(callAwaiter.resolveCall('send')); - when(fakeWebSocket.close()).thenCall(callAwaiter.resolveCall('close')); - await cp.cpToPod(namespace, pod, container, srcPath, tgtPath); - verify(fakeWebSocketInterface.connect(`${path}?${queryStr}`, null, anyFunction())).called(); + await rejects(cpPromise, /tar failed/); }); }); }); + +async function tarChunks(path: string): Promise { + const chunks: Buffer[] = []; + const pack = tar.pack(path); + for await (const chunk of pack) { + chunks.push(chunk as Buffer); + } + return chunks; +} + +function requireCapture(fakeWebSocket: WebSocketInterface): (streamNum: number, buff: Buffer) => boolean { + const [, , outputFn] = capture(fakeWebSocket.connect).last(); + if (!outputFn) { + throw new Error('expected output callback'); + } + return outputFn; +} diff --git a/src/exec.ts b/src/exec.ts index 090802b31fe..12ab5ff47a6 100644 --- a/src/exec.ts +++ b/src/exec.ts @@ -27,6 +27,8 @@ 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 {(err: any) => void} done - + * A callback called once when the command completes, is closed, or errors, optional. * @return {Promise} A promise that will return the web socket created for this command. */ public async exec( @@ -39,6 +41,7 @@ export class Exec { stdin: stream.Readable | null, tty: boolean, statusCallback?: (status: V1Status) => void, + done?: (err: any) => void, ): Promise { const query = { stdout: stdout != null, @@ -50,24 +53,21 @@ export class Exec { }; const queryStr = querystring.stringify(query); const path = `/api/v1/namespaces/${namespace}/pods/${podName}/exec?${queryStr}`; - const conn = await this.handler.connect(path, null, (streamNum: number, buff: Buffer): boolean => { - const status = WebSocketHandler.handleStandardStreams(streamNum, buff, stdout, stderr); - if (status != null) { - if (statusCallback) { - statusCallback(status); - } - return false; - } - return true; - }); - if (stdin != null) { - WebSocketHandler.handleStandardInput(conn, stdin, WebSocketHandler.StdinStream); - } + let resizeStream: stream.Readable | null = null; if (isResizable(stdout)) { this.terminalSizeQueue = new TerminalSizeQueue(); - WebSocketHandler.handleStandardInput(conn, this.terminalSizeQueue, WebSocketHandler.ResizeStream); + resizeStream = this.terminalSizeQueue; this.terminalSizeQueue.handleResizes(stdout as any as ResizableStream); } - return conn; + return WebSocketHandler.connectStandardStreams( + this.handler, + path, + stdout, + stderr, + stdin, + resizeStream, + statusCallback, + done, + ); } } diff --git a/src/exec_test.ts b/src/exec_test.ts index 405558b0c11..e4734499095 100644 --- a/src/exec_test.ts +++ b/src/exec_test.ts @@ -156,5 +156,43 @@ describe('Exec', () => { await closePromise; verify(fakeWebSocket.close()).called(); }); + + it('should call done with remote command failures', async () => { + const kc = new KubeConfig(); + const fakeWebSocketInterface: WebSocketInterface = mock(WebSocketHandler); + const fakeWebSocket: WebSocket.WebSocket = mock(WebSocket); + const exec = new Exec(kc, instance(fakeWebSocketInterface)); + const errStream = new WritableStreamBuffer(); + + const path = `/api/v1/namespaces/ns/pods/pod/exec`; + const args = `stdout=false&stderr=true&stdin=false&tty=false&command=cmd&container=container`; + when( + fakeWebSocketInterface.connect(`${path}?${args}`, null, anyFunction(), anyFunction()), + ).thenResolve(instance(fakeWebSocket)); + + let doneErr: any; + await exec.exec( + 'ns', + 'pod', + 'container', + 'cmd', + null, + errStream, + null, + false, + undefined, + (err) => { + doneErr = err; + }, + ); + + const [, , outputFn] = capture(fakeWebSocketInterface.connect).last(); + outputFn!( + WebSocketHandler.StatusStream, + Buffer.from(JSON.stringify({ status: 'Failure', message: 'command failed' })), + ); + + strictEqual(doneErr.message, 'command failed'); + }); }); }); diff --git a/src/log.ts b/src/log.ts index 2d77ec71609..7392884b4c4 100644 --- a/src/log.ts +++ b/src/log.ts @@ -4,7 +4,7 @@ import { ApiException } from './api.js'; import { KubeConfig } from './config.js'; import { HttpMethod, RequestContext } from './gen/http/http.js'; import { V1Status } from './gen/index.js'; -import { normalizeResponseHeaders } from './util.js'; +import { createDoneOnce, normalizeResponseHeaders } from './util.js'; export interface LogOptions { /** @@ -113,9 +113,11 @@ export class Log { doneOrOptions?: ((err: any) => void) | LogOptions, options?: LogOptions, ): Promise { + const done = typeof doneOrOptions === 'function' ? doneOrOptions : undefined; if (typeof doneOrOptions !== 'function') { options = doneOrOptions; } + const doneOnce = createDoneOnce(done); const path = `/api/v1/namespaces/${namespace}/pods/${podName}/log`; @@ -134,6 +136,15 @@ export class Log { await this.config.applySecurityAuthentication(ctx); const controller = new AbortController(); + controller.signal.addEventListener( + 'abort', + () => { + doneOnce( + controller.signal.reason ?? new DOMException('The operation was aborted', 'AbortError'), + ); + }, + { once: true }, + ); try { const response = await fetch(requestURL.toString(), { @@ -154,6 +165,11 @@ export class Log { ); } const nodeStream = Readable.fromWeb(response.body as any); + nodeStream.once('error', doneOnce); + stream.once('error', doneOnce); + stream.once('finish', () => doneOnce(null)); + nodeStream.once('end', () => doneOnce(null)); + nodeStream.once('close', () => doneOnce(null)); nodeStream.pipe(stream); } else if (status === 500) { const v1status = (await response.json()) as V1Status; @@ -184,10 +200,13 @@ export class Log { } } catch (err: any) { if (err instanceof ApiException) { + doneOnce(err); throw err; } - throw new ApiException(500, 'Error occurred in log request', undefined, {}); + const apiError = new ApiException(500, 'Error occurred in log request', undefined, {}); + doneOnce(apiError); + throw apiError; } return controller; diff --git a/src/log_test.ts b/src/log_test.ts index ce3f5641c1a..aced247e863 100644 --- a/src/log_test.ts +++ b/src/log_test.ts @@ -81,6 +81,30 @@ describe('Log', () => { mockAgent.assertNoPendingInterceptors(); }); + it('should call done when log stream completes', async () => { + const stream = new Writable({ + write(_chunk, _encoding, callback) { + callback(); + }, + }); + pool.intercept({ + method: 'GET', + path: '/api/v1/namespaces/default/pods/mypod/log?container=mycontainer', + }).reply(200, 'log data'); + + let doneErr: any = undefined; + const donePromise = new Promise((resolve) => { + log.log('default', 'mypod', 'mycontainer', stream, (err: any) => { + doneErr = err; + resolve(); + }); + }); + + await donePromise; + strictEqual(doneErr, null); + mockAgent.assertNoPendingInterceptors(); + }); + it('should throw an error if no active cluster', async () => { const configWithoutCluster = new KubeConfig(); const logWithoutCluster = new Log(configWithoutCluster); @@ -205,6 +229,12 @@ describe('Log', () => { strictEqual(searchParams.get('tailLines'), '1'); strictEqual(searchParams.get('timestamps'), 'true'); + searchParams = new URLSearchParams(); + options = { follow: true, pretty: false }; + AddOptionsToSearchParams(options, searchParams); + strictEqual(searchParams.get('follow'), 'true'); + strictEqual(searchParams.get('pretty'), 'false'); + const sinceTime = new Date().toISOString(); searchParams = new URLSearchParams(); options = { sinceTime }; diff --git a/src/portforward.ts b/src/portforward.ts index 07ddd872fc9..44d0e6127d5 100644 --- a/src/portforward.ts +++ b/src/portforward.ts @@ -5,6 +5,7 @@ import stream from 'node:stream'; import { AppsV1Api, CoreV1Api, V1Pod } from './gen/index.js'; import { KubeConfig } from './config.js'; import { WebSocketHandler, WebSocketInterface } from './web-socket-handler.js'; +import { createDoneOnce } from './util.js'; export class PortForward { private readonly config: KubeConfig; @@ -27,6 +28,7 @@ export class PortForward { err: stream.Writable | null, input: stream.Readable, retryCount: number = 0, + done?: (err: any) => void, ): Promise WebSocket.WebSocket | null)> { if (targetPorts.length === 0) { throw new Error('You must provide at least one port to forward to.'); @@ -43,11 +45,18 @@ export class PortForward { needsToReadPortNumber[index * 2] = true; needsToReadPortNumber[index * 2 + 1] = true; }); + const doneOnce = createDoneOnce(done); + output.once('error', doneOnce); + err?.once('error', doneOnce); const path = `/api/v1/namespaces/${namespace}/pods/${podName}/portforward?${queryStr}`; const createWebSocket = (): Promise => { - return this.handler.connect(path, null, (streamNum: number, buff: Buffer | string): boolean => { + const handleOutput = (streamNum: number, buff: Buffer | string): boolean => { if (streamNum >= targetPorts.length * 2) { - return !this.disconnectOnErr; + if (this.disconnectOnErr) { + doneOnce(new Error(`Unknown port-forward stream: ${streamNum}`)); + return false; + } + return true; } // First two bytes of each stream are the port number if (needsToReadPortNumber[streamNum]) { @@ -62,16 +71,26 @@ export class PortForward { output.write(buff); } return true; - }); + }; + return done + ? this.handler.connect(path, null, handleOutput, doneOnce) + : this.handler.connect(path, null, handleOutput); }; if (retryCount < 1) { const ws = await createWebSocket(); - WebSocketHandler.handleStandardInput(ws, input, 0); + WebSocketHandler.handleStandardInput(ws, input, 0, doneOnce); return ws; } - return WebSocketHandler.restartableHandleStandardInput(createWebSocket, input, 0, retryCount); + return WebSocketHandler.restartableHandleStandardInput( + createWebSocket, + input, + 0, + retryCount, + false, + doneOnce, + ); } /** @@ -94,6 +113,7 @@ export class PortForward { err: stream.Writable | null, input: stream.Readable, retryCount: number = 0, + done?: (err: any) => void, ): Promise WebSocket.WebSocket | null)> { const coreApi = this.config.makeApiClient(CoreV1Api); const service = await coreApi.readNamespacedService({ name: serviceName, namespace }); @@ -105,7 +125,16 @@ export class PortForward { const labelSelector = this.buildLabelSelector(service.spec.selector); const pod = await this.getFirstReadyPod(namespace, labelSelector); - return this.portForward(namespace, pod.metadata!.name!, targetPorts, output, err, input, retryCount); + return this.portForward( + namespace, + pod.metadata!.name!, + targetPorts, + output, + err, + input, + retryCount, + done, + ); } /** @@ -128,6 +157,7 @@ export class PortForward { err: stream.Writable | null, input: stream.Readable, retryCount: number = 0, + done?: (err: any) => void, ): Promise WebSocket.WebSocket | null)> { const appsApi = this.config.makeApiClient(AppsV1Api); const deployment = await appsApi.readNamespacedDeployment({ name: deploymentName, namespace }); @@ -142,7 +172,16 @@ export class PortForward { const labelSelector = this.buildLabelSelector(deployment.spec.selector.matchLabels); const pod = await this.getFirstReadyPod(namespace, labelSelector); - return this.portForward(namespace, pod.metadata!.name!, targetPorts, output, err, input, retryCount); + return this.portForward( + namespace, + pod.metadata!.name!, + targetPorts, + output, + err, + input, + retryCount, + done, + ); } /** diff --git a/src/portforward_test.ts b/src/portforward_test.ts index eca634062cf..407477f7da9 100644 --- a/src/portforward_test.ts +++ b/src/portforward_test.ts @@ -53,6 +53,28 @@ describe('PortForward', () => { strictEqual(osStream.size(), 1022); }); + it('should call done for unexpected error streams when disconnectOnErr is true', async () => { + const kc = new KubeConfig(); + const fakeWebSocket: WebSocketInterface = mock(WebSocketHandler); + const portForward = new PortForward(kc, true, instance(fakeWebSocket)); + const osStream = new WritableStreamBuffer(); + const isStream = new ReadableStreamBuffer(); + + let doneErr: any; + await portForward.portForward('ns', 'p', [8000], osStream, null, isStream, 0, (err) => { + doneErr = err; + }); + + const [, , outputFn] = capture(fakeWebSocket.connect).last(); + + strictEqual(typeof outputFn, 'function'); + if (!outputFn) { + return; + } + strictEqual(outputFn(2, Buffer.alloc(1024, 10)), false); + strictEqual(doneErr.message, 'Unknown port-forward stream: 2'); + }); + it('should correctly port-forward streams if err is null', async () => { const kc = new KubeConfig(); const fakeWebSocket: WebSocketInterface = mock(WebSocketHandler); diff --git a/src/util.ts b/src/util.ts index a27b04e2a86..8475082092e 100644 --- a/src/util.ts +++ b/src/util.ts @@ -164,6 +164,16 @@ export function normalizeResponseHeaders(response: { headers: { entries(): Itera return normalizedHeaders; } +export function createDoneOnce(done?: (err: any) => void): (err: any) => void { + let doneCalled = false; + return (err: any) => { + if (!doneCalled) { + doneCalled = true; + done?.(err); + } + }; +} + /** * Built-in Kubernetes API groups that have generated TypeScript models. * Custom resources and third-party API groups (like Knative) are not included. diff --git a/src/watch_test.ts b/src/watch_test.ts index 0b953b1787b..dba7accbc9c 100644 --- a/src/watch_test.ts +++ b/src/watch_test.ts @@ -348,7 +348,6 @@ describe('Watch', () => { ); await donePromise; - deepStrictEqual(receivedTypes, [obj.type]); deepStrictEqual(receivedObjects, [obj.object]); }); diff --git a/src/web-socket-handler.ts b/src/web-socket-handler.ts index d0abf6bab1e..d7136784659 100644 --- a/src/web-socket-handler.ts +++ b/src/web-socket-handler.ts @@ -3,6 +3,7 @@ import stream from 'node:stream'; import { V1Status } from './api.js'; import { KubeConfig } from './config.js'; +import { createDoneOnce } from './util.js'; const protocols = [ 'v5.channel.k8s.io', @@ -17,6 +18,7 @@ export interface WebSocketInterface { path: string, textHandler: ((text: string) => boolean) | null, binaryHandler: ((stream: number, buff: Buffer) => boolean) | null, + done?: (err: any) => void, ): Promise; } @@ -85,6 +87,7 @@ export class WebSocketHandler implements WebSocketInterface { ws: WebSocket.WebSocket, stdin: stream.Readable, streamNum: number = 0, + done?: (err: any) => void, ): boolean { stdin.on('data', (data) => { ws.send(copyChunkForWebSocket(streamNum, data, stdin.readableEncoding)); @@ -100,10 +103,60 @@ export class WebSocketHandler implements WebSocketInterface { } ws.close(); }); + stdin.on('error', (err) => { + done?.(err); + ws.close(); + }); // Keep the stream open return true; } + public static statusError(status: V1Status): Error | null { + if (status.status === 'Failure' || status.reason === 'NonZeroExitCode') { + return new Error(status.message || status.reason || 'Remote command failed'); + } + return null; + } + + public static async connectStandardStreams( + handler: WebSocketInterface, + path: string, + stdout: stream.Writable | null, + stderr: stream.Writable | null, + stdin: stream.Readable | null, + resizeStream: stream.Readable | null, + statusCallback?: (status: V1Status) => void, + done?: (err: any) => void, + ): Promise { + const doneOnce = createDoneOnce(done); + stdout?.once('error', doneOnce); + stderr?.once('error', doneOnce); + + const handleOutput = (streamNum: number, buff: Buffer): boolean => { + const status = WebSocketHandler.handleStandardStreams(streamNum, buff, stdout, stderr); + if (status != null) { + if (statusCallback) { + statusCallback(status); + } + doneOnce(WebSocketHandler.statusError(status)); + return false; + } + return true; + }; + + const conn = done + ? await handler.connect(path, null, handleOutput, doneOnce) + : await handler.connect(path, null, handleOutput); + + if (stdin != null) { + WebSocketHandler.handleStandardInput(conn, stdin, WebSocketHandler.StdinStream, doneOnce); + } + if (resizeStream != null) { + WebSocketHandler.handleStandardInput(conn, resizeStream, WebSocketHandler.ResizeStream, doneOnce); + } + return conn; + } + public static async processData( data: string | Buffer, ws: WebSocket.WebSocket | null, @@ -139,6 +192,7 @@ export class WebSocketHandler implements WebSocketInterface { retryCount: number = 3, // kind of hacky, but otherwise we can't wait for the writes to flush before testing. addFlushForTesting: boolean = false, + done?: (err: any) => void, ): () => WebSocket.WebSocket | null { if (retryCount < 0) { throw new Error("retryCount can't be lower than 0."); @@ -147,16 +201,23 @@ export class WebSocketHandler implements WebSocketInterface { let ws: WebSocket.WebSocket | null = null; stdin.on('data', (data) => { - queue = queue.then(async () => { - ws = await WebSocketHandler.processData( - data, - ws, - createWS, - streamNum, - retryCount, - stdin.readableEncoding, - ); - }); + queue = queue + .then(async () => { + ws = await WebSocketHandler.processData( + data, + ws, + createWS, + streamNum, + retryCount, + stdin.readableEncoding, + ); + }) + .catch((err) => { + done?.(err); + if (ws !== null) { + ws.close(); + } + }); }); if (addFlushForTesting) { @@ -170,6 +231,12 @@ export class WebSocketHandler implements WebSocketInterface { ws.close(); } }); + stdin.on('error', (err) => { + done?.(err); + if (ws !== null) { + ws.close(); + } + }); return () => ws; } @@ -213,6 +280,7 @@ export class WebSocketHandler implements WebSocketInterface { path: string, textHandler: ((text: string) => boolean) | null, binaryHandler: ((stream: number, buff: Buffer) => boolean) | null, + done?: (err: any) => void, ): Promise { const cluster = this.config.getCurrentCluster(); if (!cluster) { @@ -233,6 +301,7 @@ export class WebSocketHandler implements WebSocketInterface { ? this.socketFactory(uri, protocols, opts) : new WebSocket(uri, protocols, opts); let resolved = false; + const doneOnce = createDoneOnce(done); client.onopen = () => { resolved = true; @@ -242,25 +311,47 @@ export class WebSocketHandler implements WebSocketInterface { client.onerror = (err) => { if (!resolved) { reject(err); + } else { + doneOnce(err); } }; + client.onclose = () => { + doneOnce(null); + }; + client.onmessage = ({ data }: { data: WebSocket.Data }) => { - // TODO: support ArrayBuffer and Buffer[] data types? - if (typeof data === 'string') { - if (data.charCodeAt(0) === WebSocketHandler.CloseStream) { - WebSocketHandler.closeStream(data.charCodeAt(1), this.streams); - } - if (textHandler && !textHandler(data)) { - client.close(); - } - } else if (data instanceof Buffer) { - const streamNum = data.readUint8(0); - if (streamNum === WebSocketHandler.CloseStream) { - WebSocketHandler.closeStream(data.readInt8(1), this.streams); + try { + // TODO: support ArrayBuffer and Buffer[] data types? + if (typeof data === 'string') { + if (data.charCodeAt(0) === WebSocketHandler.CloseStream) { + WebSocketHandler.closeStream(data.charCodeAt(1), this.streams); + return; + } + if (textHandler && !textHandler(data)) { + client.close(); + } + } else if (data instanceof Buffer) { + if (data.length < 1) { + return; + } + const streamNum = data.readUint8(0); + if (streamNum === WebSocketHandler.CloseStream) { + if (data.length > 1) { + WebSocketHandler.closeStream(data.readInt8(1), this.streams); + } + return; + } + if (binaryHandler && !binaryHandler(streamNum, data.slice(1))) { + client.close(); + } } - if (binaryHandler && !binaryHandler(streamNum, data.slice(1))) { + } catch (err) { + doneOnce(err); + try { client.close(); + } catch { + // Ignore close errors while handling an existing stream error. } } }; diff --git a/src/web-socket-handler_test.ts b/src/web-socket-handler_test.ts index 22f4ce1917f..deb3a8ec646 100644 --- a/src/web-socket-handler_test.ts +++ b/src/web-socket-handler_test.ts @@ -291,6 +291,133 @@ describe('WebSocket', () => { strictEqual(datum, fill); } }); + it('should call done once after connection errors and closes', async () => { + const kc = new KubeConfig(); + const host = 'foo.company.com'; + const server = `https://${host}`; + kc.clusters = [ + { + name: 'cluster', + server, + } as Cluster, + ] as Cluster[]; + kc.contexts = [ + { + cluster: 'cluster', + user: 'user', + } as Context, + ] as Context[]; + kc.users = [ + { + name: 'user', + } as User, + ]; + + const mockWs = {} as WebSocket.WebSocket; + const handler = new WebSocketHandler(kc, (): WebSocket.WebSocket => { + return mockWs; + }); + + let doneCount = 0; + let doneErr: any; + const promise = handler.connect('/some/path', null, null, (err: any) => { + doneCount += 1; + doneErr = err; + }); + await setImmediatePromise(); + mockWs.onopen!({ target: mockWs, type: 'open' }); + await promise; + + const errEvt = { + error: {}, + message: 'some message', + type: 'some type', + target: mockWs, + }; + mockWs.onerror!(errEvt); + mockWs.onclose!({ target: mockWs, type: 'close', wasClean: false, code: 1006, reason: '' }); + + strictEqual(doneCount, 1); + strictEqual(doneErr, errEvt); + }); + it('should connect standard streams through the shared utility', async () => { + const stdout = new WritableStreamBuffer(); + const stderr = new WritableStreamBuffer(); + const stdin = new ReadableStreamBuffer(); + const resize = new ReadableStreamBuffer(); + const sent: Buffer[] = []; + const ws = { + protocol: 'v5.channel.k8s.io', + send: (data) => { + sent.push(data as Buffer); + }, + close: () => {}, + } as WebSocket.WebSocket; + + let pathOut = ''; + let binaryHandler: ((stream: number, buff: Buffer) => boolean) | null = null; + let doneHandler: ((err: any) => void) | undefined; + const handler = { + connect: async ( + path: string, + textHandler: ((text: string) => boolean) | null, + binary: ((stream: number, buff: Buffer) => boolean) | null, + done?: (err: any) => void, + ): Promise => { + strictEqual(textHandler, null); + pathOut = path; + binaryHandler = binary; + doneHandler = done; + return ws; + }, + }; + + let statusOut: V1Status | undefined; + let doneCount = 0; + let doneErr: any = undefined; + const conn = await WebSocketHandler.connectStandardStreams( + handler, + '/exec', + stdout, + stderr, + stdin, + resize, + (status) => { + statusOut = status; + }, + (err) => { + doneCount++; + doneErr = err; + }, + ); + + strictEqual(conn, ws); + strictEqual(pathOut, '/exec'); + strictEqual(typeof binaryHandler, 'function'); + strictEqual(typeof doneHandler, 'function'); + + strictEqual(binaryHandler!(WebSocketHandler.StdoutStream, Buffer.from('out')), true); + strictEqual(binaryHandler!(WebSocketHandler.StderrStream, Buffer.from('err')), true); + strictEqual(stdout.getContentsAsString(), 'out'); + strictEqual(stderr.getContentsAsString(), 'err'); + + stdin.emit('data', 'input'); + resize.emit('data', 'resize'); + deepStrictEqual(sent[0], Buffer.from('\x00input')); + deepStrictEqual(sent[1], Buffer.from('\x04resize')); + + const status = { status: 'Success', message: 'ok' } as V1Status; + strictEqual( + binaryHandler!(WebSocketHandler.StatusStream, Buffer.from(JSON.stringify(status))), + false, + ); + deepStrictEqual(statusOut, status); + strictEqual(doneCount, 1); + strictEqual(doneErr, null); + + doneHandler!(new Error('late error')); + strictEqual(doneCount, 1); + }); it('handles multi-byte characters', () => { return new Promise((resolve) => { const stream = new Readable({ read() {} });