keytrans: Detect version changes sooner

RELEASE_NOTES.md

@@ -1,2 +1,4 @@
 v0.93.3
 
+- keytrans: Detect version changes sooner
+

rust/keytrans/src/lib.rs

@@ -312,6 +312,20 @@ pub struct MonitoringData {
     pub owned: bool,
     /// Search key
     pub search_key: Vec<u8>,
+    /// Greatest counter observed in any proof step received for this search
+    /// key.
+    ///
+    /// It includes counter values only present in frontier nodes that
+    /// are _not_ used in the monitoring path (only covers the ancestor nodes)
+    /// Tracked separately from `ptrs` to keep the monitoring algorithm
+    /// unmodified and following the spec precisely.
+    ///
+    /// Max observed version is only used for the early version change
+    /// detection in `libsignal_net_chat::api::keytrans::monitor_and_search`.
+    /// Without it the version change will only be detected by monitor after
+    /// the tree root has moved to a node that includes the update, which
+    /// requires the log to grow by a lot.
+    pub max_observed_version: u32,
 }
 
 impl Debug for MonitoringData {
@@ -322,6 +336,7 @@ impl Debug for MonitoringData {
             ptrs,
             owned,
             search_key,
+            max_observed_version,
         } = self;
 
         let redact_bytes = |bytes: &[u8]| {
@@ -339,6 +354,7 @@ impl Debug for MonitoringData {
             .field("ptrs", &ptrs)
             .field("owned", &owned)
             .field("search_key", &redact_bytes(search_key))
+            .field("max_observed_version", &max_observed_version)
             .finish()
     }
 }
@@ -367,22 +383,33 @@ impl MonitoringData {
 
     /// The greatest known version of the search key.
     pub fn greatest_version(&self) -> u32 {
-        self.ptrs
+        *self
+            .ptrs
             .values()
+            .chain([&self.max_observed_version])
             .max()
-            .copied()
             .expect("at least one version must be present")
     }
 }
 
 impl MonitoringData {
     fn into_stored(self, search_key: Vec<u8>) -> StoredMonitoringData {
+        let Self {
+            index,
+            pos,
+            ptrs,
+            owned,
+            // Prefer the search key provided as argument.
+            search_key: _,
+            max_observed_version,
+        } = self;
         StoredMonitoringData {
-            index: self.index.into(),
-            pos: self.pos,
-            ptrs: self.ptrs,
-            owned: self.owned,
+            index: index.into(),
+            pos,
+            ptrs,
+            owned,
             search_key,
+            max_observed_version,
         }
     }
 }
@@ -395,6 +422,7 @@ impl From<StoredMonitoringData> for MonitoringData {
             ptrs,
             owned,
             search_key,
+            max_observed_version,
         } = value;
         Self {
             index: index.try_into().expect("must be the right size"),
@@ -402,6 +430,7 @@ impl From<StoredMonitoringData> for MonitoringData {
             ptrs,
             owned,
             search_key,
+            max_observed_version,
         }
     }
 }

rust/keytrans/src/proto/store.proto

@@ -21,6 +21,7 @@ message StoredMonitoringData {
   map<uint64, uint32> ptrs = 3;
   bool owned = 4;
   bytes search_key = 5;
+  uint32 max_observed_version = 6;
 }
 
 message StoredAccountData {

rust/keytrans/src/verify.rs

@@ -807,6 +807,7 @@ impl MonitoringDataWrapper {
                 ptrs: HashMap::from([(ver_pos, version)]),
                 owned,
                 search_key: vec![],
+                max_observed_version: version,
             });
         }
     }
@@ -860,6 +861,10 @@ impl MonitoringDataWrapper {
             }
         }
 
+        if version > data.max_observed_version {
+            data.max_observed_version = version;
+        }
+
         if !data.owned && owned {
             data.owned = true;
         }
@@ -898,6 +903,15 @@ impl MonitoringDataWrapper {
 
         data.ptrs = ptrs;
 
+        // Keep the max_observed_version up to date and tracking the maximum
+        // version across _all_ nodes, not just the ancestor ones used for
+        // `ptrs`.
+        if let Some(max_version) = tree_mapping.versions().flatten().max()
+            && max_version > data.max_observed_version
+        {
+            data.max_observed_version = max_version;
+        }
+
         Ok(())
     }
 
