ProxyPty.ts

/*
 * Copyright 2018 Simon Edwards <simon@simonzone.com>
 *
 * This source code is licensed under the MIT license which is detailed in the LICENSE.txt file.
 */
import * as child_process from 'child_process';
import {Event, BufferSizeChange, Pty, Logger} from 'extraterm-extension-api';
import { EventEmitter } from 'extraterm-event-emitter';

export interface EnvironmentMap {
  [key:string]: string;
}

export interface PtyOptions {
  exe?: string;
  args?: string[];
  cols?: number;
  rows?: number;
  cwd?: string;
  env?: EnvironmentMap;
  extraEnv?: EnvironmentMap;
}


const DEBUG_FINE = false;

const MAXIMUM_WRITE_BUFFER_SIZE = 64 * 1024;


const TYPE_CREATE = "create";
const TYPE_CREATED = "created";
const TYPE_WRITE = "write";
const TYPE_OUTPUT = "output";
const TYPE_RESIZE = "resize";
const TYPE_CLOSE = "close";
const TYPE_CLOSED = "closed";
const TYPE_TERMINATE = "terminate";
const TYPE_PERMIT_DATA_SIZE = "permit-data-size";
const TYPE_OUTPUT_WRITTEN = "output-written";

interface ProxyMessage {
  type: string;
  id: number;
}

interface CreatePtyMessage extends ProxyMessage {
  argv: string[];
  rows: number;
  columns: number;
  env: { [key: string]: string; };
  extraEnv: { [key: string]: string; };
  cwd: string;

  // the id field is not user for this message type.
}

interface CreatedPtyMessage extends ProxyMessage {
}

interface WriteMessage extends ProxyMessage {
  data: string;
}

interface ResizeMessage extends ProxyMessage {
  rows: number;
  columns: number;
}

interface CloseMessage extends ProxyMessage {
}

interface PermitDataSizeMessage extends ProxyMessage {
  size: number;
}

// Output generated by the process on the other side of the pty.
interface OutputMessage extends ProxyMessage {
  data: string;
}

interface OutputWrittenMessage extends ProxyMessage {
  chars: number;
}

interface ClosedMessage extends ProxyMessage {
}

interface TerminateMessage extends ProxyMessage {  
}

const NULL_ID = -1;


class ProxyPty implements Pty {
  
  private _id: number = NULL_ID;
  private _writeFunc: (id: number, msg: ProxyMessage) => void = null;
  
  // Pre-open write queue.
  private _writeQueue: ProxyMessage[] = [];
  private _live = true;
  private _outstandingWriteDataCount = 0;
  
  private _onDataEventEmitter = new EventEmitter<string>();
  private _onExitEventEmitter = new EventEmitter<void>();
  private _onAvailableWriteBufferSizeChangeEventEmitter = new EventEmitter<BufferSizeChange>();

  onData: Event<string>;
  onExit: Event<void>;
  onAvailableWriteBufferSizeChange: Event<BufferSizeChange>;

  constructor(writeFunc) {
    this._writeFunc = writeFunc;

    this.onData = this._onDataEventEmitter.event;
    this.onExit = this._onExitEventEmitter.event;
    this.onAvailableWriteBufferSizeChange = this._onAvailableWriteBufferSizeChangeEventEmitter.event;
  }
  
  getId(): number {
    return this._id;
  }
  
  setId(id: number): void {
    this._id = id;
    if (this._live) {
      this._writeQueue.forEach( (msg) => {
        msg.id = this._id;
        this._writeFunc(this._id, msg);
      });
      this._writeQueue = [];
    }
  }
  
  private _writeMessage(id: number, msg: ProxyMessage): void {
    if (this._live) {
      if (this._id === -1) {
        // We don't know what the ID of the pty in the proxy is.
        // Queue up this message for later.
        this._writeQueue.push(msg);
      } else {
        this._writeFunc(this._id, msg);
      }
    }
  }
  
  _charsWritten(chars: number): void {
    this._outstandingWriteDataCount -= chars;
    this._onAvailableWriteBufferSizeChangeEventEmitter.fire(
      {totalBufferSize: MAXIMUM_WRITE_BUFFER_SIZE, availableDelta: chars});
  }

  write(data: string): void {
    this._outstandingWriteDataCount += data.length;
    const msg: WriteMessage = { type: TYPE_WRITE, id: this._id, data: data };
    this._writeMessage(this._id, msg);
  }

  getAvailableWriteBufferSize(): number {
    return MAXIMUM_WRITE_BUFFER_SIZE - this._outstandingWriteDataCount;
  }

  resize(cols: number, rows: number): void {
    const msg: ResizeMessage = { type: TYPE_RESIZE, id: this._id, rows: rows, columns: cols };
    this._writeMessage(this._id, msg);
  }
  
