enhance(rtc,e2ee): support en/decrypt on assets

src/main/frontend/common/crypt.cljs

@@ -5,13 +5,31 @@
 
 (defonce subtle (.. js/crypto -subtle))
 
+(defn <export-aes-key
+  [aes-key]
+  (assert (instance? js/CryptoKey aes-key))
+  (p/let [exported (.exportKey subtle "raw" aes-key)]
+    (js/Uint8Array. exported)))
+
+(defn <import-aes-key
+  [exported-aes-key]
+  (assert (instance? js/Uint8Array exported-aes-key))
+  (.importKey subtle
+              "raw"
+              exported-aes-key
+              "AES-GCM"
+              true
+              #js ["encrypt" "decrypt"]))
+
 (defn <export-public-key
   [public-key]
+  (assert (instance? js/CryptoKey public-key))
   (p/let [exported (.exportKey subtle "spki" public-key)]
     (js/Uint8Array. exported)))
 
 (defn <import-public-key
   [exported-public-key]
+  (assert (instance? js/Uint8Array exported-public-key))
   (.importKey subtle "spki" exported-public-key
               #js {:name "RSA-OAEP" :hash "SHA-256"}
               true
@@ -54,7 +72,7 @@
 (defn <encrypt-private-key
   "Encrypts a private key with a password."
   [password private-key]
-  (assert (string? password))
+  (assert (and (string? password) (instance? js/CryptoKey private-key)))
   (p/let [salt (js/crypto.getRandomValues (js/Uint8Array. 16))
           iv (js/crypto.getRandomValues (js/Uint8Array. 12))
           password-key (.importKey subtle "raw"
@@ -118,7 +136,9 @@
 (defn <encrypt-aes-key
   "Encrypts an AES key with a public key."
   [public-key aes-key]
-  (p/let [exported-aes-key (.exportKey subtle "raw" aes-key)
+  (assert (and (instance? js/CryptoKey public-key)
+               (instance? js/CryptoKey aes-key)))
+  (p/let [exported-aes-key (<export-aes-key aes-key)
           encrypted-aes-key (.encrypt subtle
                                       #js {:name "RSA-OAEP"}
                                       public-key
@@ -146,9 +166,36 @@
               (log/error "decrypt-aes-key" e)
               (ex-info "decrypt-aes-key" {} e)))))
 
+(defn <encrypt-uint8array
+  [aes-key arr]
+  (assert (and (instance? js/CryptoKey aes-key) (instance? js/Uint8Array arr)))
+  (p/let [iv (js/crypto.getRandomValues (js/Uint8Array. 12))
+          encrypted-data (.encrypt subtle
+                                   #js {:name "AES-GCM" :iv iv}
+                                   aes-key
+                                   arr)]
+    [iv (js/Uint8Array. encrypted-data)]))
+
+(defn <decrypt-uint8array
+  [aes-key encrypted-data-vector]
+  (->
+   (p/let [[iv-data encrypted-data] encrypted-data-vector
+           _ (assert (instance? js/Uint8Array encrypted-data))
+           iv (js/Uint8Array. iv-data)
+           decrypted-data (.decrypt subtle
+                                    #js {:name "AES-GCM" :iv iv}
+                                    aes-key
+                                    encrypted-data)]
+     (js/Uint8Array. decrypted-data))
+   (p/catch
+    (fn [e]
+      (log/error "decrypt-uint8array" e)
+      (ex-info "decrypt-uint8array" {} e)))))
+
 (defn <encrypt-text
   "Encrypts text with an AES key."
   [aes-key text]
+  (assert (and (string? text) (instance? js/CryptoKey aes-key)))
   (p/let [iv (js/crypto.getRandomValues (js/Uint8Array. 12))
           encoded-text (.encode (js/TextEncoder.) text)
           encrypted-data (.encrypt subtle
@@ -159,16 +206,20 @@
 
 (defn <decrypt-text
   "Decrypts text with an AES key."
-  [aes-key encrypted-text-data]
-  (p/let [[iv-data encrypted-data-from-db] encrypted-text-data
-          iv (js/Uint8Array. iv-data)
-          encrypted-data (js/Uint8Array. encrypted-data-from-db)
-          decrypted-data (.decrypt subtle
-                                   #js {:name "AES-GCM" :iv iv}
-                                   aes-key
-                                   encrypted-data)
-          decoded-text (.decode (js/TextDecoder.) decrypted-data)]
-    decoded-text))
+  [aes-key encrypted-text-data-vector]
+  (-> (p/let [[iv-data encrypted-data] encrypted-text-data-vector
+              iv (js/Uint8Array. iv-data)
+              encrypted-data (js/Uint8Array. encrypted-data)
+              decrypted-data (.decrypt subtle
+                                       #js {:name "AES-GCM" :iv iv}
+                                       aes-key
+                                       encrypted-data)
+              decoded-text (.decode (js/TextDecoder.) decrypted-data)]
+        decoded-text)
+      (p/catch
+       (fn [e]
+         (log/error "decrypt-text" e)
+         (ex-info "decrypt-text" {} e)))))
 
 (defn <decrypt-text-if-encrypted
   "return nil if not a encrypted-package"

src/main/frontend/handler/assets.cljs

@@ -1,6 +1,7 @@
 (ns ^:no-doc frontend.handler.assets
   (:require [cljs-http-missionary.client :as http]
             [clojure.string :as string]
+            [frontend.common.crypt :as crypt]
             [frontend.common.missionary :as c.m]
             [frontend.common.thread-api :as thread-api :refer [def-thread-api]]
             [frontend.config :as config]
@@ -10,6 +11,7 @@
             [logseq.common.config :as common-config]
             [logseq.common.path :as path]
             [logseq.common.util :as common-util]
+            [logseq.db :as ldb]
             [logseq.db.frontend.asset :as db-asset]
             [medley.core :as medley]
             [missionary.core :as m]
@@ -243,6 +245,18 @@
       :assets/asset-file-write-finish
       (fn [m] (assoc-in m [repo asset-block-id-str] (common-util/time-ms)))))))
 
+(comment
+  ;; en/decrypt assets
+  (def repo (state/get-current-repo))
+  (p/let [aes-key (crypt/<generate-aes-key)
+          asset (<read-asset repo "6903201e-9573-4914-ae88-7d3f1d095d1f" "png")
+          encrypted-asset (crypt/<encrypt-uint8array aes-key asset)
+          decrypted-asset (crypt/<decrypt-uint8array aes-key encrypted-asset)]
+    (def asset asset)
+    (def xxxx encrypted-asset)
+    (prn :decrypted (.-length decrypted-asset)
+         :origin (.-length asset))))
+
 (defn <unlink-asset
   [repo asset-block-id asset-type]
   (let [file-path (path/path-join (config/get-repo-dir repo)
@@ -251,14 +265,18 @@
     (p/catch (fs/unlink! repo file-path {}) (constantly nil))))
 
 (defn new-task--rtc-upload-asset
-  [repo asset-block-uuid-str asset-type checksum put-url]
+  [repo aes-key asset-block-uuid-str asset-type checksum put-url]
   (assert (and asset-type checksum))
   (m/sp
     (let [asset-file (c.m/<? (<read-asset repo asset-block-uuid-str asset-type))
+          asset-file* (if (not aes-key)
+                        asset-file
+                        (ldb/write-transit-str
+                         (c.m/<? (crypt/<encrypt-uint8array aes-key asset-file))))
           *progress-flow (atom nil)
           http-task (http/put put-url {:headers {"x-amz-meta-checksum" checksum
                                                  "x-amz-meta-type" asset-type}
-                                       :body asset-file
+                                       :body asset-file*
                                        :with-credentials? false
                                        :*progress-flow *progress-flow})]
       (c.m/run-task :upload-asset-progress
@@ -273,7 +291,7 @@
           {:ex-data {:type :rtc.exception/upload-asset-failed :data (dissoc r :body)}})))))
 
 (defn new-task--rtc-download-asset
-  [repo asset-block-uuid-str asset-type get-url]
+  [repo aes-key asset-block-uuid-str asset-type get-url]
   (m/sp
     (let [*progress-flow (atom nil)
           http-task (http/get get-url {:with-credentials? false
@@ -291,8 +309,22 @@
         (let [{:keys [status body] :as r} (m/? http-task)]
           (if-not (http/unexceptional-status? status)
             {:ex-data {:type :rtc.exception/download-asset-failed :data (dissoc r :body)}}
-            (do (c.m/<? (<write-asset repo asset-block-uuid-str asset-type body))
-                nil)))
+            (let [asset-file
+                  (if (not aes-key)
+                    body
+                    (try
+                      (let [asset-file-untransited (ldb/read-transit-str (.decode (js/TextDecoder.) body))]
+                        (c.m/<? (crypt/<decrypt-uint8array aes-key asset-file-untransited)))
+                      (catch js/SyntaxError _
+                        body)
+                      (catch :default e
+                        ;; if decrypt failed, write origin-body
+                        (if (= "decrypt-uint8array" (ex-message e))
+                          body
+                          (throw e)))))]
+              (c.m/<? (<write-asset repo asset-block-uuid-str asset-type asset-file))
+              nil)))
+
         (catch Cancelled e
           (progress-canceler)
           (throw e))))))
@@ -310,12 +342,16 @@
   (<get-asset-file-metadata repo asset-block-id asset-type))
 
 (def-thread-api :thread-api/rtc-upload-asset
-  [repo asset-block-uuid-str asset-type checksum put-url]
-  (new-task--rtc-upload-asset repo asset-block-uuid-str asset-type checksum put-url))
+  [repo exported-aes-key asset-block-uuid-str asset-type checksum put-url]
+  (m/sp
+    (let [aes-key (when exported-aes-key (c.m/<? (crypt/<import-aes-key exported-aes-key)))]
+      (m/? (new-task--rtc-upload-asset repo aes-key asset-block-uuid-str asset-type checksum put-url)))))
 
 (def-thread-api :thread-api/rtc-download-asset
-  [repo asset-block-uuid-str asset-type get-url]
-  (new-task--rtc-download-asset repo asset-block-uuid-str asset-type get-url))
+  [repo exported-aes-key asset-block-uuid-str asset-type get-url]
+  (m/sp
+    (let [aes-key (when exported-aes-key (c.m/<? (crypt/<import-aes-key exported-aes-key)))]
+      (m/? (new-task--rtc-download-asset repo aes-key asset-block-uuid-str asset-type get-url)))))
 
 (comment
   ;; read asset

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

@@ -7,6 +7,7 @@
     indicates need to upload the asset to server"
   (:require [clojure.set :as set]
             [datascript.core :as d]
+            [frontend.common.crypt :as crypt]
             [frontend.common.missionary :as c.m]
             [frontend.worker.rtc.client-op :as client-op]
             [frontend.worker.rtc.exception :as r.ex]
@@ -119,44 +120,52 @@
 
 (defn- new-task--concurrent-download-assets
   "Concurrently download assets with limited max concurrent count"
-  [repo asset-uuid->url asset-uuid->asset-type]
-  (->> (fn [[asset-uuid url]]
-         (m/sp
-           (let [r (c.m/<?
-                    (worker-state/<invoke-main-thread :thread-api/rtc-download-asset
-                                                      repo (str asset-uuid)
-                                                      (get asset-uuid->asset-type asset-uuid) url))]
-             (when-let [edata (:ex-data r)]
-               ;; if download-url return 404, ignore this asset
-               (when (not= 404 (:status (:data edata)))
-                 (throw (ex-info "download asset failed" r)))))))
-       (c.m/concurrent-exec-flow 5 (m/seed asset-uuid->url))
-       (m/reduce (constantly nil))))
+  [repo aes-key asset-uuid->url asset-uuid->asset-type]
+  (m/sp
+    (let [exported-aes-key (when aes-key (c.m/<? (crypt/<export-aes-key aes-key)))]
+      (m/?
+       (->> (fn [[asset-uuid url]]
+              (m/sp
+                (let [r (c.m/<?
+                         (worker-state/<invoke-main-thread :thread-api/rtc-download-asset
+                                                           repo exported-aes-key (str asset-uuid)
+                                                           (get asset-uuid->asset-type asset-uuid) url))]
+                  (when-let [edata (:ex-data r)]
+                    ;; if download-url return 404, ignore this asset
+                    (when (not= 404 (:status (:data edata)))
+                      (throw (ex-info "download asset failed" r)))))))
+            (c.m/concurrent-exec-flow 5 (m/seed asset-uuid->url))
+            (m/reduce (constantly nil)))))))
 
 (defn- new-task--concurrent-upload-assets
   "Concurrently upload assets with limited max concurrent count"
-  [repo conn asset-uuid->url asset-uuid->asset-metadata]
-  (->> (fn [[asset-uuid url]]
-         (m/sp
-           (let [[asset-type checksum] (get asset-uuid->asset-metadata asset-uuid)
-                 r (c.m/<?
-                    (worker-state/<invoke-main-thread :thread-api/rtc-upload-asset
-                                                      repo (str asset-uuid) asset-type checksum url))]
-             (when (:ex-data r)
-               (throw (ex-info "upload asset failed" r)))
-             ;; asset might be deleted by the user before uploaded successfully
-             (when (d/entity @conn [:block/uuid asset-uuid])
-               (ldb/transact! conn
-                              [{:block/uuid asset-uuid
-                                :logseq.property.asset/remote-metadata {:checksum checksum :type asset-type}}]
-                            ;; Don't generate rtc ops again, (block-ops & asset-ops)
-                              {:persist-op? false}))
-             (client-op/remove-asset-op repo asset-uuid))))
-       (c.m/concurrent-exec-flow 3 (m/seed asset-uuid->url))
-       (m/reduce (constantly nil))))
+  [repo conn aes-key asset-uuid->url asset-uuid->asset-metadata]
+  (m/sp
+    (let [exported-aes-key (when aes-key (c.m/<? (crypt/<export-aes-key aes-key)))]
+      (m/?
+       (->> (fn [[asset-uuid url]]
+              (m/sp
+                (let [[asset-type checksum] (get asset-uuid->asset-metadata asset-uuid)
+                      _ (prn :xxx exported-aes-key)
+                      r (c.m/<?
+                         (worker-state/<invoke-main-thread :thread-api/rtc-upload-asset
+                                                           repo exported-aes-key (str asset-uuid)
+                                                           asset-type checksum url))]
+                  (when (:ex-data r)
+                    (throw (ex-info "upload asset failed" r)))
+                  ;; asset might be deleted by the user before uploaded successfully
+                  (when (d/entity @conn [:block/uuid asset-uuid])
+                    (ldb/transact! conn
+                                   [{:block/uuid asset-uuid
+                                     :logseq.property.asset/remote-metadata {:checksum checksum :type asset-type}}]
+                                   ;; Don't generate rtc ops again, (block-ops & asset-ops)
+                                   {:persist-op? false}))
+                  (client-op/remove-asset-op repo asset-uuid))))
+            (c.m/concurrent-exec-flow 3 (m/seed asset-uuid->url))
+            (m/reduce (constantly nil)))))))
 
 (defn- new-task--push-local-asset-updates
-  [repo get-ws-create-task conn graph-uuid major-schema-version add-log-fn]
+  [repo get-ws-create-task conn graph-uuid major-schema-version aes-key add-log-fn]
   (m/sp
     (when-let [asset-ops (not-empty (client-op/get-all-asset-ops repo))]
       (let [upload-asset-uuids (keep
@@ -198,7 +207,7 @@
                    :asset-uuid->url))]
         (when (seq asset-uuid->url)
           (add-log-fn :rtc.asset.log/upload-assets {:asset-uuids (keys asset-uuid->url)}))
-        (m/? (new-task--concurrent-upload-assets repo conn asset-uuid->url asset-uuid->asset-metadata))
+        (m/? (new-task--concurrent-upload-assets repo conn aes-key asset-uuid->url asset-uuid->asset-metadata))
         (when (seq remove-asset-uuids)
           (add-log-fn :rtc.asset.log/remove-assets {:asset-uuids remove-asset-uuids})
           (m/? (ws-util/send&recv get-ws-create-task
@@ -213,7 +222,7 @@
                           (concat (keys asset-uuid->url) remove-asset-uuids))))))
 
 (defn- new-task--pull-remote-asset-updates
-  [repo get-ws-create-task conn graph-uuid add-log-fn asset-update-ops]
+  [repo get-ws-create-task conn graph-uuid aes-key add-log-fn asset-update-ops]
   (m/sp
     (when (seq asset-update-ops)
       (let [update-asset-uuids (keep (fn [op]
@@ -252,7 +261,7 @@
                                                     repo (str asset-uuid) asset-type)))
         (when (seq asset-uuid->url)
           (add-log-fn :rtc.asset.log/download-assets {:asset-uuids (keys asset-uuid->url)}))
-        (m/? (new-task--concurrent-download-assets repo asset-uuid->url asset-uuid->asset-type))))))
+        (m/? (new-task--concurrent-download-assets repo aes-key asset-uuid->url asset-uuid->asset-type))))))
 
 (defn- get-all-asset-blocks
   [db]
@@ -267,7 +276,7 @@
        db))
 
 (defn- new-task--initial-download-missing-assets
-  [repo get-ws-create-task graph-uuid conn add-log-fn]
+  [repo get-ws-create-task graph-uuid conn aes-key add-log-fn]
   (m/sp
     (let [local-all-asset-file-paths
           (c.m/<? (worker-state/<invoke-main-thread :thread-api/get-all-asset-file-paths repo))
@@ -279,10 +288,10 @@
                        (set/difference local-all-asset-uuids local-all-asset-file-uuids)))]
         (add-log-fn :rtc.asset.log/initial-download-missing-assets {:count (count asset-update-ops)})
         (m/? (new-task--pull-remote-asset-updates
-              repo get-ws-create-task conn graph-uuid add-log-fn asset-update-ops))))))
+              repo get-ws-create-task conn graph-uuid aes-key add-log-fn asset-update-ops))))))
 
 (defn create-assets-sync-loop
-  [repo get-ws-create-task graph-uuid major-schema-version conn *auto-push?]
+  [repo get-ws-create-task graph-uuid major-schema-version conn *auto-push? *aes-key]
   (let [started-dfv (m/dfv)
         add-log-fn (fn [type message]
                      (assert (map? message) message)
@@ -295,18 +304,21 @@
       (m/sp
         (try
           (log/info :rtc-asset :loop-starting)
+          ;; check aes-key exists
+          (when (ldb/get-graph-rtc-e2ee? @conn) (assert @*aes-key))
           (started-dfv true)
-          (m/? (new-task--initial-download-missing-assets repo get-ws-create-task graph-uuid conn add-log-fn))
+          (m/? (new-task--initial-download-missing-assets
+                repo get-ws-create-task graph-uuid conn @*aes-key add-log-fn))
           (->>
            (let [event (m/?> mixed-flow)]
              (case (:type event)
                :remote-updates
                (when-let [asset-update-ops (not-empty (:value event))]
                  (m/? (new-task--pull-remote-asset-updates
-                       repo get-ws-create-task conn graph-uuid add-log-fn asset-update-ops)))
+                       repo get-ws-create-task conn graph-uuid @*aes-key add-log-fn asset-update-ops)))
                :local-update-check
                (m/? (new-task--push-local-asset-updates
-                     repo get-ws-create-task conn graph-uuid major-schema-version add-log-fn))))
+                     repo get-ws-create-task conn graph-uuid major-schema-version @*aes-key add-log-fn))))
            m/ap
            (m/reduce {} nil)
            m/?)

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

@@ -225,7 +225,8 @@
                             get-ws-create-task0 graph-uuid major-schema-version repo conn date-formatter
                             *last-calibrate-t *online-users *server-schema-version *aes-key add-log-fn)
         {:keys [assets-sync-loop-task]}
-        (r.asset/create-assets-sync-loop repo get-ws-create-task graph-uuid major-schema-version conn *auto-push?)
+        (r.asset/create-assets-sync-loop
+         repo get-ws-create-task graph-uuid major-schema-version conn *auto-push? *aes-key)
         mixed-flow                 (create-mixed-flow repo get-ws-create-task *auto-push? *online-users)]
     (assert (some? *current-ws))
     {:rtc-state-flow       (create-rtc-state-flow (create-ws-state-flow *current-ws))
@@ -369,7 +370,7 @@
                          :fail (fn [e]
                                  (reset! *last-stop-exception e)
                                  (log/info :rtc-loop-task e)
-                                 (when-not (instance? Cancelled e)
+                                 (when-not (or (instance? Cancelled e) (= "missionary.Cancelled" (ex-message e)))
                                    (println (.-stack e)))
                                  (when (= :rtc.exception/ws-timeout (some-> e ex-data :type))
                                    ;; if fail reason is websocket-timeout, try to restart rtc