Change how dynamic captcha stores data

src/internet_identity/src/storage/registration_rates.rs

@@ -9,10 +9,10 @@ use std::time::Duration;
 
 pub struct RegistrationRates<M: Memory> {
     /// Min heap of registration timestamps to calculate the reference registration rate.
-    /// [prune_data] removes values that  are in the past whenever new data is added.
+    /// [prune_data] removes values that are older than now minus the reference rate interval whenever new data is added.
     reference_rate_data: MinHeap<Timestamp, M>,
     /// Min heap of registration timestamps to calculate the current registration rate.
-    /// [prune_data] removes values that  are in the past whenever new data is added.
+    /// [prune_data] removes values that are older than now minus the current rate interval whenever new data is added.
     current_rate_data: MinHeap<Timestamp, M>,
 }
 
@@ -47,16 +47,17 @@ impl<M: Memory> RegistrationRates<M> {
     }
     pub fn new_registration(&mut self) {
         self.prune_expired();
-        let Some(data_retention) = dynamic_captcha_config() else {
+        // Return if no dynamic captcha
+        if dynamic_captcha_config().is_none() {
             return;
         };
 
         let now = time();
         self.reference_rate_data
-            .push(&(now + data_retention.reference_rate_retention_ns))
+            .push(&now)
             .expect("out of memory");
         self.current_rate_data
-            .push(&(now + data_retention.current_rate_retention_ns))
+            .push(&now)
             .expect("out of memory");
     }
 
@@ -67,12 +68,10 @@ impl<M: Memory> RegistrationRates<M> {
         let reference_rate_per_second = calculate_registration_rate(
             now,
             &self.reference_rate_data,
-            config.reference_rate_retention_ns,
         );
         let current_rate_per_second = calculate_registration_rate(
             now,
             &self.current_rate_data,
-            config.current_rate_retention_ns,
         );
         let captcha_threshold_rate = reference_rate_per_second * config.threshold_multiplier;
         let rates = NormalizedRegistrationRates {
@@ -84,8 +83,13 @@ impl<M: Memory> RegistrationRates<M> {
     }
 
     fn prune_expired(&mut self) {
-        prune_data(&mut self.reference_rate_data);
-        prune_data(&mut self.current_rate_data);
+        let Some(data_retention) = dynamic_captcha_config() else {
+            return;
+        };
+        println!("in da prun_expired");
+        println!("reference_rate {}", data_retention.reference_rate_retention_ns.to_string());
+        prune_data(&mut self.reference_rate_data, data_retention.reference_rate_retention_ns);
+        prune_data(&mut self.current_rate_data, data_retention.current_rate_retention_ns);
     }
 }
 
@@ -111,17 +115,13 @@ impl<M: Memory> RegistrationRates<M> {
 fn calculate_registration_rate<M: Memory>(
     now: u64,
     data: &MinHeap<Timestamp, M>,
-    data_retention_ns: u64,
 ) -> f64 {
     data
-        // get the oldest expiration timestamp
+        // get the oldest value
         .peek()
-        // calculate the registration timestamp from expiration
-        .map(|ts| ts - data_retention_ns)
         // calculate the time window length with respect to the current time
         .map(|ts| now - ts)
-        // the value _could_ be 0 if the oldest timestamp was added in the same execution round
-        // -> filter to avoid division by 0
+        // There could be negative values if they haven't been pruned yet.
         .filter(|val| *val != 0)
         // use the value to calculate the rate per second
         .map(|val| rate_per_second(data.len(), val))
@@ -151,7 +151,7 @@ fn dynamic_captcha_config() -> Option<DynamicCaptchaConfig> {
     }
 }
 
-fn prune_data<M: Memory>(data: &mut MinHeap<Timestamp, M>) {
+fn prune_data<M: Memory>(data: &mut MinHeap<Timestamp, M>, interval_ns: u64) {
     const MAX_TO_PRUNE: usize = 100;
 
     let now = time();
@@ -162,7 +162,7 @@ fn prune_data<M: Memory>(data: &mut MinHeap<Timestamp, M>) {
 
         // The timestamps are sorted because the expiration is constant and time() is monotonic
         // -> we can stop pruning once we reach a not expired timestamp
-        if timestamp > now {
+        if timestamp > (now - interval_ns) {
             break;
         }
         data.pop();
@@ -281,6 +281,69 @@ mod test {
         );
     }
 
+    #[test]
+    fn previous_config_should_not_affect_new_config() {
+        let mut registration_rates = setup();
+        // initialize reference rate with 1 registration / second
+        for _ in 0..1000 {
+            registration_rates.new_registration();
+            TIME.with_borrow_mut(|t| *t += Duration::from_secs(1).as_nanos() as u64);
+        }
+
+        // raise current rate to 3 registrations / second
+        for _ in 0..100 {
+            registration_rates.new_registration();
+            registration_rates.new_registration();
+            registration_rates.new_registration();
+            TIME.with_borrow_mut(|t| *t += Duration::from_secs(1).as_nanos() as u64);
+        }
+
+        assert_eq!(
+            registration_rates.registration_rates().unwrap(),
+            NormalizedRegistrationRates {
+                current_rate_per_second: 3.0,
+                reference_rate_per_second: 1.2, // increases too, but more slowly
+                captcha_threshold_rate: 1.44,   // reference rate * 1.2 (as per config)
+            }
+        );
+
+        // Change the config.
+        state::persistent_state_mut(|ps| {
+            ps.captcha_config = CaptchaConfig {
+                max_unsolved_captchas: 500,
+                captcha_trigger: CaptchaTrigger::Dynamic {
+                    threshold_pct: 10,
+                    current_rate_sampling_interval_s: 100,
+                    reference_rate_sampling_interval_s: 100,
+                },
+            }
+        });
+
+        // Needed to prune the data after new config
+        // It prunes 100 for each new registration
+        // we added 1100 before. Remove them with 110, plus 2 to remove the new 110 added.
+        for _ in 0..112 {
+            registration_rates.new_registration();
+            TIME.with_borrow_mut(|t| *t += Duration::from_secs(1).as_nanos() as u64);
+        }
+
+        // Set current rate to 2 registrations per seconds for 100 seconds
+        // which is the current and reference rate.
+        for _ in 0..100 {
+            registration_rates.new_registration();
+            registration_rates.new_registration();
+            TIME.with_borrow_mut(|t| *t += Duration::from_secs(1).as_nanos() as u64);
+        }
+        assert_eq!(
+            registration_rates.registration_rates().unwrap(),
+            NormalizedRegistrationRates {
+                current_rate_per_second: 2.0,
+                reference_rate_per_second: 2.0, // increases too, but more slowly
+                captcha_threshold_rate: 2.2,   // reference rate * 1.1 (as per config)
+            }
+        );
+    }
+
     #[test]
     fn should_prune_old_data() {
         let mut registration_rates = setup();