From 2a8e3f41d7a10a06e52d477cb0710f354aaf3281 Mon Sep 17 00:00:00 2001
From: Klesti Selimaj <klesti.selimay@icloud.com>
Date: Mon, 17 Aug 2026 06:33:21 +0200
Subject: [PATCH] Init

---
 src/App.tsx        | 50 ++------------------------------------------
 src/protocol/ws.ts | 52 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 52 insertions(+), 50 deletions(-)

diff --git a/src/App.tsx b/src/App.tsx
index 0bf448c..aceff92 100644
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -4,7 +4,6 @@ import { invoke } from "@tauri-apps/api/core";
 import "./App.css";
 import { EnclaveWebSocket } from "./protocol/ws";
 import * as ed from "@noble/ed25519";
-import { base58 } from "@scure/base";
 import { sha512 } from "@noble/hashes/sha2.js";
 
 ed.hashes.sha512 = sha512;
@@ -20,55 +19,10 @@ function App() {
     initialized.current = true;
 
     (async () => {
-      const ws = new EnclaveWebSocket("ws://localhost:3415");
-
-      const timestamp = Date.now();
-      const hostname = "localhost:3415";
-
-      const msg = new TextEncoder().encode(`${timestamp}@${hostname}`);
-
+      const ws = new EnclaveWebSocket("localhost:3415");
       ws.clientSecretKey = ed.utils.randomSecretKey();
-
       ws.clientPublicKey = ed.getPublicKey(ws.clientSecretKey);
-
-      const publicKeyString = base58.encode(ws.clientPublicKey);
-
-      ws.send({
-        method: "Initialize",
-        public_key: publicKeyString,
-        signature: base58.encode(ed.sign(msg, ws.clientSecretKey)),
-
-        timestamp,
-        hostname,
-      });
-
-      const initialized = await ws.read();
-
-      if (initialized.method !== "Initialized") {
-        console.error("Invalid method from server, closing. ", initialized);
-        ws.websocket.close();
-        return;
-      }
-
-      if (initialized.hostname !== hostname) {
-        console.error("Server is trying to be a middle man", initialized);
-        ws.websocket.close();
-        return;
-      }
-
-      const sigMsg = new TextEncoder().encode(
-        `${initialized.timestamp}@${hostname}@${publicKeyString}`,
-      );
-
-      ws.serverPublicKey = base58.decode(initialized.public_key);
-      const signature = base58.decode(initialized.signature);
-
-      if (!ed.verify(signature, sigMsg, ws.serverPublicKey)) {
-        console.error("Invalid signature", initialized);
-        ws.websocket.close();
-        return;
-      }
-
+      ws.init();
       console.log("OK");
     })();
   }, []);
diff --git a/src/protocol/ws.ts b/src/protocol/ws.ts
index 153ecd3..f59b44b 100644
--- a/src/protocol/ws.ts
+++ b/src/protocol/ws.ts
@@ -1,7 +1,13 @@
 import { ClientMethod, ServerMethod } from "./protocol";
+import { base58 } from "@scure/base";
+import * as ed from "@noble/ed25519";
+import { sha512 } from "@noble/hashes/sha2.js";
+
+ed.hashes.sha512 = sha512;
 
 export class EnclaveWebSocket {
   public websocket: WebSocket;
+  public hostname: string;
   onOpenQueue: Array<() => void>;
 
   public clientPublicKey: Uint8Array;
@@ -9,17 +15,59 @@ export class EnclaveWebSocket {
 
   public serverPublicKey?: Uint8Array;
 
-  public constructor(url: string | URL) {
+  public constructor(hostname: string) {
     this.clientPublicKey = new Uint8Array();
     this.clientSecretKey = new Uint8Array();
     this.onOpenQueue = new Array();
 
-    this.websocket = new WebSocket(url);
+    this.hostname = hostname;
+    this.websocket = new WebSocket("ws://" + hostname);
     this.websocket.onopen = () => {
       this.onOpenQueue.forEach((fun) => fun());
     };
   }
 
+  public async init() {
+    const publicKeyString = base58.encode(this.clientPublicKey);
+
+    const timestamp = Date.now();
+
+    const msg = new TextEncoder().encode(`${timestamp}@${this.hostname}`);
+
+    this.send({
+      method: "Initialize",
+      public_key: publicKeyString,
+      signature: base58.encode(ed.sign(msg, this.clientSecretKey)),
+
+      timestamp,
+      hostname: this.hostname,
+    });
+
+    const initialized = await this.read();
+
+    if (initialized.method !== "Initialized") {
+      this.websocket.close();
+      throw Error("Invalid method from server, closing. ");
+    }
+
+    if (initialized.hostname !== this.hostname) {
+      this.websocket.close();
+      throw Error("Server is trying to be a middle man");
+    }
+
+    const sigMsg = new TextEncoder().encode(
+      `${initialized.timestamp}@${this.hostname}@${publicKeyString}`,
+    );
+
+    this.serverPublicKey = base58.decode(initialized.public_key);
+    const signature = base58.decode(initialized.signature);
+
+    if (!ed.verify(signature, sigMsg, this.serverPublicKey)) {
+      this.websocket.close();
+      throw Error("Invalid signature");
+    }
+  }
+
   public async send(method: ServerMethod) {
     if (this.websocket.readyState !== WebSocket.OPEN) {
       return new Promise<void>((ok) => {