secure connection works

src/connection.ts

@@ -1,13 +1,13 @@
-import Websock from "./websock";
-import * as message from "./message.js";
-import * as rendezvous from "./rendezvous.js";
-import { loadVp9, loadOpus } from "./codec";
-import * as globals from "./globals";
+import Websock from './websock';
+import * as message from './message.js';
+import * as rendezvous from './rendezvous.js';
+import { loadVp9, loadOpus } from './codec';
+import * as globals from './globals';
 
 const PORT = 21116;
-const HOST = "rs-sg.rustdesk.com";
-const licenceKey = "";
-const SCHEMA = "ws://";
+const HOST = 'rs-sg.rustdesk.com';
+const licenceKey = '';
+const SCHEMA = 'ws://';
 
 export default class Connection {
   _msgs: any[];
@@ -38,7 +38,7 @@ export default class Connection {
     await ws.open();
     const connType = rendezvous.ConnType.DEFAULT_CONN;
     const natType = rendezvous.NatType.SYMMETRIC;
-    const punchHoleRequest = rendezvous.PunchHoleRequest.fromJSON({
+    const punchHoleRequest = rendezvous.PunchHoleRequest.fromPartial({
       id,
       licenceKey,
       connType,
@@ -72,9 +72,9 @@ export default class Connection {
     const uuid = rr.uuid;
     const ws = new Websock(uri);
     await ws.open();
-    console.log("Connected to relay server");
+    console.log('Connected to relay server');
     this._ws = ws;
-    const requestRelay = rendezvous.RequestRelay.fromJSON({
+    const requestRelay = rendezvous.RequestRelay.fromPartial({
       licenceKey,
       uuid,
     });
@@ -85,8 +85,60 @@ export default class Connection {
   async secure(pk: Uint8Array | undefined) {
     if (pk) {
       const RS_PK = 'OeVuKk5nlHiXp+APNn0Y3pC1Iwpwn44JGqrQCsWqmBw=';
-      let pk_id = await globals.verify(pk, RS_PK);
+      try {
+        pk = await globals.verify(pk, RS_PK).catch();
+        if (pk?.length != 32) {
+          pk = undefined;
+        }
+      } catch (e) {
+        console.error(e);
+        pk = undefined;
+      }
+      if (!pk) console.error('Handshake failed: invalid public key from rendezvous server');
     }
+    if (!pk) {
+      // send an empty message out in case server is setting up secure and waiting for first message
+      await this._ws?.sendMessage({});
+      return;
+    }
+    const msg = this._ws?.parseMessage(await this._ws?.next());
+    let signedId: any = msg?.signedId;
+    if (!signedId) {
+      console.error("Handshake failed: invalid message type");
+      await this._ws?.sendMessage({});
+      return;
+    }
+    try {
+      signedId = await globals.verify(signedId.id, Uint8Array.from(pk!));
+    } catch (e) {
+      console.error(e);
+      // fall back to non-secure connection in case pk mismatch
+      console.error("pk mismatch, fall back to non-secure");
+      const publicKey = message.PublicKey.fromPartial({});
+      await this._ws?.sendMessage({ publicKey });
+      return;
+    }
+    signedId = new TextDecoder().decode(signedId!);
+    const tmp = signedId.split('\0');
+    const id = tmp[0];
+    let theirPk = tmp[1];
+    if (id != this._id!) {
+      console.error("Handshake failed: sign failure");
+      await this._ws?.sendMessage({});
+      return;
+    }
+    theirPk = globals.decodeBase64(theirPk);
+    if (theirPk.length != 32) {
+      console.error("Handshake failed: invalid public box key length from peer");
+      await this._ws?.sendMessage({});
+      return;
+    }
+    const [mySk, asymmetricValue] = globals.genBoxKeyPair();
+    const secretKey = globals.genSecretKey();
+    const symmetricValue = globals.seal(secretKey, theirPk, mySk);
+    const publicKey = message.PublicKey.fromPartial({ asymmetricValue, symmetricValue });
+    await this._ws?.sendMessage({ publicKey });
+    this._ws?.setSecretKey(secretKey)
   }
 }
 
@@ -97,17 +149,17 @@ async function testDelay() {
 }
 
 function getDefaultUri(isRelay: Boolean = false): string {
-  const host = localStorage.getItem("host");
-  return SCHEMA + (host || HOST) + ":" + (PORT + (isRelay ? 3 : 2));
+  const host = localStorage.getItem('host');
+  return SCHEMA + (host || HOST) + ':' + (PORT + (isRelay ? 3 : 2));
 }
 
 function getrUriFromRs(uri: string): string {
-  if (uri.indexOf(":") > 0) {
-    const tmp = uri.split(":");
+  if (uri.indexOf(':') > 0) {
+    const tmp = uri.split(':');
     const port = parseInt(tmp[1]);
-    uri = tmp[0] + ":" + (port + 2);
+    uri = tmp[0] + ':' + (port + 2);
   } else {
-    uri += ":" + (PORT + 3);
+    uri += ':' + (PORT + 3);
   }
   return SCHEMA + uri;
 }

src/globals.js

@@ -23,8 +23,38 @@ export async function verify(signed, pk) {
     await _sodium.ready;
     sodium = _sodium;
   }
-  pk = sodium.from_base64(pk, sodium.base64_variants.ORIGINAL);
+  if (typeof pk == 'string') {
+    pk = decodeBase64(pk);
+  }
   return sodium.crypto_sign_open(signed, pk);
 }
 
+export function decodeBase64(pk) {
+  return sodium.from_base64(pk, sodium.base64_variants.ORIGINAL);
+}
+
+export function genBoxKeyPair() {
+  const pair = sodium.crypto_box_keypair();
+  const sk = pair.privateKey;
+  const pk = pair.publicKey;
+  return [sk, pk];
+}
+
+export function genSecretKey() {
+  return sodium.crypto_secretbox_keygen();
+}
+
+export function seal(unsigned, theirPk, ourSk) {
+  const nonce = Uint8Array.from(Array(24).fill(0));
+  return sodium.crypto_box_easy(unsigned, nonce, theirPk, ourSk);
+}
+
+export function encrypt(unsigned, nonce, key) {
+  return sodium.crypto_secretbox_easy(unsigned, nonce, key);
+}
+
+export function decrypt(signed, nonce, key) {
+  return sodium.crypto_secretbox_open_easy(signed, nonce, key);
+}
+
 window.startConn = startConn;

src/websock.ts

@@ -10,6 +10,7 @@ export default class Websock {
   _buf: Uint8Array[];
   _status: any;
   _latency: number;
+  _secretKey: any;
 
   constructor(uri: string) {
     this._eventHandlers = {
@@ -30,16 +31,20 @@ export default class Websock {
     return this._latency;
   }
 
+  setSecretKey(key: any) {
+    this._secretKey = key;
+  }
+
   sendMessage(data: any) {
     this._websocket.send(
-      message.Message.encode(message.Message.fromJSON(data)).finish()
+      message.Message.encode(message.Message.fromPartial(data)).finish()
     );
   }
 
   sendRendezvous(data: any) {
     this._websocket.send(
       rendezvous.RendezvousMessage.encode(
-        rendezvous.RendezvousMessage.fromJSON(data)
+        rendezvous.RendezvousMessage.fromPartial(data)
       ).finish()
     );
   }