@@ -964,6 +978,14 @@ impl VersionExtractor<'_> {
             .map(|step| Ok(get_proto_field(&step.prefix, "prefix")?.counter))
             .transpose()
     }
+
+    pub fn versions(&self) -> impl Iterator<Item = Option<u32>> {
+        self.0.values().map(|step| {
+            get_proto_field(&step.prefix, "prefix")
+                .ok()
+                .map(|x| x.counter)
+        })
+    }
 }
 
 fn into_sorted_pairs<K: Ord + Copy, V>(map: HashMap<K, V>) -> (Vec<K>, Vec<V>) {
@@ -1125,6 +1147,7 @@ mod test {
             ptrs: HashMap::from_iter([(16777215, 2)]),
             owned: true,
             search_key: vec![],
+            max_observed_version: 2,
         }));
         // These values were obtained by running the integration test in
         // rust/net/chat/src/api/keytrans.rs and extracting positions and versions
@@ -1205,6 +1228,7 @@ mod test {
             ptrs: HashMap::from([(10, 1)]),
             owned: true,
             search_key: vec![],
+            max_observed_version: 1,
         }));
 
         let steps = proof_steps([(11, 1), (15, 2)]);
@@ -1223,6 +1247,7 @@ mod test {
             ptrs: HashMap::from([(10, 1)]),
             owned: true,
             search_key: vec![],
+            max_observed_version: 1,
         }));
         // later position contains a smaller version
         let steps = proof_steps([(11, 0)]);
@@ -1239,6 +1264,7 @@ mod test {
             ptrs: HashMap::from([(10, 1)]),
             owned: true,
             search_key: vec![],
+            max_observed_version: 1,
         }));
 
         let steps = HashMap::from_iter([
@@ -1259,6 +1285,7 @@ mod test {
             ptrs: HashMap::from([(10, 1), (11, 2)]),
             owned: true,
             search_key: vec![],
+            max_observed_version: 2,
         }));
         let steps = proof_steps([(11, 3)]);
         let result = wrapper.update(16, &steps);
@@ -1330,4 +1357,48 @@ mod test {
         let result = extractor.get(1);
         assert_matches!(result, Err(Error::RequiredFieldMissing(s)) => assert!(s.contains("prefix")));
     }
+
+    // A counter increment that's only visible at a frontier proof (e.g. a
+    // recent tombstone for a contact's old E.164) must still be reflected in
+    // `MonitoringData::greatest_version`, so that the version-change detector
+    // in `monitor_then_search` triggers a follow-up search.
+    //
+    // For a tree of size 20 with `first_pos = 0`:
+    // ```text
+    //   root(0, 20)                = 15
+    //   monitoring_path(10, 0, 20) = [11, 15]      (ancestors > 10)
+    //   full_monitoring_path       = [11, 15, 19]  (adds frontier node 19)
+    // ```
+    // A tombstone at log position 17 is not yet in the prefix tree at sizes
+    // 12 or 16, so the ancestor proofs at positions 11 and 15 report the
+    // stored counter (0). It _is_ in the prefix tree at size 20, so the
+    // frontier proof at position 19 reports the updated counter (1). The
+    // ancestor-only walk in `find_updated_mapping` never visits the frontier
+    // node, so `ptrs` does not change. The greatest known version of
+    // the search key, however, should reflect the higher counter so that
+    // version change detection still works.
+    #[test]
+    fn frontier_only_version_increment_is_detected() {
+        let mut wrapper = MonitoringDataWrapper::new(Some(MonitoringData {
+            index: [0; 32],
+            pos: 0,
+            ptrs: HashMap::from([(10, 0)]),
+            owned: false,
+            search_key: vec![],
+            max_observed_version: 0,
+        }));
+
+        let steps = proof_steps([
+            // Ancestor nodes
+            (11, 0),
+            (15, 0),
+            // Frontier node with updated version counter
+            (19, 1),
+        ]);
+
+        wrapper.update(20, &steps).expect("valid test data");
+        let data = wrapper.inner.expect("valid test data");
+
+        assert_eq!(1, data.greatest_version(),);
+    }
 }