From 3ba152252ecf13c0a1bd63b2cad77c0eba501176 Mon Sep 17 00:00:00 2001
From: LittleSheep <littlesheep.code@hotmail.com>
Date: Thu, 9 May 2024 23:35:13 +0800
Subject: [PATCH] :sparkles: E2EE Key Exchange

---
 .idea/workspace.xml           | 76 ++++++++++++++++++++--------------
 pkg/cmd/main.go               |  1 +
 pkg/models/unified.go         | 21 ++++++++++
 pkg/server/notify_ws.go       | 25 ------------
 pkg/server/startup.go         |  6 +--
 pkg/server/ws.go              | 77 +++++++++++++++++++++++++++++++++++
 pkg/services/e2ee.go          | 77 +++++++++++++++++++++++++++++++++++
 pkg/services/notifications.go |  5 ++-
 8 files changed, 229 insertions(+), 59 deletions(-)
 create mode 100644 pkg/models/unified.go
 delete mode 100644 pkg/server/notify_ws.go
 create mode 100644 pkg/server/ws.go
 create mode 100644 pkg/services/e2ee.go

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index d783639..58e924a 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -4,10 +4,13 @@
     <option name="autoReloadType" value="ALL" />
   </component>
   <component name="ChangeListManager">
-    <list default="true" id="3fefb2c4-b6f9-466b-a523-53352e8d6f95" name="更改" comment=":zap: Use map to improve message delivery time">
+    <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" />
       <change beforePath="$PROJECT_DIR$/.idea/workspace.xml" beforeDir="false" afterPath="$PROJECT_DIR$/.idea/workspace.xml" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/models/notifications.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/models/notifications.go" afterDir="false" />
-      <change beforePath="$PROJECT_DIR$/pkg/server/notifications_api.go" beforeDir="false" afterPath="$PROJECT_DIR$/pkg/server/notifications_api.go" 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" />
     </list>
     <option name="SHOW_DIALOG" value="false" />
@@ -26,13 +29,13 @@
   <component name="Git.Settings">
     <option name="RECENT_BRANCH_BY_REPOSITORY">
       <map>
-        <entry key="$PROJECT_DIR$" value="master" />
+        <entry key="$PROJECT_DIR$" value="features/kex" />
       </map>
     </option>
     <option name="RECENT_GIT_ROOT_PATH" value="$PROJECT_DIR$" />
   </component>
   <component name="ProblemsViewState">
