Merge pull request #240 from lizardruss/master

k8s-ci-robot · web-flow · commit e5ec9456cf8c · 2019-05-02T12:24:27.000-07:00
Support terminal resizes
diff --git a/package.json b/package.json
@@ -67,6 +67,7 @@
     "@types/chai": "^4.1.6",
     "@types/mocha": "^5.2.5",
     "@types/mock-fs": "^3.6.30",
+    "@types/stream-buffers": "^3.0.3",
     "chai": "^4.2.0",
     "jasmine": "^3.3.0",
     "mocha": "^5.2.0",
diff --git a/src/attach.ts b/src/attach.ts
@@ -3,11 +3,14 @@ import querystring = require('querystring');
 import stream = require('stream');
 
 import { KubeConfig } from './config';
+import { isResizable, ResizableStream, TerminalSizeQueue } from './terminal-size-queue';
 import { WebSocketHandler, WebSocketInterface } from './web-socket-handler';
 
 export class Attach {
     public 'handler': WebSocketInterface;
 
+    private terminalSizeQueue?: TerminalSizeQueue;
+
     public constructor(config: KubeConfig, websocketInterface?: WebSocketInterface) {
         if (websocketInterface) {
             this.handler = websocketInterface;
@@ -43,9 +46,13 @@ export class Attach {
             },
         );
         if (stdin != null) {
-            WebSocketHandler.handleStandardInput(conn, stdin);
+            WebSocketHandler.handleStandardInput(conn, stdin, WebSocketHandler.StdinStream);
+        }
+        if (isResizable(stdout)) {
+            this.terminalSizeQueue = new TerminalSizeQueue();
+            WebSocketHandler.handleStandardInput(conn, this.terminalSizeQueue, WebSocketHandler.ResizeStream);
+            this.terminalSizeQueue.handleResizes((stdout as any) as ResizableStream);
         }
-
         return conn;
     }
 }
