Correctly return the next ID of (signed, kyber and normal) pre-keys (#…

…254)

Co-authored-by: boxdot <d@zerovolt.org>

Cargo.toml

@@ -3,7 +3,7 @@ members = ["presage", "presage-cli", "presage-store-sled", "presage-store-cipher
 resolver = "2"
 
 [patch.crates-io]
-curve25519-dalek = { git = 'https://github.com/signalapp/curve25519-dalek', tag = 'signal-curve25519-4.1.1' }
+curve25519-dalek = { git = 'https://github.com/signalapp/curve25519-dalek', tag = 'signal-curve25519-4.1.3' }
 
 # [patch."https://github.com/whisperfish/libsignal-service-rs.git"]
 # libsignal-service = { path = "../libsignal-service-rs/libsignal-service" }

presage-cli/src/main.rs

@@ -16,6 +16,7 @@ use notify_rust::Notification;
 use presage::libsignal_service::configuration::SignalServers;
 use presage::libsignal_service::content::Reaction;
 use presage::libsignal_service::models::Contact;
+use presage::libsignal_service::pre_keys::PreKeysStore;
 use presage::libsignal_service::prelude::phonenumber::PhoneNumber;
 use presage::libsignal_service::prelude::Uuid;
 use presage::libsignal_service::proto::data_message::Quote;
@@ -164,6 +165,8 @@ enum Cmd {
         master_key: GroupMasterKeyBytes,
     },
     RequestContactsSync,
+    #[clap(about = "Print various statistics useful for debugging")]
+    Stats,
 }
 
 enum Recipient {
@@ -731,6 +734,54 @@ async fn run<S: Store>(subcommand: Cmd, config_store: S) -> anyhow::Result<()> {
                 print_message(&manager, false, &msg);
             }
         }
+        Cmd::Stats => {
+            let manager = Manager::load_registered(config_store).await?;
+
+            #[allow(unused)]
+            #[derive(Debug)]
+            struct Stats {
+                aci_next_pre_key_id: u32,
+                aci_next_signed_pre_keys_id: u32,
+                aci_next_kyber_pre_keys_id: u32,
+                aci_signed_pre_keys_count: usize,
+                aci_kyber_pre_keys_count: usize,
+                aci_kyber_pre_keys_count_last_resort: usize,
+                pni_next_pre_key_id: u32,
+                pni_next_signed_pre_keys_id: u32,
+                pni_next_kyber_pre_keys_id: u32,
+                pni_signed_pre_keys_count: usize,
+                pni_kyber_pre_keys_count: usize,
+                pni_kyber_pre_keys_count_last_resort: usize,
+            }
+
+            let aci = manager.store().aci_protocol_store();
+            let pni = manager.store().pni_protocol_store();
+
+            const LAST_RESORT: bool = true;
+
+            let stats = Stats {
+                aci_next_pre_key_id: aci.next_pre_key_id().await.unwrap(),
+                aci_next_signed_pre_keys_id: aci.next_signed_pre_key_id().await.unwrap(),
+                aci_next_kyber_pre_keys_id: aci.next_pq_pre_key_id().await.unwrap(),
+                aci_signed_pre_keys_count: aci.signed_pre_keys_count().await.unwrap(),
+                aci_kyber_pre_keys_count: aci.kyber_pre_keys_count(!LAST_RESORT).await.unwrap(),
+                aci_kyber_pre_keys_count_last_resort: aci
+                    .kyber_pre_keys_count(LAST_RESORT)
+                    .await
+                    .unwrap(),
+                pni_next_pre_key_id: pni.next_pre_key_id().await.unwrap(),
+                pni_next_signed_pre_keys_id: pni.next_signed_pre_key_id().await.unwrap(),
+                pni_next_kyber_pre_keys_id: pni.next_pq_pre_key_id().await.unwrap(),
+                pni_signed_pre_keys_count: pni.signed_pre_keys_count().await.unwrap(),
+                pni_kyber_pre_keys_count: pni.kyber_pre_keys_count(!LAST_RESORT).await.unwrap(),
+                pni_kyber_pre_keys_count_last_resort: pni
+                    .kyber_pre_keys_count(LAST_RESORT)
+                    .await
+                    .unwrap(),
+            };
+
+            println!("{stats:#?}")
+        }
     }
     Ok(())
 }

presage-store-sled/src/lib.rs

@@ -33,9 +33,6 @@ use sled::IVec;
 
 const SLED_TREE_STATE: &str = "state";
 
-const SLED_KEY_NEXT_SIGNED_PRE_KEY_ID: &str = "next_signed_pre_key_id";
-const SLED_KEY_NEXT_PQ_PRE_KEY_ID: &str = "next_pq_pre_key_id";
-const SLED_KEY_PRE_KEYS_OFFSET_ID: &str = "pre_keys_offset_id";
 const SLED_KEY_REGISTRATION: &str = "registration";
 const SLED_KEY_SCHEMA_VERSION: &str = "schema_version";
 #[cfg(feature = "encryption")]
