enhance(rtc,e2ee): Add encryption to the upload-graph process

src/main/frontend/common/crypt.cljs

@@ -1,5 +1,6 @@
 (ns frontend.common.crypt
-  (:require [promesa.core :as p]))
+  (:require [logseq.db :as ldb]
+            [promesa.core :as p]))
 
 (defonce subtle (.. js/crypto -subtle))
 
@@ -118,7 +119,7 @@
                                    #js {:name "AES-GCM" :iv iv}
                                    aes-key
                                    encoded-text)]
-    [iv encrypted-data]))
+    [iv (js/Uint8Array. encrypted-data)]))
 
 (defn <decrypt-text
   "Decrypts text with an AES key."
@@ -133,6 +134,54 @@
           decoded-text (.decode (js/TextDecoder.) decrypted-data)]
     decoded-text))
 
+(defn <decrypt-text-if-encrypted
+  "return nil if not a encrypted-package"
+  [aes-key maybe-encrypted-package]
+  (when (and (vector? maybe-encrypted-package)
+             (<= 2 (count maybe-encrypted-package)))
+    (<decrypt-text aes-key maybe-encrypted-package)))
+
+(defn <encrypt-map
+  [aes-key encrypt-attr-set m]
+  (assert (map? m))
+  (reduce
+   (fn [map-p encrypt-attr]
+     (p/let [m map-p]
+       (if-let [v (get m encrypt-attr)]
+         (p/let [v' (p/chain (<encrypt-text aes-key v) ldb/write-transit-str)]
+           (assoc m encrypt-attr v'))
+         m)))
+   (p/promise m) encrypt-attr-set))
+
+(defn <encrypt-av-coll
+  "see also `rtc-schema/av-schema`"
+  [aes-key encrypt-attr-set av-coll]
+  (p/all
+   (mapv
+    (fn [[a v & others]]
+      (p/let [v' (if (and (contains? encrypt-attr-set a)
+                          (string? v))
+                   (p/chain (<encrypt-text aes-key v) ldb/write-transit-str)
+                   v)]
+        (apply conj [a v'] others)))
+    av-coll)))
+
+(defn <decrypt-map
+  [aes-key encrypt-attr-set m]
+  (assert (map? m))
+  (reduce
+   (fn [map-p encrypt-attr]
+     (p/let [m map-p]
+       (if-let [v (get m encrypt-attr)]
+         (if (string? v)
+           (p/let [v' (<decrypt-text-if-encrypted aes-key (ldb/read-transit-str v))]
+             (if v'
+               (assoc m encrypt-attr v')
+               m))
+           m)
+         m)))
+   (p/promise m) encrypt-attr-set))
+
 (comment
   (let [array-buffers-equal?
         (fn [^js/ArrayBuffer buf1 ^js/ArrayBuffer buf2]

src/main/frontend/handler.cljs

@@ -17,6 +17,7 @@
             [frontend.db.restore :as db-restore]
             [frontend.error :as error]
             [frontend.handler.command-palette :as command-palette]
+            [frontend.handler.crypt]
             [frontend.handler.db-based.vector-search-flows :as vector-search-flows]
             [frontend.handler.events :as events]
             [frontend.handler.events.ui]

src/main/frontend/handler/crypt.cljs

@@ -0,0 +1,6 @@
+(ns frontend.handler.crypt
+  (:require [frontend.common.thread-api :refer [def-thread-api]]))
+
+(def-thread-api :thread-api/request-e2ee-password
+  []
+  {:password "test-password"})

src/main/frontend/worker/rtc/crypt.cljs

@@ -7,6 +7,7 @@
             [frontend.common.crypt :as crypt]
             [frontend.common.missionary :as c.m]
             [frontend.worker.rtc.ws-util :as ws-util]
+            [frontend.worker.state :as worker-state]
             [missionary.core :as m]
             [promesa.core :as p]))
 
@@ -52,45 +53,70 @@
                         (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."
-  [token user-uuid password]
-  (m/sp
-    (let [key-pair (c.m/<? (<get-item (user-rsa-key-pair-idb-key user-uuid)))]
-      (if key-pair
-        (let [private-key (c.m/<? (crypt/<decrypt-private-key password (:encrypted-private-key key-pair)))]
-          {:public-key (:public-key key-pair)
-           :private-key private-key})
-        (let [{:keys [get-ws-create-task]} (ws-util/gen-get-ws-create-map--memoized (ws-util/get-ws-url token))
-              response (m/? (ws-util/send&recv get-ws-create-task
-                                               {:action "fetch-user-rsa-key-pair"
-                                                :user-uuid user-uuid}))]
-          (if (:ex-data response)
-            (throw (ex-info (:ex-message response)
-                            (assoc (:ex-data response)
-                                   :type :rtc.exception/fetch-user-rsa-key-pair-error)))
-            (let [retrieved-key-pair (:body response)]
-              (c.m/<? (<set-item! (user-rsa-key-pair-idb-key user-uuid) retrieved-key-pair))
-              (let [private-key (c.m/<? (crypt/<decrypt-private-key password (:encrypted-private-key retrieved-key-pair)))]
-                {:public-key (:public-key retrieved-key-pair)
-                 :private-key private-key}))))))))
+  "Fetches the user's RSA key pair, from indexeddb or server.
+  Return nil if not exists"
+  [get-ws-create-task user-uuid]
+  (letfn [(select-keys-fn [m] (select-keys m [:public-key :encrypted-private-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 [response (m/? (ws-util/send&recv get-ws-create-task
+                                                 {:action "fetch-user-rsa-key-pair"
+                                                  :user-uuid user-uuid}))]
+            (if (:ex-data response)
+              (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)]
+                (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)))))))))
 
 (defn task--fetch-graph-aes-key
-  "Fetches the AES key for a graph, from indexeddb or server."
-  [token graph-uuid private-key]
+  "Fetches the AES key for a graph, from indexeddb or server.
+  Return nil if not exists"
+  [get-ws-create-task graph-uuid private-key]
   (m/sp
     (let [encrypted-aes-key (c.m/<? (<get-item (graph-encrypted-aes-key-idb-key graph-uuid)))]
       (if encrypted-aes-key
         (c.m/<? (crypt/<decrypt-aes-key private-key encrypted-aes-key))
-        (let [{:keys [get-ws-create-task]} (ws-util/gen-get-ws-create-map--memoized (ws-util/get-ws-url token))
-              response (m/? (ws-util/send&recv get-ws-create-task
+        (let [response (m/? (ws-util/send&recv get-ws-create-task
                                                {:action "fetch-graph-encrypted-aes-key"
                                                 :graph-uuid graph-uuid}))]
           (if (:ex-data response)
             (throw (ex-info (:ex-message response) (assoc (:ex-data response)
                                                           :type :rtc.exception/fetch-graph-aes-key-error)))
-            (let [fetched-encrypted-aes-key (:body response)]
-              (c.m/<? (<set-item! (graph-encrypted-aes-key-idb-key graph-uuid) fetched-encrypted-aes-key))
-              (c.m/<? (crypt/<decrypt-aes-key private-key fetched-encrypted-aes-key)))))))))
+            (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))
+                  aes-key)))))))))
+
+(defn task--persist-graph-encrypted-aes-key
+  [graph-uuid encrypted-aes-key]
+  (m/sp
+    (c.m/<? (<set-item! (graph-encrypted-aes-key-idb-key graph-uuid) encrypted-aes-key))))
+
+(defn task--generate-graph-aes-key
+  []
+  (m/sp (c.m/<? (crypt/<generate-aes-key))))
+
+(defn task--get-user-public-key
+  [get-ws-create-task user-uuid]
+  (m/sp
+   (:public-key (m/? (task--fetch-user-rsa-key-pair get-ws-create-task user-uuid)))))
+
+(defn task--get-rsa-key-pair
+  [get-ws-create-task user-uuid]
+  (m/sp
+    (let [{:keys [password]} (c.m/<? (worker-state/<invoke-main-thread :thread-api/request-e2ee-password))
+          {:keys [public-key encrypted-private-key]}
+          (m/? (task--fetch-user-rsa-key-pair get-ws-create-task user-uuid))
+          private-key (c.m/<? (crypt/<decrypt-private-key password encrypted-private-key))]
+      {:public-key public-key
+       :private-key private-key})))
 
 (comment
   (do

src/main/frontend/worker/rtc/exception.cljs

@@ -29,9 +29,8 @@ the server will put it to s3 and return its presigned-url to clients."}
 
   :rtc.exception/fetch-user-rsa-key-pair-error {:doc "Failed to fetch user RSA key pair from server"}
   :rtc.exception/fetch-graph-aes-key-error {:doc "Failed to fetch graph AES key from server"}
-  :rtc.exception/upload-graph-encrypted-aes-key-error {:doc "Failed to upload graph encrypted AES key to server"}
   :rtc.exception/upload-user-rsa-key-pair-error {:doc "Failed to upload user RSA key pair to server"}
-  )
+  :rtc.exception/not-found-user-rsa-key-pair {:doc "user rsa-key-pair not found"})
 
 (def ex-ws-already-disconnected
   (ex-info "websocket conn is already disconnected" {:type :rtc.exception/ws-already-disconnected}))

