✨ E2EE Key Exchange

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

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
+}

pkg/server/notify_ws.go

@@ -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)
-}

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{

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)
+}

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)
+			}
+		}
+	}
+}

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