From 059a7d3df06e03260427df56863190dfade7193b Mon Sep 17 00:00:00 2001
From: rcmerci <rcmerci@gmail.com>
Date: Mon, 27 Oct 2025 23:42:55 +0800
Subject: [PATCH] enhance(rtc,e2ee): init/reset user rsa-pair

---
 src/main/frontend/common/crypt.cljs           |  2 +
 src/main/frontend/db/rtc/debug_ui.cljs        | 12 ++++-
 src/main/frontend/worker/rtc/crypt.cljs       | 54 ++++++++++++++-----
 .../rtc/full_upload_download_graph.cljs       |  3 +-
 .../frontend/worker/rtc/malli_schema.cljs     |  3 +-
 5 files changed, 58 insertions(+), 16 deletions(-)

diff --git a/src/main/frontend/common/crypt.cljs b/src/main/frontend/common/crypt.cljs
index bbc19ab460..935a4966de 100644
--- a/src/main/frontend/common/crypt.cljs
+++ b/src/main/frontend/common/crypt.cljs
@@ -53,6 +53,7 @@
 (defn <encrypt-private-key
   "Encrypts a private key with a password."
   [password private-key]
+  (assert (string? password))
   (p/let [salt (js/crypto.getRandomValues (js/Uint8Array. 16))
           iv (js/crypto.getRandomValues (js/Uint8Array. 12))
           password-key (.importKey subtle "raw"
@@ -79,6 +80,7 @@
 (defn <decrypt-private-key
   "Decrypts a private key with a password."
   [password encrypted-key-data]
+  (assert (and (vector? encrypted-key-data) (= 3 (count encrypted-key-data))))
   (p/let [[salt-data iv-data encrypted-private-key-data] encrypted-key-data
           salt (js/Uint8Array. salt-data)
           iv (js/Uint8Array. iv-data)
diff --git a/src/main/frontend/db/rtc/debug_ui.cljs b/src/main/frontend/db/rtc/debug_ui.cljs
index 7a2a97f8e0..18d76a1a4b 100644
--- a/src/main/frontend/db/rtc/debug_ui.cljs
+++ b/src/main/frontend/db/rtc/debug_ui.cljs
@@ -278,8 +278,16 @@
                        (when-let [token (state/get-auth-id-token)]
                          (p/let [r (state/<invoke-db-worker :thread-api/init-user-rsa-key-pair token (user/user-uuid))]
                            (when (instance? ExceptionInfo r)
-                             (log/error :thread-api/init-user-rsa-key-pair r)))))}
-          (shui/tabler-icon "upload") "init upload user rsa-key-pair")]
+                             (log/error :init-user-rsa-key-pair r)))))}
+          (shui/tabler-icon "upload") "init upload user rsa-key-pair")
+         (shui/button
+          {:size :sm
+           :on-click (fn [_]
+                       (when-let [token (state/get-auth-id-token)]
+                         (p/let [r (state/<invoke-db-worker :thread-api/force-reset-user-rsa-key-pair token (user/user-uuid))]
+                           (when (instance? ExceptionInfo r)
+                             (log/error :force-reset-user-rsa-key-pair r)))))}
+          (shui/tabler-icon "upload") "force-reset upload user rsa-key-pair")]
         [:div.pb-4
          [:pre.select-text
           (-> keys-state
diff --git a/src/main/frontend/worker/rtc/crypt.cljs b/src/main/frontend/worker/rtc/crypt.cljs
index 105a846f1f..18d3bec5b2 100644
--- a/src/main/frontend/worker/rtc/crypt.cljs
+++ b/src/main/frontend/worker/rtc/crypt.cljs
@@ -43,7 +43,8 @@
 
 (defn task--upload-user-rsa-key-pair
   "Uploads the user's RSA key pair to the server."
-  [get-ws-create-task user-uuid public-key encrypted-private-key]
+  [get-ws-create-task user-uuid public-key encrypted-private-key & {:keys [force-reset]
+                                                                    :or {force-reset false}}]
   (m/sp
     (let [exported-public-key-str (ldb/write-transit-str (c.m/<? (crypt/<export-public-key public-key)))
           encrypted-private-key-str (ldb/write-transit-str encrypted-private-key)
@@ -51,20 +52,27 @@
                                            {:action "upload-user-rsa-key-pair"
                                             :user-uuid user-uuid
                                             :public-key exported-public-key-str
-                                            :encrypted-private-key encrypted-private-key-str}))]
+                                            :encrypted-private-key encrypted-private-key-str
+                                            :force-reset force-reset}))]
       (when (:ex-data response)
         (throw (ex-info (:ex-message response)
                         (assoc (:ex-data response) :type :rtc.exception/upload-user-rsa-key-pair-error)))))))
 
 (defn task--fetch-user-rsa-key-pair
   "Fetches the user's RSA key pair, from indexeddb or server.
+  Return {:public-key CryptoKey, :encrypted-private-key [array,array,array]}
   Return nil if not exists"
   [get-ws-create-task user-uuid]
-  (letfn [(select-keys-fn [m] (select-keys m [:public-key :encrypted-private-key]))]
+  (letfn [(task--import-public-key [public-key-transit-str]
+            (m/sp
+              (when-let [exported-public-key (ldb/read-transit-str public-key-transit-str)]
+                (c.m/<? (crypt/<import-public-key exported-public-key)))))]
     (m/sp
       (let [key-pair (c.m/<? (<get-item (user-rsa-key-pair-idb-key user-uuid)))]
         (if key-pair
-          (select-keys-fn key-pair)
+          (let [public-key (m/? (task--import-public-key (:public-key key-pair)))]
+            {:public-key public-key
+             :encrypted-private-key (ldb/read-transit-str (:encrypted-private-key key-pair))})
           (let [response (m/? (ws-util/send&recv get-ws-create-task
                                                  {:action "fetch-user-rsa-key-pair"
                                                   :user-uuid user-uuid}))]
@@ -72,11 +80,14 @@
               (throw (ex-info (:ex-message response)
                               (assoc (:ex-data response)
                                      :type :rtc.exception/fetch-user-rsa-key-pair-error)))
-              (let [{:keys [public-key encrypted-private-key] :as key-pair} (select-keys-fn response)]
+              (let [{:keys [public-key encrypted-private-key]} response]
                 (when (and public-key encrypted-private-key)
                   (c.m/<? (<set-item! (user-rsa-key-pair-idb-key user-uuid)
-                                      (clj->js key-pair)))
-                  key-pair)))))))))
+                                      (clj->js
+                                       {:public-key public-key
+                                        :encrypted-private-key encrypted-private-key})))
+                  {:public-key (m/? (task--import-public-key public-key))
+                   :encrypted-private-key (ldb/read-transit-str encrypted-private-key)})))))))))
 
 (defn task--fetch-graph-aes-key
   "Fetches the AES key for a graph, from indexeddb or server.
@@ -93,9 +104,9 @@
             (throw (ex-info (:ex-message response) (assoc (:ex-data response)
                                                           :type :rtc.exception/fetch-graph-aes-key-error)))
             (let [{:keys [encrypted-aes-key]} response]
-              (when encrypted-aes-key
-                (let [aes-key (c.m/<? (crypt/<decrypt-aes-key private-key encrypted-aes-key))]
-                  (c.m/<? (<set-item! (graph-encrypted-aes-key-idb-key graph-uuid) encrypted-aes-key))
+              (when-let [encrypted-aes-key* (ldb/read-transit-str encrypted-aes-key)]
+                (let [aes-key (c.m/<? (crypt/<decrypt-aes-key private-key encrypted-aes-key*))]
+                  (c.m/<? (<set-item! (graph-encrypted-aes-key-idb-key graph-uuid) encrypted-aes-key*))
                   aes-key)))))))))
 
 (defn task--persist-graph-encrypted-aes-key
@@ -134,11 +145,30 @@
       (let [{:keys [get-ws-create-task]} (ws-util/gen-get-ws-create-map--memoized (ws-util/get-ws-url token))]
         (when-not (m/? (task--fetch-user-rsa-key-pair get-ws-create-task user-uuid))
           (let [{:keys [publicKey privateKey]} (c.m/<? (crypt/<generate-rsa-key-pair))
-                password (c.m/<? (worker-state/<invoke-main-thread :thread-api/request-e2ee-password))
+                {:keys [password]} (c.m/<? (worker-state/<invoke-main-thread :thread-api/request-e2ee-password))
                 encrypted-private-key (c.m/<? (crypt/<encrypt-private-key password privateKey))]
             (m/? (task--upload-user-rsa-key-pair get-ws-create-task user-uuid publicKey encrypted-private-key))
             ;; fetch again
-            (m/? (task--fetch-user-rsa-key-pair get-ws-create-task user-uuid)))))
+            (c.m/<? (<remove-item! (user-rsa-key-pair-idb-key user-uuid)))
+            (m/? (task--fetch-user-rsa-key-pair get-ws-create-task user-uuid))
+            (c.m/<? (<get-item (user-rsa-key-pair-idb-key user-uuid))))))
+      (catch Cancelled _)
+      (catch :default e e))))
+
+(def-thread-api :thread-api/force-reset-user-rsa-key-pair
+  [token user-uuid]
+  (m/sp
+    (try
+      (let [{:keys [get-ws-create-task]} (ws-util/gen-get-ws-create-map--memoized (ws-util/get-ws-url token))
+            {:keys [publicKey privateKey]} (c.m/<? (crypt/<generate-rsa-key-pair))
+            {:keys [password]} (c.m/<? (worker-state/<invoke-main-thread :thread-api/request-e2ee-password))
+            encrypted-private-key (c.m/<? (crypt/<encrypt-private-key password privateKey))]
+        (m/? (task--upload-user-rsa-key-pair get-ws-create-task user-uuid publicKey encrypted-private-key
+                                             :force-reset true))
+        ;; fetch again
+        (c.m/<? (<remove-item! (user-rsa-key-pair-idb-key user-uuid)))
+        (m/? (task--fetch-user-rsa-key-pair get-ws-create-task user-uuid))
+        (c.m/<? (<get-item (user-rsa-key-pair-idb-key user-uuid))))
       (catch Cancelled _)
       (catch :default e e))))
 