diff --git a/src/attach_test.ts b/src/attach_test.ts
@@ -3,8 +3,10 @@ import WebSocket = require('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';
 import { Attach } from './attach';
 import { KubeConfig } from './config';
+import { TerminalSize } from './terminal-size-queue';
 import { WebSocketHandler, WebSocketInterface } from './web-socket-handler';
 
 describe('Attach', () => {
@@ -46,9 +48,11 @@ describe('Attach', () => {
 
         it('should correctly attach to streams', async () => {
             const kc = new KubeConfig();
-            const fakeWebSocket: WebSocketInterface = mock(WebSocketHandler);
-            const attach = new Attach(kc, instance(fakeWebSocket));
-            const osStream = new WritableStreamBuffer();
+            const fakeWebSocketInterface: WebSocketInterface = mock(WebSocketHandler);
+            const fakeWebSocket: WebSocket = mock(WebSocket);
+            const callAwaiter: CallAwaiter = new CallAwaiter();
+            const attach = new Attach(kc, instance(fakeWebSocketInterface));
+            const osStream = new ResizableWriteableStreamBuffer();
             const errStream = new WritableStreamBuffer();
             const isStream = new ReadableStreamBuffer();
 
@@ -59,11 +63,15 @@ describe('Attach', () => {
             const path = `/api/v1/namespaces/${namespace}/pods/${pod}/attach`;
             const args = `container=${container}&stderr=true&stdin=true&stdout=true&tty=false`;
 
-            const fakeConn: WebSocket = mock(WebSocket);
-            when(fakeWebSocket.connect(`${path}?${args}`, null, anyFunction())).thenResolve(fakeConn);
+            const fakeConn: WebSocket = instance(fakeWebSocket);
+            when(fakeWebSocketInterface.connect(`${path}?${args}`, null, anyFunction())).thenResolve(
+                fakeConn,
+            );
+            when(fakeWebSocket.send(anything())).thenCall(callAwaiter.resolveCall('send'));
+            when(fakeWebSocket.close()).thenCall(callAwaiter.resolveCall('close'));
 
             await attach.attach(namespace, pod, container, osStream, errStream, isStream, false);
-            const [, , outputFn] = capture(fakeWebSocket.connect).last();
+            const [, , outputFn] = capture(fakeWebSocketInterface.connect).last();
 
             /* tslint:disable:no-unused-expression */
             expect(outputFn).to.not.be.null;
@@ -90,12 +98,34 @@ describe('Attach', () => {
                 expect(buff[i]).to.equal(20);
             }
 
+            const initialTerminalSize: TerminalSize = { height: 0, width: 0 };
+            await callAwaiter.awaitCall('send');
+            verify(
+                fakeWebSocket.send(
+                    matchBuffer(WebSocketHandler.ResizeStream, JSON.stringify(initialTerminalSize)),
+                ),
+            ).called();
+
             const msg = 'This is test data';
+            const inputPromise = callAwaiter.awaitCall('send');
             isStream.put(msg);
-            verify(fakeConn.send(msg));
-
+            await inputPromise;
+            verify(fakeWebSocket.send(matchBuffer(WebSocketHandler.StdinStream, msg))).called();
+
+            const terminalSize: TerminalSize = { height: 80, width: 120 };
+            const resizePromise = callAwaiter.awaitCall('send');
+            osStream.rows = terminalSize.height;
+            osStream.columns = terminalSize.width;
+            osStream.emit('resize');
+            await resizePromise;
+            verify(
+                fakeWebSocket.send(matchBuffer(WebSocketHandler.ResizeStream, JSON.stringify(terminalSize))),
+            ).called();
+
+            const closePromise = callAwaiter.awaitCall('close');
             isStream.stop();
-            verify(fakeConn.close());
+            await closePromise;
+            verify(fakeWebSocket.close()).called();
         });
     });
 });
diff --git a/src/exec.ts b/src/exec.ts
@@ -4,11 +4,14 @@ import stream = require('stream');
 
 import { V1Status } from './api';
 import { KubeConfig } from './config';
+import { isResizable, ResizableStream, TerminalSizeQueue } from './terminal-size-queue';
 import { WebSocketHandler, WebSocketInterface } from './web-socket-handler';
 
 export class Exec {
     public 'handler': WebSocketInterface;
 
+    private terminalSizeQueue?: TerminalSizeQueue;
+
     public constructor(config: KubeConfig, wsInterface?: WebSocketInterface) {
         if (wsInterface) {
             this.handler = wsInterface;
@@ -66,7 +69,12 @@ export class Exec {
             },
         );
         if (stdin != null) {
-            WebSocketHandler.handleStandardInput(conn, stdin);
+            WebSocketHandler.handleStandardInput(conn, stdin, WebSocketHandler.StdinStream);
+        }
+        if (isResizable(stdout)) {
+            this.terminalSizeQueue = new TerminalSizeQueue();
+            WebSocketHandler.handleStandardInput(conn, this.terminalSizeQueue, WebSocketHandler.ResizeStream);
+            this.terminalSizeQueue.handleResizes((stdout as any) as ResizableStream);
         }
         return conn;
     }
diff --git a/src/exec_test.ts b/src/exec_test.ts
@@ -1,11 +1,13 @@
 import { expect } from 'chai';
 import WebSocket = require('isomorphic-ws');
 import { ReadableStreamBuffer, WritableStreamBuffer } from 'stream-buffers';
-import { anyFunction, capture, instance, mock, verify, when } from 'ts-mockito';
+import { anyFunction, anything, capture, instance, mock, verify, when } from 'ts-mockito';
 
+import { CallAwaiter, matchBuffer, ResizableWriteableStreamBuffer } from '../test';
 import { V1Status } from './api';
 import { KubeConfig } from './config';
 import { Exec } from './exec';
+import { TerminalSize } from './terminal-size-queue';
 import { WebSocketHandler, WebSocketInterface } from './web-socket-handler';
 
 describe('Exec', () => {
@@ -55,9 +57,11 @@ describe('Exec', () => {
 
         it('should correctly attach to streams', async () => {
             const kc = new KubeConfig();
-            const fakeWebSocket: WebSocketInterface = mock(WebSocketHandler);
-            const exec = new Exec(kc, instance(fakeWebSocket));
-            const osStream = new WritableStreamBuffer();
+            const fakeWebSocketInterface: WebSocketInterface = mock(WebSocketHandler);
+            const fakeWebSocket: WebSocket = mock(WebSocket);
+            const callAwaiter: CallAwaiter = new CallAwaiter();
+            const exec = new Exec(kc, instance(fakeWebSocketInterface));
+            const osStream = new ResizableWriteableStreamBuffer();
             const errStream = new WritableStreamBuffer();
             const isStream = new ReadableStreamBuffer();
 
@@ -71,8 +75,12 @@ describe('Exec', () => {
 
             let statusOut = {} as V1Status;
 
-            const fakeConn: WebSocket = mock(WebSocket);
-            when(fakeWebSocket.connect(`${path}?${args}`, null, anyFunction())).thenResolve(fakeConn);
+            const fakeConn: WebSocket = instance(fakeWebSocket);
+            when(fakeWebSocketInterface.connect(`${path}?${args}`, null, anyFunction())).thenResolve(
+                fakeConn,
+            );
+            when(fakeWebSocket.send(anything())).thenCall(callAwaiter.resolveCall('send'));
+            when(fakeWebSocket.close()).thenCall(callAwaiter.resolveCall('close'));
 
             await exec.exec(
                 namespace,
@@ -88,7 +96,7 @@ describe('Exec', () => {
                 },
             );
 
-            const [, , outputFn] = capture(fakeWebSocket.connect).last();
+            const [, , outputFn] = capture(fakeWebSocketInterface.connect).last();
 
             /* tslint:disable:no-unused-expression */
             expect(outputFn).to.not.be.null;
@@ -115,20 +123,41 @@ describe('Exec', () => {
                 expect(buff[i]).to.equal(20);
             }
 
+            const initialTerminalSize: TerminalSize = { height: 0, width: 0 };
+            await callAwaiter.awaitCall('send');
+            verify(
+                fakeWebSocket.send(
+                    matchBuffer(WebSocketHandler.ResizeStream, JSON.stringify(initialTerminalSize)),
+                ),
+            ).called();
+
             const msg = 'This is test data';
+            const inputPromise = callAwaiter.awaitCall('send');
             isStream.put(msg);
-            verify(fakeConn.send(msg));
+            await inputPromise;
+            verify(fakeWebSocket.send(matchBuffer(WebSocketHandler.StdinStream, msg))).called();
+
+            const terminalSize: TerminalSize = { height: 80, width: 120 };
+            const resizePromise = callAwaiter.awaitCall('send');
+            osStream.rows = terminalSize.height;
+            osStream.columns = terminalSize.width;
+            osStream.emit('resize');
+            await resizePromise;
+            verify(
+                fakeWebSocket.send(matchBuffer(WebSocketHandler.ResizeStream, JSON.stringify(terminalSize))),
+            ).called();
 
             const statusIn = {
                 code: 100,
                 message: 'this is a test',
             } as V1Status;
-
             outputFn(WebSocketHandler.StatusStream, Buffer.from(JSON.stringify(statusIn)));
             expect(statusOut).to.deep.equal(statusIn);
 
+            const closePromise = callAwaiter.awaitCall('close');
             isStream.stop();
-            verify(fakeConn.close());
+            await closePromise;
+            verify(fakeWebSocket.close()).called();
         });
     });
 });
diff --git a/src/terminal-size-queue.ts b/src/terminal-size-queue.ts
@@ -0,0 +1,49 @@
+import { Readable, ReadableOptions } from 'stream';
+
+export interface ResizableStream {
+    columns: number;
+    rows: number;
+    on(event: 'resize', cb: () => void);
+}
+
+export interface TerminalSize {
+    height: number;
+    width: number;
+}
+
+export class TerminalSizeQueue extends Readable {
+    constructor(opts: ReadableOptions = {}) {
+        super({
+            ...opts,
+            // tslint:disable-next-line:no-empty
+            read() {},
+        });
+    }
+
+    public handleResizes(writeStream: ResizableStream) {
+        // Set initial size
+        this.resize(getTerminalSize(writeStream));
+
+        // Handle future size updates
+        writeStream.on('resize', () => this.resize(getTerminalSize(writeStream)));
+    }
+
+    private resize(size: TerminalSize) {
+        this.push(JSON.stringify(size));
+    }
+}
+
+export function isResizable(stream: any) {
+    if (stream == null) {
+        return false;
+    }
+
+    const hasRows = 'rows' in stream;
+    const hasColumns = 'columns' in stream;
+    const hasOn = typeof stream.on === 'function';
+    return hasRows && hasColumns && hasOn;
+}
+
+function getTerminalSize(writeStream: ResizableStream): TerminalSize {
+    return { height: writeStream.rows!, width: writeStream.columns! };
+}
diff --git a/src/web-socket-handler.ts b/src/web-socket-handler.ts
@@ -19,6 +19,7 @@ export class WebSocketHandler implements WebSocketInterface {
     public static readonly StdoutStream = 1;
     public static readonly StderrStream = 2;
     public static readonly StatusStream = 3;
+    public static readonly ResizeStream = 4;
 
     public static handleStandardStreams(
         streamNum: number,
diff --git a/test/call-awaiter.ts b/test/call-awaiter.ts
@@ -0,0 +1,13 @@
+import { EventEmitter } from 'events';
+
+export class CallAwaiter extends EventEmitter {
+  public awaitCall(event: string) {
+    return new Promise<any[]>((resolve) => {
+      this.once(event, resolve);
+    });
+  }
+
+  public resolveCall(event: string) {
+    return (...args: any[]) => this.emit(event, ...args);
+  }
+}
diff --git a/test/index.ts b/test/index.ts
@@ -0,0 +1,3 @@
+export * from './call-awaiter';
+export * from './match-buffer';
+export * from './resizable-writeable-stream-buffer';
diff --git a/test/match-buffer.ts b/test/match-buffer.ts
@@ -0,0 +1,26 @@
+import { Matcher } from 'ts-mockito/lib/matcher/type/Matcher';
+
+export function matchBuffer(channel: number, contents: string): Matcher {
+  return new StringBufferMatcher(channel, contents);
+}
+
+class StringBufferMatcher extends Matcher {
+  constructor(private channel: number, private contents: string) {
+      super();
+  }
+
+  public match(value: any): boolean {
+      if (value instanceof Buffer) {
+          const buffer = value as Buffer;
+          const channel: number = buffer.readInt8(0);
+          const contents: string = buffer.toString('utf8', 1);
+          return this.channel === channel && this.contents === contents;
+      }
+
+      return false;
+  }
+
+  public toString(): string {
+      return `buffer did not contain "${this.contents}"`;
+  }
+}
diff --git a/test/resizable-writeable-stream-buffer.ts b/test/resizable-writeable-stream-buffer.ts
@@ -0,0 +1,6 @@
+import { WritableStreamBuffer } from 'stream-buffers';
+
+export class ResizableWriteableStreamBuffer extends WritableStreamBuffer implements NodeJS.WritableStream {
+  public columns: number = 0;
+  public rows: number = 0;
+}

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+export * from './call-awaiter';`
	`2`	`+export * from './match-buffer';`
	`3`	`+export * from './resizable-writeable-stream-buffer';`