Quick hack / experiment for a AWS IoT SecureTunnelling SSH client using node.js

index.js

const { readFileSync } = require('fs');

const { Client } = require('ssh2');
var WebSocketClient = require('websocket').client;
var protobuf = require("protobufjs");

var messageProto;
protobuf.load("message.proto", function(err, root) {
    if (err)
        throw err;
    messageProto = root.lookupType("com.amazonaws.iot.securedtunneling.Message");
});

var EventEmitter = require('events').EventEmitter;

var region = "eu-central-1";
var accessToken = "TODO";

class MySocket extends EventEmitter {
    constructor() {
        super();
        this._paused = true;
        this.writable = true;
        this.connecting = true;
        this._readableState = { "ended": false };
        this._buffer_to_ssh = Buffer.alloc(0);
        this._buffer_from_websocket = Buffer.alloc(0);

        // TODO: Which listeners do I have to support?
        this.on("newListener", (event, listener) => {
            console.log("new listener: " + event);
        });
        var ws = new WebSocketClient();
        ws.connect("wss://data.tunneling.iot." + region + ".amazonaws.com/tunnel?local-proxy-mode=source&access-token=" + accessToken, "aws.iot.securetunneling-1.0");
        ws.on('connect', (conn) => {
            this._ws_conn = conn;
            conn.on('error', (err) => { console.log("Websocket error: " + err); });
            conn.on('message', (message) => { this._receive_websocket(message.binaryData); });

            console.log("Connected");
            this._send_proto_message({"type": 2, "streamId": 1});

            this.connecting = false;
            this.emit('connect');
        });
        ws.on('connectFailed', (err) => { this.emit('error', err); });
    }

    _receive_websocket(data) {
        this._buffer_from_websocket = Buffer.concat([this._buffer_from_websocket, data]);
        this._parse_websocket();
    }

    _parse_websocket() {
        if (this._buffer_from_websocket.length < 2) {
            // Wait for more data
            return;
        }
        var length = this._buffer_from_websocket.readUInt16BE();
        if (this._buffer_from_websocket.length < 2 + length) {
            // Wait for more data
            return;
        }
        // We got a whole packet!
        var this_packet = this._buffer_from_websocket.subarray(2, 2 + length);
        this._buffer_from_websocket = this._buffer_from_websocket.subarray(2 + length);

        var message = messageProto.decode(this_packet);
        console.log("Received message on websocket: " + JSON.stringify(message));
        if (message.type == "DATA" || message.type == 1) {
            this._buffer_to_ssh = Buffer.concat([this._buffer_to_ssh, message.payload]);
            if (!this._paused) {
                this.resume();
            }
        } else if (message.type == "STREAM_RESET" || message.type == 3) {
            console.log("Connection closed");
            this.emit('close', false);
        } else {
            this.emit("error", new Error("Unexpected message type received: " + message.type));
        }

        // Check if we have more packets in the buffer
        this._parse_websocket();
    }

    destroy() {
        // TODO
    }

    end() {
        // TODO
    }

    pause() {
        this._paused = true;
    }

    resume() {
        this._paused = false;
        if (this._buffer_to_ssh.length > 0) {
            var buffer = this._buffer_to_ssh;
            this._buffer_to_ssh = Buffer.alloc(0);
            this.emit('data', buffer);
        }
    }

    write(data) {
        this._send_proto_message({"type": 1, "streamId": 1, "payload": data});
    }

    _send_proto_message(data) {
        console.log("Sending on websocket: " + JSON.stringify(data));
        var encoded = messageProto.encode(data).finish();
        var length = Buffer.alloc(2);
        length.writeUInt16BE(encoded.length);
        this._ws_conn.sendBytes(Buffer.concat([length, encoded]));
    }
};

const sock = new MySocket();

const conn = new Client();
conn.on('ready', () => {
  console.log('Client :: ready');
  conn.shell((err, stream) => {
    if (err) throw err;
    stream.on('close', () => {
      console.log('Stream :: close');
      conn.end();
    }).on('data', (data) => {
      console.log('OUTPUT: ' + data);
    });
    // I wanted to test something with a longer connection. Hence I did this.
    if (false) {
      stream.end('ls -l\nexit\n');
    } else {
      stream.write('ls -l\nsleep 300\nsurvived\n');
    }
  });
}).on('error', (err) => {
    console.log("ERROR: " + err);
}).connect({
  //host: 'localhost',
  //port: 2222,
  sock: sock,
  keepaliveInterval: 45 * 1000, // Dunno what a good value is, but AFAIR if we do not send something every minute, the AWS kills the connection
  debug: (msg) => {console.log("DEBUG: " + msg);},
  username: 'Name of the user',
  privateKey: readFileSync('/tmp/key')
});

message.proto

syntax = "proto3";

package com.amazonaws.iot.securedtunneling;

option java_outer_classname = "Protobuf";
option optimize_for = LITE_RUNTIME;

message Message {
    Type    type         = 1;
    int32   streamId     = 2;
    bool    ignorable    = 3;
    bytes   payload      = 4;

    enum Type {
        UNKNOWN = 0;
        DATA = 1;
        STREAM_START = 2;
        STREAM_RESET = 3;
        SESSION_RESET = 4;
    }
}

package.json

{
  "name": "foo",
  "version": "1.0.0",
  "main": "index.js",
  "scripts": {
    "test": "echo \"Error: no test specified\" && exit 1"
  },
  "author": "",
  "license": "ISC",
  "description": "",
  "dependencies": {
    "protobufjs": "^7.4.0",
    "ssh2": "^1.16.0",
    "websocket": "^1.0.35"
  }
}