Normalize streaming lifecycle behavior (#16)

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<WebSocket.WebSocket> {
         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,
+        );
     }
 }

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<void>((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<void>((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)]);
     }
 }

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<Buffer[]> {
+    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;
+}

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<WebSocket>} 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<WebSocket.WebSocket> {
         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,
+        );
     }
 }