-    <option name="selectedTabId" value="ProjectErrors" />
+    <option name="selectedTabId" value="CurrentFile" />
   </component>
   <component name="ProjectColorInfo">{
   &quot;customColor&quot;: &quot;&quot;,
@@ -43,32 +46,32 @@
     <option name="hideEmptyMiddlePackages" value="true" />
     <option name="showLibraryContents" value="true" />
   </component>
-  <component name="PropertiesComponent">{
-  &quot;keyToString&quot;: {
-    &quot;DefaultGoTemplateProperty&quot;: &quot;Go File&quot;,
-    &quot;Go 构建.Backend.executor&quot;: &quot;Run&quot;,
-    &quot;RunOnceActivity.ShowReadmeOnStart&quot;: &quot;true&quot;,
-    &quot;RunOnceActivity.go.formatter.settings.were.checked&quot;: &quot;true&quot;,
-    &quot;RunOnceActivity.go.migrated.go.modules.settings&quot;: &quot;true&quot;,
-    &quot;RunOnceActivity.go.modules.automatic.dependencies.download&quot;: &quot;true&quot;,
-    &quot;RunOnceActivity.go.modules.go.list.on.any.changes.was.set&quot;: &quot;true&quot;,
-    &quot;git-widget-placeholder&quot;: &quot;master&quot;,
-    &quot;go.import.settings.migrated&quot;: &quot;true&quot;,
-    &quot;go.sdk.automatically.set&quot;: &quot;true&quot;,
-    &quot;last_opened_file_path&quot;: &quot;/Users/littlesheep/Documents/Projects/Hydrogen/Passport/pkg/server/ui&quot;,
-    &quot;node.js.detected.package.eslint&quot;: &quot;true&quot;,
-    &quot;node.js.selected.package.eslint&quot;: &quot;(autodetect)&quot;,
-    &quot;nodejs_package_manager_path&quot;: &quot;npm&quot;,
-    &quot;run.code.analysis.last.selected.profile&quot;: &quot;pProject Default&quot;,
-    &quot;settings.editor.selected.configurable&quot;: &quot;preferences.lookFeel&quot;,
-    &quot;vue.rearranger.settings.migration&quot;: &quot;true&quot;
+  <component name="PropertiesComponent"><![CDATA[{
+  "keyToString": {
+    "DefaultGoTemplateProperty": "Go File",
+    "Go 构建.Backend.executor": "Run",
+    "RunOnceActivity.ShowReadmeOnStart": "true",
+    "RunOnceActivity.go.formatter.settings.were.checked": "true",
+    "RunOnceActivity.go.migrated.go.modules.settings": "true",
+    "RunOnceActivity.go.modules.automatic.dependencies.download": "true",
+    "RunOnceActivity.go.modules.go.list.on.any.changes.was.set": "true",
+    "git-widget-placeholder": "master",
+    "go.import.settings.migrated": "true",
+    "go.sdk.automatically.set": "true",
+    "last_opened_file_path": "/Users/littlesheep/Documents/Projects/Hydrogen/Passport/pkg/server/ui",
+    "node.js.detected.package.eslint": "true",
+    "node.js.selected.package.eslint": "(autodetect)",
+    "nodejs_package_manager_path": "npm",
+    "run.code.analysis.last.selected.profile": "pProject Default",
+    "settings.editor.selected.configurable": "preferences.lookFeel",
+    "vue.rearranger.settings.migration": "true"
   },
-  &quot;keyToStringList&quot;: {
-    &quot;DatabaseDriversLRU&quot;: [
-      &quot;postgresql&quot;
+  "keyToStringList": {
+    "DatabaseDriversLRU": [
+      "postgresql"
     ]
   }
-}</component>
+}]]></component>
   <component name="RecentsManager">
     <key name="CopyFile.RECENT_KEYS">
       <recent name="$PROJECT_DIR$/pkg/server/ui" />
@@ -122,7 +125,19 @@
       <map>
         <entry key="MAIN">
           <value>
-            <State />
+            <State>
+              <option name="FILTERS">
+                <map>
+                  <entry key="branch">
+                    <value>
+                      <list>
+                        <option value="master" />
+                      </list>
+                    </value>
+                  </entry>
+                </map>
+              </option>
+            </State>
           </value>
         </entry>
       </map>
@@ -152,7 +167,8 @@
     <MESSAGE value=":bug: Dumb man make dumb mistake again" />
     <MESSAGE value=":bug: Fix new realm owner missing permissions" />
     <MESSAGE value=":zap: Use map to improve message delivery time" />
-    <option name="LAST_COMMIT_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" />
   </component>
   <component name="VgoProject">
     <settings-migrated>true</settings-migrated>
diff --git a/pkg/cmd/main.go b/pkg/cmd/main.go
index f451f18..ac2eec5 100644
--- a/pkg/cmd/main.go
+++ b/pkg/cmd/main.go
@@ -72,6 +72,7 @@ func main() {
 	quartz.AddFunc("@every 60m", services.DoAutoSignoff)
 	quartz.AddFunc("@every 60m", services.DoAutoAuthCleanup)
 	quartz.AddFunc("@every 60m", services.DoAutoDatabaseCleanup)
+	quartz.AddFunc("@every 5m", services.KexCleanup)
 	quartz.Start()
 
 	// Messages
diff --git a/pkg/models/unified.go b/pkg/models/unified.go
new file mode 100644
index 0000000..bea3048
--- /dev/null
+++ b/pkg/models/unified.go
@@ -0,0 +1,21 @@
+package models
+
+import jsoniter "github.com/json-iterator/go"
+
+type UnifiedCommand struct {
+	Action  string `json:"w"`
+	Message string `json:"m"`
+	Payload any    `json:"p"`
+}
+
+func UnifiedCommandFromError(err error) UnifiedCommand {
+	return UnifiedCommand{
+		Action:  "error",
+		Message: err.Error(),
+	}
+}
+
+func (v UnifiedCommand) Marshal() []byte {
+	data, _ := jsoniter.Marshal(v)
+	return data
+}
diff --git a/pkg/server/notify_ws.go b/pkg/server/notify_ws.go
deleted file mode 100644
index 2c1615f..0000000
--- a/pkg/server/notify_ws.go
+++ /dev/null
@@ -1,25 +0,0 @@
-package server
-
-import (
-	"git.solsynth.dev/hydrogen/passport/pkg/models"
-	"git.solsynth.dev/hydrogen/passport/pkg/services"
-	"github.com/gofiber/contrib/websocket"
-)
-
-func listenNotifications(c *websocket.Conn) {
-	user := c.Locals("principal").(models.Account)
-
-	// Push connection
-	services.ClientRegister(user, c)
-
-	// Event loop
-	var err error
-	for {
-		if _, _, err = c.ReadMessage(); err != nil {
-			break
-		}
-	}
-
-	// Pop connection
-	services.ClientUnregister(user, c)
-}
diff --git a/pkg/server/startup.go b/pkg/server/startup.go
index 1b9a766..d119d19 100644
--- a/pkg/server/startup.go
+++ b/pkg/server/startup.go
@@ -1,13 +1,13 @@
 package server
 
 import (
+	"github.com/gofiber/contrib/websocket"
 	"net/http"
 	"strings"
 
 	"git.solsynth.dev/hydrogen/passport/pkg"
 	"git.solsynth.dev/hydrogen/passport/pkg/i18n"
 	"git.solsynth.dev/hydrogen/passport/pkg/server/ui"
-	"github.com/gofiber/contrib/websocket"
 	"github.com/gofiber/fiber/v2"
 	"github.com/gofiber/fiber/v2/middleware/cors"
 	"github.com/gofiber/fiber/v2/middleware/favicon"
@@ -74,8 +74,6 @@ func NewServer() {
 			notify.Post("/subscribe", authMiddleware, addNotifySubscriber)
 			notify.Put("/batch/read", authMiddleware, markNotificationReadBatch)
 			notify.Put("/:notificationId/read", authMiddleware, markNotificationRead)
-
-			notify.Get("/listen", authMiddleware, websocket.New(listenNotifications))
 		}
 
 		me := api.Group("/users/me").Name("Myself Operations")
@@ -136,6 +134,8 @@ func NewServer() {
 		{
 			developers.Post("/notify", notifyUser)
 		}
+
+		api.Get("/ws", authMiddleware, websocket.New(listenWebsocket))
 	}
 
 	A.Use(favicon.New(favicon.Config{
diff --git a/pkg/server/ws.go b/pkg/server/ws.go
new file mode 100644
index 0000000..2800136
--- /dev/null
+++ b/pkg/server/ws.go
@@ -0,0 +1,77 @@
+package server
+
+import (
+	"fmt"
+	"git.solsynth.dev/hydrogen/passport/pkg/models"
+	"git.solsynth.dev/hydrogen/passport/pkg/services"
+	"github.com/gofiber/contrib/websocket"
+	jsoniter "github.com/json-iterator/go"
+	"github.com/samber/lo"
+)
+
+func listenWebsocket(c *websocket.Conn) {
+	user := c.Locals("principal").(models.Account)
+
+	// Push connection
+	services.ClientRegister(user, c)
+
+	// Event loop
+	var task models.UnifiedCommand
+
+	var messageType int
+	var payload []byte
+	var packet []byte
+	var err error
+
+	for {
+		if messageType, packet, err = c.ReadMessage(); err != nil {
+			break
+		} else if err := jsoniter.Unmarshal(packet, &task); err != nil {
+			_ = c.WriteMessage(messageType, models.UnifiedCommand{
+				Action:  "error",
+				Message: "unable to unmarshal your command, requires json request",
+			}.Marshal())
+			continue
+		} else {
+			payload, _ = jsoniter.Marshal(task.Payload)
+		}
+
+		var message *models.UnifiedCommand
+		switch task.Action {
+		case "kex.request":
+			var req struct {
+				RequestID string `json:"request_id"`
+				KeypairID string `json:"keypair_id"`
+				OwnerID   uint   `json:"owner_id"`
+				Deadline  int64  `json:"deadline"`
+			}
+			_ = jsoniter.Unmarshal(payload, &req)
+			if len(req.RequestID) <= 0 || len(req.KeypairID) <= 0 || req.OwnerID <= 0 {
+				message = lo.ToPtr(models.UnifiedCommandFromError(fmt.Errorf("invalid request")))
+			}
+			services.KexRequest(c, req.RequestID, req.KeypairID, user.ID, req.Deadline)
+		case "kex.provide":
+			var req struct {
+				RequestID string `json:"request_id"`
+				KeypairID string `json:"keypair_id"`
+				PublicKey []byte `json:"public_key"`
+			}
+			_ = jsoniter.Unmarshal(payload, &req)
+			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)
+		default:
+			message = lo.ToPtr(models.UnifiedCommandFromError(fmt.Errorf("unknown action")))
+		}
+
+		if message != nil {
+			if err = c.WriteMessage(messageType, message.Marshal()); err != nil {
+				break
+			}
+		}
+	}
+
+	// Pop connection
+	services.ClientUnregister(user, c)
+}
diff --git a/pkg/services/e2ee.go b/pkg/services/e2ee.go
new file mode 100644
index 0000000..34aaf40
--- /dev/null
+++ b/pkg/services/e2ee.go
@@ -0,0 +1,77 @@
+package services
+
+import (
+	"git.solsynth.dev/hydrogen/passport/pkg/models"
+	"github.com/gofiber/contrib/websocket"
+	"github.com/gofiber/fiber/v2"
+	"time"
+)
+
+type kexRequest struct {
+	OwnerID  uint
+	Conn     *websocket.Conn
+	Deadline time.Time
+}
+
+var kexRequests = make(map[string]map[string]kexRequest)
+
+func KexRequest(conn *websocket.Conn, requestId, keypairId string, ownerId uint, deadline int64) {
+	if kexRequests[keypairId] == nil {
+		kexRequests[keypairId] = make(map[string]kexRequest)
+	}
+
+	ddl := time.Now().Add(time.Second * time.Duration(deadline))
+	request := kexRequest{
+		OwnerID:  ownerId,
+		Conn:     conn,
+		Deadline: ddl,
+	}
+
+	flag := false
+	for conn := range wsConn[ownerId] {
+		if conn.WriteMessage(1, models.UnifiedCommand{
+			Action: "kex.request",
+			Payload: fiber.Map{
+				"request_id": requestId,
+				"keypair_id": keypairId,
+				"owner_id":   ownerId,
+				"deadline":   deadline,
+			},
+		}.Marshal()) == nil {
+			flag = true
+		}
+	}
+
+	if flag {
+		kexRequests[keypairId][requestId] = request
+	}
+}
+
+func KexProvide(userId uint, requestId, keypairId string, pkt []byte) {
+	if kexRequests[keypairId] == nil {
+		return
+	}
+
+	val, ok := kexRequests[keypairId][requestId]
+	if !ok {
+		return
+	} else if val.OwnerID != userId {
+		return
+	} else {
+		_ = val.Conn.WriteMessage(1, pkt)
+	}
+}
+
+func KexCleanup() {
+	if len(kexRequests) <= 0 {
+		return
+	}
+
+	for kp, data := range kexRequests {
+		for idx, req := range data {
+			if req.Deadline.Unix() <= time.Now().Unix() {
+				delete(kexRequests[kp], idx)
+			}
+		}
+	}
+}
diff --git a/pkg/services/notifications.go b/pkg/services/notifications.go
index 8f9de11..2a7d7d3 100644
--- a/pkg/services/notifications.go
+++ b/pkg/services/notifications.go
@@ -60,7 +60,10 @@ func NewNotification(notification models.Notification) error {
 func PushNotification(notification models.Notification) error {
 	raw, _ := jsoniter.Marshal(notification)
 	for conn := range wsConn[notification.RecipientID] {
-		_ = conn.WriteMessage(1, raw)
+		_ = conn.WriteMessage(1, models.UnifiedCommand{
+			Action:  "notifications.new",
+			Payload: raw,
+		}.Marshal())
 	}
 
 	var subscribers []models.NotificationSubscriber