From ae1b90ce6fb2b2d2e273aae070b6377802957914 Mon Sep 17 00:00:00 2001
From: Vecna <vvecna@uwaterloo.ca>
Date: Sun, 27 Oct 2024 22:43:37 -0400
Subject: [PATCH] Store 0 connection counts if cc does not appear in bridge-ips

---
 src/lib.rs               |  9 +++++++
 src/tests/extra_infos.rs | 56 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 9738548..34ec89e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -339,6 +339,15 @@ pub fn add_extra_info_to_db(db: &Db, extra_info: ExtraInfo) {
                 .insert(country.to_string(), bridge_country_info);
         }
     }
+
+    // Store 0 if we have no connections from a country today, but we
+    // have seen connections from the country in the past.
+    for (country, info) in &mut bridge_info.info_by_country {
+        if !extra_info.bridge_ips.contains_key(country) {
+            info.add_info(BridgeInfoType::BridgeIps, extra_info.date, 0);
+        }
+    }
+
     // Commit changes to database
     db.insert(fingerprint, bincode::serialize(&bridge_info).unwrap())
         .unwrap();
diff --git a/src/tests/extra_infos.rs b/src/tests/extra_infos.rs
index 6c432b1..37dbe6e 100644
--- a/src/tests/extra_infos.rs
+++ b/src/tests/extra_infos.rs
@@ -136,4 +136,60 @@ router-digest F30B38390C375E1EE74BFED844177804442569E0"#;
     assert!(db.contains_key(bridge_to_test).unwrap());
     let _bridge_info: BridgeInfo =
         bincode::deserialize(&db.get(bridge_to_test).unwrap().unwrap()).unwrap();
+
+    // Add another day's data, with no Russia data this time
+    let extra_info_str_2 = r#"@type bridge-extra-info 1.3
+extra-info ElephantBridgeDE2 72E12B89136B45BBC81D1EF0AC7DDDBB91B148DB
+bridge-stats-end 2024-04-06 06:51:44 (86400 s)
+bridge-ips us=32,??=8,au=8,br=8,by=8,cn=8,de=8,eg=8,eu=8,gb=8,ge=8,hr=8,ie=8,ir=8,kp=8,lt=8,mt=8,nl=8,pl=8,ro=8,sg=8,tn=8,tr=8,vn=8"#;
+
+    let extra_info_set_2 = ExtraInfo::parse_file(&extra_info_str_2);
+
+    // Add our extra-info to the server's records
+    {
+        use hyper::{Body, Client, Method, Request};
+        let client = Client::new();
+        let req = Request::builder()
+            .method(Method::POST)
+            .uri("http://localhost:8004/add".parse::<hyper::Uri>().unwrap())
+            .body(Body::from(
+                serde_json::to_string(&extra_info_set_2).unwrap(),
+            ))
+            .unwrap();
+        client.request(req).await.unwrap();
+    }
+
+    // Update extra-infos (add new record)
+    update_extra_infos(&db, "http://localhost:8004/")
+        .await
+        .unwrap();
+
+    let bridge_info: BridgeInfo =
+        bincode::deserialize(&db.get(bridge_to_test).unwrap().unwrap()).unwrap();
+
+    // We should not have any data on Canada
+    assert!(!bridge_info.info_by_country.contains_key("ca"));
+
+    // We should have data on Russia
+    let russia_info = bridge_info.info_by_country.get("ru").unwrap();
+    assert_eq!(
+        *russia_info
+            .info_by_day
+            .get(&2460406)
+            .unwrap()
+            .get(&BridgeInfoType::BridgeIps)
+            .unwrap(),
+        40
+    );
+    // Because we have historical data on Russia, but it wasn't listed
+    // today, we should see 0
+    assert_eq!(
+        *russia_info
+            .info_by_day
+            .get(&2460407)
+            .unwrap()
+            .get(&BridgeInfoType::BridgeIps)
+            .unwrap(),
+        0
+    );
 }