@@ -77,11 +74,13 @@ pub enum SchemaVersion {
     V4 = 4,
     /// ACI and PNI identity key pairs are moved into dedicated storage keys from registration data
     V5 = 5,
+    /// Reset pre-keys after fixing persistence
+    V6 = 6,
 }
 
 impl SchemaVersion {
     fn current() -> SchemaVersion {
-        Self::V5
+        Self::V6
     }
 
     /// return an iterator on all the necessary migration steps from another version
@@ -96,6 +95,7 @@ impl SchemaVersion {
             3 => SchemaVersion::V3,
             4 => SchemaVersion::V4,
             5 => SchemaVersion::V5,
+            6 => SchemaVersion::V6,
             _ => unreachable!("oops, this not supposed to happen!"),
         })
     }
@@ -392,6 +392,41 @@ fn migrate(
                         log::error!("failed to run v4 -> v5 migration: {error}");
                     }
                 }
+                SchemaVersion::V6 => {
+                    debug!("migrating from schema v5 to v6: new keys encoding in ACI and PNI protocol stores");
+                    let db = store.db.read().expect("poisoned");
+
+                    let trees = [
+                        AciSledStore::signed_pre_keys(),
+                        AciSledStore::pre_keys(),
+                        AciSledStore::kyber_pre_keys(),
+                        AciSledStore::kyber_pre_keys_last_resort(),
+                        PniSledStore::signed_pre_keys(),
+                        PniSledStore::pre_keys(),
+                        PniSledStore::kyber_pre_keys(),
+                        PniSledStore::kyber_pre_keys_last_resort(),
+                    ];
+
+                    for tree_name in trees {
+                        let tree = db.open_tree(tree_name)?;
+                        let num_keys_before = tree.len();
+                        let mut data = Vec::new();
+                        for (k, v) in tree.iter().filter_map(|kv| kv.ok()) {
+                            if let Some(key) = std::str::from_utf8(&k)
+                                .ok()
+                                .and_then(|s| s.parse::<u32>().ok())
+                            {
+                                data.push((key, v));
+                            }
+                        }
+                        tree.clear()?;
+                        for (k, v) in data {
+                            let _ = tree.insert(k.to_be_bytes(), v);
+                        }
+                        let num_keys_after = tree.len();
+                        debug!("migrated keys in {tree_name}: before {num_keys_before} -> after {num_keys_after}");
+                    }
+                }
                 _ => return Err(SledStoreError::MigrationConflict),
             }
 
@@ -469,8 +504,8 @@ impl StateStore for SledStore {
         self.clear_profiles()?;
 
         // drop all keys
-        self.aci_protocol_store().clear()?;
-        self.pni_protocol_store().clear()?;
+        self.aci_protocol_store().clear(true)?;
+        self.pni_protocol_store().clear(true)?;
 
         Ok(())
     }
@@ -503,13 +538,33 @@ mod tests {
         content::{ContentBody, Metadata},
         prelude::Uuid,
         proto::DataMessage,
+        protocol::PreKeyId,
         ServiceAddress,
     };
     use presage::store::ContentsStore;
+    use protocol::SledPreKeyId;
     use quickcheck::{Arbitrary, Gen};
+    use quickcheck_macros::quickcheck;
+
+    use crate::SchemaVersion;
 
-    use super::SledStore;
+    use super::*;
 
+    #[test]
+    fn test_migration_steps() {
+        let steps: Vec<_> = SchemaVersion::steps(SchemaVersion::V0).collect();
+        assert_eq!(
+            steps,
+            [
+                SchemaVersion::V1,
+                SchemaVersion::V2,
+                SchemaVersion::V3,
+                SchemaVersion::V4,
+                SchemaVersion::V5,
+                SchemaVersion::V6,
+            ]
+        )
+    }
     #[derive(Debug, Clone)]
     struct Thread(presage::store::Thread);
 
@@ -567,6 +622,14 @@ mod tests {
         }
     }
 
+    #[quickcheck]
+    fn compare_pre_keys(mut pre_key_id: u32, mut next_pre_key_id: u32) {
+        if pre_key_id > next_pre_key_id {
+            std::mem::swap(&mut pre_key_id, &mut next_pre_key_id);
+        }
+        assert!(PreKeyId::from(pre_key_id).sled_key() <= PreKeyId::from(next_pre_key_id).sled_key())
+    }
+
     #[quickcheck_async::tokio]
     async fn test_store_messages(thread: Thread, content: Content) -> anyhow::Result<()> {
         let db = SledStore::temporary()?;