Merge branch 'feature/support-existing-rtcdatachannel-or-websocket-squashed' of https://github.com/TimSBSquare/noVNC

5 files changed, 99 insertions, 36 deletions

diff --git a/core/rfb.js b/core/rfb.js
index e3786cb..3390276 100644
--- a/core/rfb.js
+++ b/core/rfb.js
@@ -66,20 +66,25 @@ const extendedClipboardActionPeek    = 1 << 26;
 const extendedClipboardActionNotify  = 1 << 27;
 const extendedClipboardActionProvide = 1 << 28;
 
-
 export default class RFB extends EventTargetMixin {
-    constructor(target, url, options) {
+    constructor(target, urlOrChannel, options) {
         if (!target) {
             throw new Error("Must specify target");
         }
-        if (!url) {
-            throw new Error("Must specify URL");
+        if (!urlOrChannel) {
+            throw new Error("Must specify URL, WebSocket or RTCDataChannel");
         }
 
         super();
 
         this._target = target;
-        this._url = url;
+
+        if (typeof urlOrChannel === "string") {
+            this._url = urlOrChannel;
+        } else {
+            this._url = null;
+            this._rawChannel = urlOrChannel;
+        }
 
         // Connection details
         options = options || {};
@@ -275,6 +280,8 @@ export default class RFB extends EventTargetMixin {
                     break;
             }
             this._sock.off('close');
+            // Delete reference to raw channel to allow cleanup.
+            this._rawChannel = null;
         });
         this._sock.on('error', e => Log.Warn("WebSocket on-error event"));
 
@@ -501,16 +508,23 @@ export default class RFB extends EventTargetMixin {
     _connect() {
         Log.Debug(">> RFB.connect");
 
-        Log.Info("connecting to " + this._url);
-
-        try {
-            // WebSocket.onopen transitions to the RFB init states
-            this._sock.open(this._url, this._wsProtocols);
-        } catch (e) {
-            if (e.name === 'SyntaxError') {
-                this._fail("Invalid host or port (" + e + ")");
-            } else {
-                this._fail("Error when opening socket (" + e + ")");
+        if (this._url) {
+            try {
+                Log.Info(`connecting to ${this._url}`);
+                this._sock.open(this._url, this._wsProtocols);
+            } catch (e) {
+                if (e.name === 'SyntaxError') {
+                    this._fail("Invalid host or port (" + e + ")");
+                } else {
+                    this._fail("Error when opening socket (" + e + ")");
+                }
+            }
+        } else {
+            try {
+                Log.Info(`attaching ${this._rawChannel} to Websock`);
+                this._sock.attach(this._rawChannel);
+            } catch (e) {
+                this._fail("Error attaching channel (" + e + ")");
             }
         }
 
diff --git a/core/websock.js b/core/websock.js
index 8e606bc..89ccdd9 100644
--- a/core/websock.js
+++ b/core/websock.js
@@ -1,10 +1,10 @@
 /*
- * Websock: high-performance binary WebSockets
+ * Websock: high-performance buffering wrapper
  * Copyright (C) 2019 The noVNC Authors
  * Licensed under MPL 2.0 (see LICENSE.txt)
  *
- * Websock is similar to the standard WebSocket object but with extra
- * buffer handling.
+ * Websock is similar to the standard WebSocket / RTCDataChannel object
+ * but with extra buffer handling.
  *
  * Websock has built-in receive queue buffering; the message event
  * does not contain actual data but is simply a notification that
@@ -19,9 +19,37 @@ import * as Log from './util/logging.js';
 // at the moment.  It may be valuable to turn it on in the future.
 const MAX_RQ_GROW_SIZE = 40 * 1024 * 1024;  // 40 MiB
 
+// Constants pulled from RTCDataChannelState enum
+// https://developer.mozilla.org/en-US/docs/Web/API/RTCDataChannel/readyState#RTCDataChannelState_enum
+const DataChannel = {
+    CONNECTING: "connecting",
+    OPEN: "open",
+    CLOSING: "closing",
+    CLOSED: "closed"
+};
+
+const ReadyStates = {
+    CONNECTING: [WebSocket.CONNECTING, DataChannel.CONNECTING],
+    OPEN: [WebSocket.OPEN, DataChannel.OPEN],
+    CLOSING: [WebSocket.CLOSING, DataChannel.CLOSING],
+    CLOSED: [WebSocket.CLOSED, DataChannel.CLOSED],
+};
+
+// Properties a raw channel must have, WebSocket and RTCDataChannel are two examples
+const rawChannelProps = [
+    "send",
+    "close",
+    "binaryType",
+    "onerror",
+    "onmessage",
+    "onopen",
+    "protocol",
+    "readyState",
+];
+
 export default class Websock {
     constructor() {
-        this._websocket = null;  // WebSocket object
+        this._websocket = null;  // WebSocket or RTCDataChannel object
 
         this._rQi = 0;           // Receive queue index
         this._rQlen = 0;         // Next write position in the receive queue
@@ -140,7 +168,7 @@ export default class Websock {
     // Send Queue
 
     flush() {
-        if (this._sQlen > 0 && this._websocket.readyState === WebSocket.OPEN) {
+        if (this._sQlen > 0 && ReadyStates.OPEN.indexOf(this._websocket.readyState) >= 0) {
             this._websocket.send(this._encodeMessage());
             this._sQlen = 0;
         }
@@ -177,13 +205,26 @@ export default class Websock {
     }
 
     open(uri, protocols) {
+        this.attach(new WebSocket(uri, protocols));
+    }
+
+    attach(rawChannel) {
         this.init();
 
-        this._websocket = new WebSocket(uri, protocols);
-        this._websocket.binaryType = 'arraybuffer';
+        // Must get object and class methods to be compatible with the tests.
+        const channelProps = [...Object.keys(rawChannel), ...Object.getOwnPropertyNames(Object.getPrototypeOf(rawChannel))];
+        for (let i = 0; i < rawChannelProps.length; i++) {
+            const prop = rawChannelProps[i];
+            if (channelProps.indexOf(prop) < 0) {
+                throw new Error('Raw channel missing property: ' + prop);
+            }
+        }
 
+        this._websocket = rawChannel;
+        this._websocket.binaryType = "arraybuffer";
         this._websocket.onmessage = this._recvMessage.bind(this);
-        this._websocket.onopen = () => {
+
+        const onOpen = () => {
             Log.Debug('>> WebSock.onopen');
             if (this._websocket.protocol) {
                 Log.Info("Server choose sub-protocol: " + this._websocket.protocol);
@@ -192,11 +233,21 @@ export default class Websock {
             this._eventHandlers.open();
             Log.Debug("<< WebSock.onopen");
         };
+
+        // If the readyState cannot be found this defaults to assuming it's not open.
+        const isOpen = ReadyStates.OPEN.indexOf(this._websocket.readyState) >= 0;
+        if (isOpen) {
+            onOpen();
+        } else {
+            this._websocket.onopen = onOpen;
+        }
+
         this._websocket.onclose = (e) => {
             Log.Debug(">> WebSock.onclose");
             this._eventHandlers.close(e);
             Log.Debug("<< WebSock.onclose");
         };
+
         this._websocket.onerror = (e) => {
             Log.Debug(">> WebSock.onerror: " + e);
             this._eventHandlers.error(e);
@@ -206,8 +257,8 @@ export default class Websock {
 
     close() {
         if (this._websocket) {
-            if ((this._websocket.readyState === WebSocket.OPEN) ||
-                    (this._websocket.readyState === WebSocket.CONNECTING)) {
+            if (ReadyStates.CONNECTING.indexOf(this._websocket.readyState) >= 0 ||
+                ReadyStates.OPEN.indexOf(this._websocket.readyState) >= 0) {
                 Log.Info("Closing WebSocket connection");
                 this._websocket.close();
             }
diff --git a/docs/API.md b/docs/API.md
index d78360a..aa5aea7 100644
--- a/docs/API.md
+++ b/docs/API.md
@@ -165,9 +165,9 @@ connection to a specified VNC server.
     existing contents of the `HTMLElement` will be untouched, but new
     elements will be added during the lifetime of the `RFB` object.
 
-**`url`**
+**`urlOrDataChannel`**
   - A `DOMString` specifying the VNC server to connect to. This must be
-    a valid WebSocket URL.
+    a valid WebSocket URL. This can also be a `WebSocket` or `RTCDataChannel`.
 
 **`options`** *Optional*
   - An `Object` specifying extra details about how the connection
diff --git a/tests/fake.websocket.js b/tests/fake.websocket.js
index 4db3c59..8fb02c5 100644
--- a/tests/fake.websocket.js
+++ b/tests/fake.websocket.js
@@ -6,6 +6,10 @@ export default class FakeWebSocket {
         this.binaryType = "arraybuffer";
         this.extensions = "";
 
+        this.onerror = null;
+        this.onmessage = null;
+        this.onopen = null;
+
         if (!protocols || typeof protocols === 'string') {
             this.protocol = protocols;
         } else {
diff --git a/tests/test.websock.js b/tests/test.websock.js
index a8e17e6..6f35ec3 100644
--- a/tests/test.websock.js
+++ b/tests/test.websock.js
@@ -254,13 +254,7 @@ describe('Websock', function () {
         beforeEach(function () {
             sock = new Websock();
             // eslint-disable-next-line no-global-assign
-            WebSocket = sinon.spy();
-            WebSocket.OPEN = oldWS.OPEN;
-            WebSocket.CONNECTING = oldWS.CONNECTING;
-            WebSocket.CLOSING = oldWS.CLOSING;
-            WebSocket.CLOSED = oldWS.CLOSED;
-
-            WebSocket.prototype.binaryType = 'arraybuffer';
+            WebSocket = sinon.spy(FakeWebSocket);
         });
 
         describe('opening', function () {
@@ -278,7 +272,7 @@ describe('Websock', function () {
 
         describe('closing', function () {
             beforeEach(function () {
-                sock.open('ws://');
+                sock.open('ws://localhost');
                 sock._websocket.close = sinon.spy();
             });
 
@@ -324,7 +318,7 @@ describe('Websock', function () {
                 sock.on('open', sinon.spy());
                 sock.on('close', sinon.spy());
                 sock.on('error', sinon.spy());
-                sock.open('ws://');
+                sock.open('ws://localhost');
             });
 
             it('should call _recvMessage on a message', function () {