src/main/frontend/worker/rtc/full_upload_download_graph.cljs

@@ -4,15 +4,15 @@
   (:require [cljs-http-missionary.client :as http]
             [clojure.set :as set]
             [datascript.core :as d]
+            [frontend.common.crypt :as crypt]
             [frontend.common.missionary :as c.m]
             [frontend.common.thread-api :as thread-api]
             [frontend.worker-common.util :as worker-util]
-            [frontend.worker.crypt :as crypt]
             [frontend.worker.db-metadata :as worker-db-metadata]
             [frontend.worker.rtc.client-op :as client-op]
             [frontend.worker.rtc.const :as rtc-const]
+            [frontend.worker.rtc.crypt :as rtc-crypt]
             [frontend.worker.rtc.db :as rtc-db]
-            [frontend.worker.rtc.encrypt :as rtc-encrypt]
             [frontend.worker.rtc.log-and-state :as rtc-log-and-state]
             [frontend.worker.rtc.ws-util :as ws-util]
             [frontend.worker.shared-service :as shared-service]
@@ -130,7 +130,7 @@
            result []]
       (if-not block
         result
-        (let [block' (c.m/<? (rtc-encrypt/<encrypt-map encrypt-key encrypt-attr-set block))]
+        (let [block' (c.m/<? (crypt/<encrypt-map encrypt-key encrypt-attr-set block))]
           (recur rest-blocks (conj result block')))))))
 
 (comment
@@ -140,56 +140,64 @@
   (def canceler ((m/sp
                    (let [k (c.m/<? (rtc-encrypt/<salt+password->key salt "password"))]
                      (m/? (task--encrypt-blocks k #{:block/title :block/name} blocks))))
-                 #(def encrypted-blocks %) prn))
-  )
+                 #(def encrypted-blocks %) prn)))
 
 (defn new-task--upload-graph
   [get-ws-create-task repo conn remote-graph-name major-schema-version]
   (m/sp
-    (rtc-log-and-state/rtc-log :rtc.log/upload {:sub-type :fetching-presigned-put-url
-                                                :message "fetching presigned put-url"})
-    (let [[{:keys [url key]} all-blocks-str]
-          (m/?
-           (m/join
-            vector
-            (ws-util/send&recv get-ws-create-task {:action "presign-put-temp-s3-obj"})
-            (m/sp
-              (let [all-blocks (export-as-blocks
-                                @conn
-                                :ignore-attr-set rtc-const/ignore-attrs-when-init-upload
-                                :ignore-entity-set rtc-const/ignore-entities-when-init-upload)
-                    encrypt-key (c.m/<? (rtc-encrypt/<get-encrypt-key repo))
-                    _ (assert (some? encrypt-key))
-                    encrypted-blocks (c.m/<? (task--encrypt-blocks encrypt-key rtc-const/encrypt-attr-set all-blocks))]
-                (ldb/write-transit-str encrypted-blocks)))))]
-      (rtc-log-and-state/rtc-log :rtc.log/upload {:sub-type :upload-data
-                                                  :message "uploading data"})
-      (m/? (http/put url {:body all-blocks-str :with-credentials? false}))
-      (rtc-log-and-state/rtc-log :rtc.log/upload {:sub-type :request-upload-graph
-                                                  :message "requesting upload-graph"})
-      (let [aes-key (c.m/<? (crypt/<gen-aes-key))
-            aes-key-jwk (ldb/write-transit-str (c.m/<? (crypt/<export-key aes-key)))
-            upload-resp
-            (m/? (ws-util/send&recv get-ws-create-task {:action "upload-graph"
-                                                        :s3-key key
-                                                        :schema-version (str major-schema-version)
-                                                        :graph-name remote-graph-name}))]
-        (if-let [graph-uuid (:graph-uuid upload-resp)]
-          (let [schema-version (ldb/get-graph-schema-version @conn)]
-            (ldb/transact! conn
-                           [(ldb/kv :logseq.kv/graph-uuid graph-uuid)
-                            (ldb/kv :logseq.kv/graph-local-tx "0")
-                            (ldb/kv :logseq.kv/remote-schema-version schema-version)])
-            (client-op/update-graph-uuid repo graph-uuid)
-            (client-op/remove-local-tx repo)
-            (client-op/update-local-tx repo 1)
-            (client-op/add-all-exists-asset-as-ops repo)
-            (crypt/store-graph-keys-jwk repo aes-key-jwk)
-            (c.m/<? (worker-db-metadata/<store repo (pr-str {:kv/value graph-uuid})))
-            (rtc-log-and-state/rtc-log :rtc.log/upload {:sub-type :upload-completed
-                                                        :message "upload-graph completed"})
-            {:graph-uuid graph-uuid})
-          (throw (ex-info "upload-graph failed" {:upload-resp upload-resp})))))))
+    (rtc-log-and-state/rtc-log :rtc.log/upload {:sub-type :generate-aes-key
+                                                :message "generate aes-encrypt-key"})
+    (let [aes-key (m/? (rtc-crypt/task--generate-graph-aes-key))
+          user-uuid (some-> (worker-state/get-id-token)
+                            worker-util/parse-jwt
+                            :sub)
+          public-key (when user-uuid
+                       (m/? (rtc-crypt/task--get-user-public-key get-ws-create-task user-uuid)))]
+      (when-not public-key
+        (throw (ex-info "user public-key not found" {:type :rtc.exception/not-found-user-rsa-key-pair
+                                                     :user-uuid user-uuid})))
+
+      (let [encrypted-aes-key (c.m/<? (crypt/<encrypt-aes-key public-key aes-key))
+            _ (rtc-log-and-state/rtc-log :rtc.log/upload {:sub-type :fetching-presigned-put-url
+                                                          :message "fetching presigned put-url"})
+            [{:keys [url key]} all-blocks-str]
+            (m/?
+             (m/join
+              vector
+              (ws-util/send&recv get-ws-create-task {:action "presign-put-temp-s3-obj"})
+              (m/sp
+                (let [all-blocks (export-as-blocks
+                                  @conn
+                                  :ignore-attr-set rtc-const/ignore-attrs-when-init-upload
+                                  :ignore-entity-set rtc-const/ignore-entities-when-init-upload)
+                      encrypted-blocks (c.m/<? (task--encrypt-blocks aes-key rtc-const/encrypt-attr-set all-blocks))]
+                  (ldb/write-transit-str encrypted-blocks)))))]
+        (rtc-log-and-state/rtc-log :rtc.log/upload {:sub-type :upload-data
+                                                    :message "uploading data"})
+        (m/? (http/put url {:body all-blocks-str :with-credentials? false}))
+        (rtc-log-and-state/rtc-log :rtc.log/upload {:sub-type :request-upload-graph
+                                                    :message "requesting upload-graph"})
+        (let [upload-resp
+              (m/? (ws-util/send&recv get-ws-create-task {:action "upload-graph"
+                                                          :s3-key key
+                                                          :schema-version (str major-schema-version)
+                                                          :graph-name remote-graph-name}))]
+          (if-let [graph-uuid (:graph-uuid upload-resp)]
+            (let [schema-version (ldb/get-graph-schema-version @conn)]
+              (ldb/transact! conn
+                             [(ldb/kv :logseq.kv/graph-uuid graph-uuid)
+                              (ldb/kv :logseq.kv/graph-local-tx "0")
+                              (ldb/kv :logseq.kv/remote-schema-version schema-version)])
+              (client-op/update-graph-uuid repo graph-uuid)
+              (client-op/remove-local-tx repo)
+              (client-op/update-local-tx repo 1)
+              (client-op/add-all-exists-asset-as-ops repo)
+              (c.m/<? (worker-db-metadata/<store repo (pr-str {:kv/value graph-uuid})))
+              (m/? (rtc-crypt/task--persist-graph-encrypted-aes-key graph-uuid encrypted-aes-key))
+              (rtc-log-and-state/rtc-log :rtc.log/upload {:sub-type :upload-completed
+                                                          :message "upload-graph completed"})
+              {:graph-uuid graph-uuid})
+            (throw (ex-info "upload-graph failed" {:upload-resp upload-resp}))))))))
 
 (defn- fill-block-fields
   [blocks]
@@ -515,9 +523,7 @@
       (m/? (http/put url {:body all-blocks-str :with-credentials? false}))
       (rtc-log-and-state/rtc-log :rtc.log/branch-graph {:sub-type :request-branch-graph
                                                         :message "requesting branch-graph"})
-      (let [aes-key (c.m/<? (crypt/<gen-aes-key))
-            aes-key-jwk (ldb/write-transit-str (c.m/<? (crypt/<export-key aes-key)))
-            resp (m/? (ws-util/send&recv get-ws-create-task {:action "branch-graph"
+      (let [resp (m/? (ws-util/send&recv get-ws-create-task {:action "branch-graph"
                                                              :s3-key key
                                                              :schema-version (str major-schema-version)
                                                              :graph-uuid graph-uuid}))]
@@ -530,7 +536,6 @@
             (client-op/update-graph-uuid repo graph-uuid)
             (client-op/remove-local-tx repo)
             (client-op/add-all-exists-asset-as-ops repo)
-            (crypt/store-graph-keys-jwk repo aes-key-jwk)
             (c.m/<? (worker-db-metadata/<store repo (pr-str {:kv/value graph-uuid})))
             (rtc-log-and-state/rtc-log :rtc.log/branch-graph {:sub-type :completed
                                                               :message "branch-graph completed"})

src/main/frontend/worker/rtc/malli_schema.cljs

@@ -243,6 +243,20 @@
                       [:graph<->user/user-type :keyword]
                       [:user/online? :boolean]]]]]]
      ["inject-users-info" [:map]]
+
+     ;; keys manage
+     ["fetch-user-rsa-key-pair"
+      [:map
+       [:public-key [:maybe :string]]
+       [:encrypted-private-key [:maybe :string]]]]
+     ["fetch-graph-encrypted-aes-key"
+      [:map
+       [:encrypted-aes-key [:maybe :string]]]]
+     ["upload-user-rsa-key-pair"
+      [:map
+       [:public-key :string]
+       [:encrypted-private-key :string]]]
+
      [nil data-from-ws-schema-fallback]]))
 
 (def data-from-ws-coercer (m/coercer data-from-ws-schema mt/string-transformer nil
@@ -349,6 +363,8 @@
         [:graph-uuid :uuid]
         [:schema-version db-schema/major-schema-version-string-schema]
         [:asset-uuids [:sequential :uuid]]]]
+      ;; ================================================================
+      ;; TODO: cleanup
       ["get-user-devices"
        [:map]]
       ["add-user-device"
@@ -373,6 +389,18 @@
       ["sync-encrypted-aes-key"
        [:map
         [:device-uuid->encrypted-aes-key [:map-of :uuid :string]]
+        [:graph-uuid :uuid]]]
+      ;; ================================================================
+      ["upload-user-rsa-key-pair"
+       [:map
+        [:user-uuid :uuid]
+        [:public-key :string]
+        [:encrypted-private-key :string]]]
+      ["fetch-user-rsa-key-pair"
+       [:map
+        [:user-uuid :uuid]]]
+      ["fetch-graph-encrypted-aes-key"
+       [:map
         [:graph-uuid :uuid]]]])))
 
 (def data-to-ws-encoder (m/encoder data-to-ws-schema (mt/transformer