From bcfe1d39278bfdd31bae3071f220d00c813888ac Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Thu, 9 May 2024 23:48:35 +0800
Subject: [PATCH] :bug: Bug fixes on E2EE

---
 .idea/workspace.xml  | 13 +++++--------
 pkg/server/ws.go     |  2 +-
 pkg/services/e2ee.go |  8 ++++++--
 3 files changed, 12 insertions(+), 11 deletions(-)

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index 58e924a..cfe426d 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -4,14 +4,10 @@
     <option name="autoReloadType" value="ALL" />
   </component>
   <component name="ChangeListManager">
-    <list default="true" id="3fefb2c4-b6f9-466b-a523-53352e8d6f95" name="更改" comment=":recycle: Improved the notification subscriber API">
-      <change afterPath="$PROJECT_DIR$/pkg/models/unified.go" afterDir="false" />
-      <change afterPath="$PROJECT_DIR$/pkg/services/e2ee.go" afterDir="false" />
+    <list default="true" id="3fefb2c4-b6f9-466b-a523-53352e8d6f95" name="更改" comment=":sparkles: E2EE Key Exchange">
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/cmd/main.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/cmd/main.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/server/notify_ws.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/server/ws.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/server/startup.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/server/startup.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/services/notifications.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/services/notifications.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/pkg/server/ws.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/server/ws.go" afterDir="false" />
+      <change beforePath="$PROJECT_DIR$/pkg/services/e2ee.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/services/e2ee.go" afterDir="false" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
@@ -168,7 +164,8 @@
     <MESSAGE value=":bug: Fix new realm owner missing permissions" />
     <MESSAGE value=":zap: Use map to improve message delivery time" />
     <MESSAGE value=":recycle: Improved the notification subscriber API" />
-    <option name="LAST_COMMIT_MESSAGE" value=":recycle: Improved the notification subscriber API" />
+    <MESSAGE value=":sparkles: E2EE Key Exchange" />
+    <option name="LAST_COMMIT_MESSAGE" value=":sparkles: E2EE Key Exchange" />
   </component>
   <component name="VgoProject">
     <settings-migrated>true</settings-migrated>
diff --git a/pkg/server/ws.go b/pkg/server/ws.go
index 2800136..ecdf629 100644
--- a/pkg/server/ws.go
+++ b/pkg/server/ws.go
@@ -60,7 +60,7 @@ func listenWebsocket(c *websocket.Conn) {
 			if len(req.RequestID) <= 0 || len(req.KeypairID) <= 0 {
 				message = lo.ToPtr(models.UnifiedCommandFromError(fmt.Errorf("invalid request")))
 			}
-			services.KexProvide(user.ID, req.RequestID, req.KeypairID, payload)
+			services.KexProvide(user.ID, req.RequestID, req.KeypairID, packet)
 		default:
 			message = lo.ToPtr(models.UnifiedCommandFromError(fmt.Errorf("unknown action")))
 		}
diff --git a/pkg/services/e2ee.go b/pkg/services/e2ee.go
index 34aaf40..19b4d07 100644
--- a/pkg/services/e2ee.go
+++ b/pkg/services/e2ee.go
@@ -28,8 +28,12 @@ func KexRequest(conn *websocket.Conn, requestId, keypairId string, ownerId uint,
 	}
 
 	flag := false
-	for conn := range wsConn[ownerId] {
-		if conn.WriteMessage(1, models.UnifiedCommand{
+	for c := range wsConn[ownerId] {
+		if c == conn {
+			continue
+		}
+
+		if c.WriteMessage(1, models.UnifiedCommand{
 			Action: "kex.request",
 			Payload: fiber.Map{
 				"request_id": requestId,