From 723894b9eaa271e82006316f237e8978663a451f Mon Sep 17 00:00:00 2001
From: Nick Mathewson <nickm@torproject.org>
Date: Thu, 6 Jan 2022 15:11:38 -0500
Subject: [PATCH] Tests for new guardmgr functionality.

---
 crates/tor-guardmgr/src/guard.rs  | 51 +++++++++++++++++++++++++++++++
 crates/tor-guardmgr/src/sample.rs | 32 +++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/crates/tor-guardmgr/src/guard.rs b/crates/tor-guardmgr/src/guard.rs
index a95161c0c..bfafd8382 100644
--- a/crates/tor-guardmgr/src/guard.rs
+++ b/crates/tor-guardmgr/src/guard.rs
@@ -763,6 +763,33 @@ mod test {
         assert!(g.conforms_to_usage(&usage6));
     }
 
+    #[allow(clippy::redundant_clone)]
+    #[test]
+    fn trickier_usages() {
+        let g = basic_guard();
+        use crate::{GuardUsageBuilder, GuardUsageKind};
+        let data_usage = GuardUsageBuilder::new()
+            .kind(GuardUsageKind::Data)
+            .build()
+            .unwrap();
+        let dir_usage = GuardUsageBuilder::new()
+            .kind(GuardUsageKind::OneHopDirectory)
+            .build()
+            .unwrap();
+        assert!(g.conforms_to_usage(&data_usage));
+        assert!(g.conforms_to_usage(&dir_usage));
+
+        let mut g2 = g.clone();
+        g2.microdescriptor_missing = true;
+        assert!(!g2.conforms_to_usage(&data_usage));
+        assert!(g2.conforms_to_usage(&dir_usage));
+
+        let mut g3 = g.clone();
+        g3.is_dir_cache = false;
+        assert!(g3.conforms_to_usage(&data_usage));
+        assert!(!g3.conforms_to_usage(&dir_usage));
+    }
+
     #[test]
     fn retry_interval_check() {
         const MIN: Duration = Duration::from_secs(60);
@@ -953,6 +980,18 @@ mod test {
         .unwrap()
         .unwrap_if_sufficient()
         .unwrap();
+        // Same as above but omit [22] as well as MD for [23].
+        let netdir3 = testnet::construct_custom_netdir(|idx, mut node| {
+            if idx == 22 {
+                node.omit_rs = true;
+            } else if idx == 23 {
+                node.omit_md = true;
+            }
+        })
+        .unwrap()
+        .unwrap_if_sufficient()
+        .unwrap();
+
         //let params = GuardParams::default();
         let now = SystemTime::now();
 
@@ -963,6 +1002,7 @@ mod test {
             now,
         );
         assert_eq!(guard255.unlisted_since, None);
+        assert_eq!(guard255.listed_in(&netdir), Some(false));
         guard255.update_from_netdir(&netdir);
         assert_eq!(
             guard255.unlisted_since,
@@ -973,15 +1013,26 @@ mod test {
         // Try a guard that is in netdir, but not netdir2.
         let mut guard22 = Guard::new(GuardId::new([22; 32].into(), [22; 20].into()), vec![], now);
         let relay22 = guard22.id.get_relay(&netdir).unwrap();
+        assert_eq!(guard22.listed_in(&netdir), Some(true));
         guard22.update_from_netdir(&netdir);
         assert_eq!(guard22.unlisted_since, None); // It's listed.
         assert_eq!(&guard22.orports, relay22.addrs()); // Addrs are set.
+        assert_eq!(guard22.listed_in(&netdir2), Some(false));
         guard22.update_from_netdir(&netdir2);
         assert_eq!(
             guard22.unlisted_since,
             Some(netdir2.lifetime().valid_after())
         );
         assert_eq!(&guard22.orports, relay22.addrs()); // Addrs still set.
+        assert!(!guard22.microdescriptor_missing);
+
+        // Now see what happens for a guard that's in the consensus, but missing an MD.
+        let mut guard23 = Guard::new(GuardId::new([23; 32].into(), [23; 20].into()), vec![], now);
+        assert_eq!(guard23.listed_in(&netdir2), Some(true));
+        assert_eq!(guard23.listed_in(&netdir3), None);
+        guard23.update_from_netdir(&netdir3);
+        assert!(guard23.microdescriptor_missing);
+        assert!(guard23.is_dir_cache);
     }
 
     #[test]
diff --git a/crates/tor-guardmgr/src/sample.rs b/crates/tor-guardmgr/src/sample.rs
index 1b24b7db4..76134d8eb 100644
--- a/crates/tor-guardmgr/src/sample.rs
+++ b/crates/tor-guardmgr/src/sample.rs
@@ -1138,4 +1138,36 @@ mod test {
         assert_eq!(kind, ListKind::Primary);
         assert_eq!(p_id3, p_id1);
     }
+
+    #[test]
+    fn count_missing_mds() {
+        let netdir = netdir();
+        let params = GuardParams {
+            min_filtered_sample_size: 5,
+            n_primary: 2,
+            max_sample_bw_fraction: 1.0,
+            ..GuardParams::default()
+        };
+        let usage = crate::GuardUsageBuilder::default().build().unwrap();
+        let mut guards = GuardSet::new();
+        guards.extend_sample_as_needed(SystemTime::now(), &params, &netdir);
+        guards.select_primary_guards(&params);
+        assert_eq!(guards.primary.len(), 2);
+
+        let (_kind, p_id1) = guards.pick_guard(&usage, &params).unwrap();
+        guards.record_success(&p_id1, &params, SystemTime::now());
+        assert_eq!(guards.missing_primary_microdescriptors(&netdir), 0);
+
+        use tor_netdir::testnet;
+        let netdir2 = testnet::construct_custom_netdir(|idx, bld| {
+            if idx == p_id1.ed25519.as_bytes()[0] as usize {
+                bld.omit_md = true;
+            }
+        })
+        .unwrap()
+        .unwrap_if_sufficient()
+        .unwrap();
+
+        assert_eq!(guards.missing_primary_microdescriptors(&netdir2), 1);
+    }
 }