  permittedDataSize(size: number): void {
    const msg: PermitDataSizeMessage = { type: TYPE_PERMIT_DATA_SIZE, id: this._id, size: size };
    this._writeMessage(this._id, msg);
  }

  data(data: string): void {
    if (this._live) {
      this._onDataEventEmitter.fire(data);
    }
  }
  
  exit(): void {
    this._onExitEventEmitter.fire(undefined);
    this._live = false;
  }
  
  destroy(): void {    
    const msg: CloseMessage = { type: TYPE_CLOSE, id: this._id };
    this._writeMessage(this._id, msg);
  }
}

export abstract class ProxyPtyConnector {
  private _ptys: ProxyPty[] = [];
  private _messageBuffer = "";
  private _proxy: child_process.ChildProcess = null;

  constructor(private _log: Logger) {
  }

  start(): void {
    this._proxy = this._spawnServer();

    this._proxy.stdout.on('data', (data: Buffer) => {
      if (DEBUG_FINE) {
        this._log.debug("server -> main: ", data.toString("utf8"));
      }
      this._messageBuffer = this._messageBuffer + data.toString('utf8');
      this._processMessageBuffer();
    });

    this._proxy.stderr.on('data', (data: Buffer) => {
      this._log.warn('ptyproxy process stderr: ', data.toString());
    });

    this._proxy.on('close', code => {
      if (DEBUG_FINE) {
        this._log.debug('bridge process closed with code: ', code);
      }
    });
    
    this._proxy.on('exit', code => {
      if (DEBUG_FINE) {
        this._log.debug('bridge process exited with code: ', code);
      }
    });
    
    this._proxy.on('error', (err) => {
      this._log.severe("Failed to start server process. ", err);
    });
  }

  protected abstract _spawnServer(): child_process.ChildProcess;

  spawn(options: PtyOptions): Pty {
    let rows = 24;
    let columns = 80;
    const file = options.exe;
    const args = options.args;
    
    if (DEBUG_FINE) {
      this._log.debug("ptyproxy spawn file: ", file);
    }
    if (options !== undefined) {
      rows = options.rows !== undefined ? options.rows : rows;
      columns = options.cols !== undefined ? options.cols : columns;
    }
    const pty = new ProxyPty(this._sendMessage.bind(this));
    this._ptys.push(pty);
    const msg: CreatePtyMessage = { type: TYPE_CREATE, argv: [file, ...args], rows: rows, columns: columns,
      id: NULL_ID, env: options.env, extraEnv: options.extraEnv == null ? {} : options.extraEnv, cwd: options.cwd || null};
    this._sendMessage(null, msg);
    return pty;
  }
  
  destroy(): void {
    const msg: TerminateMessage = { type: TYPE_TERMINATE, id: -1 };
    this._sendMessage(null, msg);
  }

  private _processMessageBuffer(): void {
    while (true) {
      const end = this._messageBuffer.indexOf('\n');
      if (end !== -1) {
        const msgString = this._messageBuffer.slice(0, end);
        this._messageBuffer = this._messageBuffer.slice(end+1);
        const msg = <ProxyMessage> JSON.parse(msgString);
        this._processMessage(msg);
      } else {
        break;
      }
    }
  }

  private _processMessage(msg: ProxyMessage): void {
    const msgType = msg.type;
    
    if (msgType === TYPE_CREATED) {
      const createdPtyMsg = <CreatedPtyMessage> msg;
      for (let i=0; i<this._ptys.length; i++) {
        const pty = this._ptys[i];
        if (pty.getId() === NULL_ID) {
          pty.setId(createdPtyMsg.id);
          break;
        }
      }
      return;
    }
    
    if (msgType === TYPE_OUTPUT) {
      const outputMsg = <OutputMessage> msg;
      const pty = this._findPtyById(outputMsg.id);
      if (pty !== null) {
        pty.data(outputMsg.data);
      }
      return;
    }
    
    if (msgType === TYPE_CLOSED) {
      const closedMsg = <ClosedMessage> msg;
      const pty = this._findPtyById(closedMsg.id);
      if (pty !== null) {
        pty.exit();
      }
      return;
    }

    if (msgType === TYPE_OUTPUT_WRITTEN) {
      const outputWrittenMsg = <OutputWrittenMessage> msg;
      const pty = this._findPtyById(outputWrittenMsg.id);
      if (pty !== null) {
        pty._charsWritten(outputWrittenMsg.chars);
      }      
    }
  }
  
  private _findPtyById(id: number): ProxyPty {
    for (let i=0; i<this._ptys.length; i++) {
      if (this._ptys[i].getId() === id) {
        return this._ptys[i];
      }
    }
    return null;
  }
  
  private _sendMessage(id: number, msg: ProxyMessage): void {
    const msgText = JSON.stringify(msg);
    if (DEBUG_FINE) {
      this._log.debug("main -> server: ", msgText);
    }
    this._proxy.stdin.write(msgText + "\n", 'utf8');
  }
  
}