1 files changed, 61 insertions, 0 deletions
diff --git a/third_party/rust/sync15/src/client/collection_keys.rs b/third_party/rust/sync15/src/client/collection_keys.rs
new file mode 100644
index 0000000000..f51894f756
--- /dev/null
+++ b/third_party/rust/sync15/src/client/collection_keys.rs
@@ -0,0 +1,61 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+use crate::error::Result;
+use crate::record_types::CryptoKeysRecord;
+use crate::{EncryptedPayload, KeyBundle, ServerTimestamp};
+use std::collections::HashMap;
+
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct CollectionKeys {
+    pub timestamp: ServerTimestamp,
+    pub default: KeyBundle,
+    pub collections: HashMap<String, KeyBundle>,
+}
+
+impl CollectionKeys {
+    pub fn new_random() -> Result<CollectionKeys> {
+        let default = KeyBundle::new_random()?;
+        Ok(CollectionKeys {
+            timestamp: ServerTimestamp(0),
+            default,
+            collections: HashMap::new(),
+        })
+    }
+
+    pub fn from_encrypted_payload(
+        record: EncryptedPayload,
+        timestamp: ServerTimestamp,
+        root_key: &KeyBundle,
+    ) -> Result<CollectionKeys> {
+        let keys: CryptoKeysRecord = record.decrypt_into(root_key)?;
+        Ok(CollectionKeys {
+            timestamp,
+            default: KeyBundle::from_base64(&keys.default[0], &keys.default[1])?,
+            collections: keys
+                .collections
+                .into_iter()
+                .map(|kv| Ok((kv.0, KeyBundle::from_base64(&kv.1[0], &kv.1[1])?)))
+                .collect::<Result<HashMap<String, KeyBundle>>>()?,
+        })
+    }
+
+    pub fn to_encrypted_payload(&self, root_key: &KeyBundle) -> Result<EncryptedPayload> {
+        let record = CryptoKeysRecord {
+            id: "keys".into(),
+            collection: "crypto".into(),
+            default: self.default.to_b64_array(),
+            collections: self
+                .collections
+                .iter()
+                .map(|kv| (kv.0.clone(), kv.1.to_b64_array()))
+                .collect(),
+        };
+        EncryptedPayload::from_cleartext_payload(root_key, &record)
+    }
+
+    pub fn key_for_collection<'a>(&'a self, collection: &str) -> &'a KeyBundle {
+        self.collections.get(collection).unwrap_or(&self.default)
+    }
+}