Merge pull request #204 from proAlexandr/issue-187

Fix #187 (forward.portForward - Uncaught Error: not opened)

Author: k8s-ci-robot

src/portforward.ts

@@ -28,7 +28,8 @@ export class PortForward {
         output: stream.Writable,
         err: stream.Writable | null,
         input: stream.Readable,
-    ): Promise<WebSocket> {
+        retryCount: number = 0,
+    ): Promise<WebSocket | (() => WebSocket | null)> {
         if (targetPorts.length === 0) {
             throw new Error('You must provide at least one port to forward to.');
         }
@@ -45,29 +46,37 @@ export class PortForward {
             needsToReadPortNumber[index * 2 + 1] = true;
         });
         const path = `/api/v1/namespaces/${namespace}/pods/${podName}/portforward?${queryStr}`;
-        const conn = await this.handler.connect(
-            path,
-            null,
-            (streamNum: number, buff: Buffer | string): boolean => {
-                if (streamNum >= targetPorts.length * 2) {
-                    return !this.disconnectOnErr;
-                }
-                // First two bytes of each stream are the port number
-                if (needsToReadPortNumber[streamNum]) {
-                    buff = buff.slice(2);
-                    needsToReadPortNumber[streamNum] = false;
-                }
-                if (streamNum % 2 === 1) {
-                    if (err) {
-                        err.write(buff);
+        const createWebSocket = (): Promise<WebSocket> => {
+            return this.handler.connect(
+                path,
+                null,
+                (streamNum: number, buff: Buffer | string): boolean => {
+                    if (streamNum >= targetPorts.length * 2) {
+                        return !this.disconnectOnErr;
                     }
-                } else {
-                    output.write(buff);
-                }
-                return true;
-            },
-        );
-        WebSocketHandler.handleStandardInput(conn, input, 0);
-        return conn;
+                    // First two bytes of each stream are the port number
+                    if (needsToReadPortNumber[streamNum]) {
+                        buff = buff.slice(2);
+                        needsToReadPortNumber[streamNum] = false;
+                    }
+                    if (streamNum % 2 === 1) {
+                        if (err) {
+                            err.write(buff);
+                        }
+                    } else {
+                        output.write(buff);
+                    }
+                    return true;
+                },
+            );
+        };
+
+        if (retryCount < 1) {
+            const ws = await createWebSocket();
+            WebSocketHandler.handleStandardInput(ws, input, 0);
+            return ws;
+        }
+
+        return WebSocketHandler.restartableHandleStandardInput(createWebSocket, input, 0, retryCount);
     }
 }

src/web-socket-handler.ts

@@ -71,6 +71,57 @@ export class WebSocketHandler implements WebSocketInterface {
         return true;
     }
 
+    public static restartableHandleStandardInput(
+        createWS: () => Promise<WebSocket>,
+        stdin: stream.Readable | any,
+        streamNum: number = 0,
+        retryCount: number = 3,
+    ): () => WebSocket | null {
+        if (retryCount < 0) {
+            throw new Error("retryCount can't be lower than 0.");
+        }
+
+        let queue: Promise<void> = Promise.resolve();
+        let ws: WebSocket | null;
+
+        async function processData(data): Promise<void> {
+            const buff = Buffer.alloc(data.length + 1);
+
+            buff.writeInt8(streamNum, 0);
+            if (data instanceof Buffer) {
+                data.copy(buff, 1);
+            } else {
+                buff.write(data, 1);
+            }
+
+            let i = 0;
+            for (; i < retryCount; ++i) {
+                if (ws !== null && ws.readyState === WebSocket.OPEN) {
+                    ws.send(buff);
+                    break;
+                } else {
+                    ws = await createWS();
+                }
+            }
+
+            if (i >= retryCount) {
+                throw new Error("can't send data to ws");
+            }
+        }
+
+        stdin.on('data', (data) => {
+            queue = queue.then(() => processData(data));
+        });
+
+        stdin.on('end', () => {
+            if (ws) {
+                ws.close();
+            }
+        });
+
+        return () => ws;
+    }
+
     // factory is really just for test injection
     public constructor(
         readonly config: KubeConfig,