diff --git a/src/main/frontend/worker/rtc/full_upload_download_graph.cljs b/src/main/frontend/worker/rtc/full_upload_download_graph.cljs
index ede71cb97d..f4aaf8172f 100644
--- a/src/main/frontend/worker/rtc/full_upload_download_graph.cljs
+++ b/src/main/frontend/worker/rtc/full_upload_download_graph.cljs
@@ -182,7 +182,8 @@
                                                           :s3-key key
                                                           :schema-version (str major-schema-version)
                                                           :graph-name remote-graph-name
-                                                          :encrypted-aes-key encrypted-aes-key}))]
+                                                          :encrypted-aes-key
+                                                          (ldb/write-transit-str encrypted-aes-key)}))]
           (if-let [graph-uuid (:graph-uuid upload-resp)]
             (let [schema-version (ldb/get-graph-schema-version @conn)]
               (ldb/transact! conn
diff --git a/src/main/frontend/worker/rtc/malli_schema.cljs b/src/main/frontend/worker/rtc/malli_schema.cljs
index c24c505cba..4ca99e2d0a 100644
--- a/src/main/frontend/worker/rtc/malli_schema.cljs
+++ b/src/main/frontend/worker/rtc/malli_schema.cljs
@@ -396,7 +396,8 @@
        [:map
         [:user-uuid :uuid]
         [:public-key :string]
-        [:encrypted-private-key :string]]]
+        [:encrypted-private-key :string]
+        [:force-reset {:optional true} :boolean]]]
       ["fetch-user-rsa-key-pair"
        [:map
         [:user-uuid :uuid]]]