Track bridge info by country first, then date

src/lib.rs

@@ -49,11 +49,8 @@ pub struct BridgeInfo {
     /// first Julian date we started collecting data on this bridge
     pub first_seen: u32,
 
-    /// list of countries where the bridge is believed to be blocked
-    pub blocked_in: Vec<String>,
-
-    /// map of dates to data for that day
-    pub info_by_day: HashMap<u32, DailyBridgeInfo>,
+    /// map of countries to data for this bridge in that country
+    pub info_by_country: HashMap<String, BridgeCountryInfo>,
 }
 
 impl BridgeInfo {
@@ -62,8 +59,7 @@ impl BridgeInfo {
             fingerprint: fingerprint,
             nickname: nickname.to_string(),
             first_seen: get_date(),
-            blocked_in: Vec::<String>::new(),
-            info_by_day: HashMap::<u32, DailyBridgeInfo>::new(),
+            info_by_country: HashMap::<String, BridgeCountryInfo>::new(),
         }
     }
 }
@@ -76,15 +72,11 @@ impl fmt::Display for BridgeInfo {
         );
         str.push_str(format!("nickname: {}\n", self.nickname).as_str());
         str.push_str(format!("first_seen: {}\n", self.first_seen).as_str());
-        str.push_str("blocked_in:");
-        for country in &self.blocked_in {
-            str.push_str(format!("\n  {}", country).as_str());
-        }
-        str.push_str("info_by_day:");
-        for day in self.info_by_day.keys() {
-            str.push_str(format!("\n  day: {}", day).as_str());
-            let daily_info = self.info_by_day.get(day).unwrap();
-            for line in daily_info.to_string().lines() {
+        str.push_str("info_by_country:");
+        for country in self.info_by_country.keys() {
+            str.push_str(format!("\n  country: {}", country).as_str());
+            let country_info = self.info_by_country.get(country).unwrap();
+            for line in country_info.to_string().lines() {
                 str.push_str(format!("\n  {}", line).as_str());
             }
         }
@@ -99,62 +91,49 @@ pub enum BridgeInfoType {
     PositiveReports,
 }
 
-/// Information about bridge reachability, gathered daily
+/// Information about bridge reachability from a given country
 #[derive(Serialize, Deserialize)]
-pub struct DailyBridgeInfo {
-    pub info_by_country: BTreeMap<String, BTreeMap<BridgeInfoType, u32>>,
+pub struct BridgeCountryInfo {
+    pub info_by_day: BTreeMap<u32, BTreeMap<BridgeInfoType, u32>>,
+    pub blocked: bool,
 }
 
-impl DailyBridgeInfo {
+impl BridgeCountryInfo {
     pub fn new() -> Self {
         Self {
-            info_by_country: BTreeMap::<String, BTreeMap<BridgeInfoType, u32>>::new(),
+            info_by_day: BTreeMap::<u32, BTreeMap<BridgeInfoType, u32>>::new(),
+            blocked: false,
         }
     }
 
-    pub fn add_info(
-        &mut self,
-        info_type: BridgeInfoType,
-        count_per_country: &BTreeMap<String, u32>,
-    ) {
-        for country in count_per_country.keys() {
-            if self.info_by_country.contains_key(country) {
-                let info = self.info_by_country.get_mut(country).unwrap();
-                if !info.contains_key(&info_type) {
-                    info.insert(
-                        info_type,
-                        *count_per_country.get(&country.to_string()).unwrap(),
-                    );
-                } else if info_type == BridgeInfoType::BridgeIps {
-                    // Use newest value we've seen today
-                    if info.get(&info_type).unwrap() < count_per_country.get(country).unwrap() {
-                        info.insert(
-                            BridgeInfoType::BridgeIps,
-                            *count_per_country.get(&country.to_string()).unwrap(),
-                        );
-                    }
-                } else {
-                    let new_count = info.get(&info_type).unwrap()
-                        + *count_per_country.get(&country.to_string()).unwrap();
-                    info.insert(info_type, new_count);
+    pub fn add_info(&mut self, info_type: BridgeInfoType, date: u32, count: u32) {
+        if self.info_by_day.contains_key(&date) {
+            let info = self.info_by_day.get_mut(&date).unwrap();
+            if !info.contains_key(&info_type) {
+                info.insert(info_type, count);
+            } else if info_type == BridgeInfoType::BridgeIps {
+                if *info.get(&info_type).unwrap() < count {
+                    // Use highest value we've seen today
+                    info.insert(info_type, count);
                 }
             } else {
-                let mut info = BTreeMap::<BridgeInfoType, u32>::new();
-                info.insert(
-                    info_type,
-                    *count_per_country.get(&country.to_string()).unwrap(),
-                );
-                self.info_by_country.insert(country.to_string(), info);
+                // Add count to previous count for reports
+                let new_count = info.get(&info_type).unwrap() + count;
+                info.insert(info_type, new_count);
             }
+        } else {
+            let mut info = BTreeMap::<BridgeInfoType, u32>::new();
+            info.insert(info_type, count);
+            self.info_by_day.insert(date, info);
         }
     }
 }
 
-impl fmt::Display for DailyBridgeInfo {
+impl fmt::Display for BridgeCountryInfo {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         let mut str = String::from("info:");
-        for country in self.info_by_country.keys() {
-            let info = self.info_by_country.get(country).unwrap();
+        for date in self.info_by_day.keys() {
+            let info = self.info_by_day.get(date).unwrap();
             let ip_count = match info.get(&BridgeInfoType::BridgeIps) {
                 Some(v) => v,
                 None => &0,
@@ -170,8 +149,8 @@ impl fmt::Display for DailyBridgeInfo {
             if ip_count > &0 || nr_count > &0 || pr_count > &0 {
                 str.push_str(
                     format!(
-                        "\n  cc: {}\n    connections: {}\n    negative reports: {}\n    positive reports: {}",
-                        country,
+                        "\n  date: {}\n    connections: {}\n    negative reports: {}\n    positive reports: {}",
+                        date,
                         ip_count,
                         nr_count,
                         pr_count,
@@ -196,20 +175,29 @@ pub fn add_extra_info_to_db(db: &Db, extra_info: ExtraInfo) {
         Some(v) => bincode::deserialize(&v).unwrap(),
         None => BridgeInfo::new(fingerprint, &extra_info.nickname),
     };
-    // If we already have an entry, compare it with the new one. For each
-    // country:count mapping, use the greater of the two counts.
-    if bridge_info.info_by_day.contains_key(&extra_info.date) {
-        let daily_bridge_info = bridge_info.info_by_day.get_mut(&extra_info.date).unwrap();
-        daily_bridge_info.add_info(BridgeInfoType::BridgeIps, &extra_info.bridge_ips);
-    } else {
-        // No existing entry; make a new one.
-        let mut daily_bridge_info = DailyBridgeInfo {
-            info_by_country: BTreeMap::<String, BTreeMap<BridgeInfoType, u32>>::new(),
-        };
-        daily_bridge_info.add_info(BridgeInfoType::BridgeIps, &extra_info.bridge_ips);
-        bridge_info
-            .info_by_day
-            .insert(extra_info.date, daily_bridge_info);
+    for country in extra_info.bridge_ips.keys() {
+        if bridge_info.info_by_country.contains_key::<String>(country) {
+            bridge_info
+                .info_by_country
+                .get_mut(country)
+                .unwrap()
+                .add_info(
+                    BridgeInfoType::BridgeIps,
+                    extra_info.date,
+                    *extra_info.bridge_ips.get(country).unwrap(),
+                );
+        } else {
+            // No existing entry; make a new one.
+            let mut bridge_country_info = BridgeCountryInfo::new();
+            bridge_country_info.add_info(
+                BridgeInfoType::BridgeIps,
+                extra_info.date,
+                *extra_info.bridge_ips.get(country).unwrap(),
+            );
+            bridge_info
+                .info_by_country
+                .insert(country.to_string(), bridge_country_info);
+        }
     }
     // Commit changes to database
     db.insert(fingerprint, bincode::serialize(&bridge_info).unwrap())
@@ -260,11 +248,12 @@ pub fn save_negative_report_to_process(db: &Db, nr: NegativeReport) {
         Some(v) => bincode::deserialize(&v).unwrap(),
         None => BTreeMap::<String, BTreeMap<String, u32>>::new(),
     };
-    // Store to-be-processed reports with key [fingerprint]_[date]
+    // Store to-be-processed reports with key [fingerprint]_[country]_[date]
     let map_key = format!(
-        "{}_{}",
+        "{}_{}_{}",
         array_bytes::bytes2hex("", &nr.fingerprint),
-        &nr.date
+        &nr.country,
+        &nr.date,
     );
     let serialized_nr = nr.to_json();
     if reports.contains_key(&map_key) {
@@ -322,12 +311,9 @@ pub async fn update_negative_reports(db: &Db, distributors: &BTreeMap<BridgeDist
         Some(v) => bincode::deserialize(&v).unwrap(),
         None => BTreeMap::<String, BTreeMap<String, u32>>::new(),
     };
-    for bridge_date in all_negative_reports.keys() {
-        // We could parse the fingerprint and date:
-        //let fingerprint: [u8; 20] = array_bytes::hex2array(&bridge_date[0..40]).unwrap();
-        //let date: u32 = &bridge_date[41..].parse().unwrap();
-        // but instead, let's just get it from the first report
-        let reports = all_negative_reports.get(bridge_date).unwrap();
+    // Key is [fingerprint]_[country]_[date]
+    for bridge_country_date in all_negative_reports.keys() {
+        let reports = all_negative_reports.get(bridge_country_date).unwrap();
         if !reports.is_empty() {
             let first_report: SerializableNegativeReport =
                 serde_json::from_str(reports.first_key_value().unwrap().0).unwrap();
@@ -335,8 +321,7 @@ pub async fn update_negative_reports(db: &Db, distributors: &BTreeMap<BridgeDist
             let date = first_report.date;
             let country = first_report.country;
             let count_valid = verify_negative_reports(&distributors, reports).await;
-            let mut count_per_country = BTreeMap::<String, u32>::new();
-            count_per_country.insert(country, count_valid).unwrap();
+
             let mut bridge_info = match db.get(&fingerprint).unwrap() {
                 Some(v) => bincode::deserialize(&v).unwrap(),
                 // It should already exist, unless the bridge hasn't published
@@ -344,17 +329,19 @@ pub async fn update_negative_reports(db: &Db, distributors: &BTreeMap<BridgeDist
                 None => BridgeInfo::new(fingerprint, &"".to_string()),
             };
             // Add the new report count to it
-            if bridge_info.info_by_day.contains_key(&date) {
-                let daily_bridge_info = bridge_info.info_by_day.get_mut(&date).unwrap();
-                daily_bridge_info.add_info(BridgeInfoType::NegativeReports, &count_per_country);
+            if bridge_info.info_by_country.contains_key(&country) {
+                let bridge_country_info = bridge_info.info_by_country.get_mut(&country).unwrap();
+                bridge_country_info.add_info(BridgeInfoType::NegativeReports, date, count_valid);
                 // Commit changes to database
                 db.insert(fingerprint, bincode::serialize(&bridge_info).unwrap())
                     .unwrap();
             } else {
                 // No existing entry; make a new one.
-                let mut daily_bridge_info = DailyBridgeInfo::new();
-                daily_bridge_info.add_info(BridgeInfoType::NegativeReports, &count_per_country);
-                bridge_info.info_by_day.insert(date, daily_bridge_info);
+                let mut bridge_country_info = BridgeCountryInfo::new();
+                bridge_country_info.add_info(BridgeInfoType::NegativeReports, date, count_valid);
+                bridge_info
+                    .info_by_country
+                    .insert(country, bridge_country_info);
                 // Commit changes to database
                 db.insert(fingerprint, bincode::serialize(&bridge_info).unwrap())
                     .unwrap();
@@ -380,11 +367,12 @@ pub fn save_positive_report_to_process(db: &Db, pr: PositiveReport) {
         Some(v) => bincode::deserialize(&v).unwrap(),
         None => BTreeMap::<String, Vec<SerializablePositiveReport>>::new(),
     };
-    // Store to-be-processed reports with key [fingerprint]_[date]
+    // Store to-be-processed reports with key [fingerprint]_[country]_[date]
     let map_key = format!(
-        "{}_{}",
+        "{}_{}_{}",
         array_bytes::bytes2hex("", &pr.fingerprint),
-        &pr.date
+        &pr.country,
+        &pr.date,
     );
     if reports.contains_key(&map_key) {
         reports
@@ -437,20 +425,15 @@ pub async fn update_positive_reports(db: &Db, distributors: &BTreeMap<BridgeDist
         Some(v) => bincode::deserialize(&v).unwrap(),
         None => BTreeMap::<String, Vec<SerializablePositiveReport>>::new(),
     };
-    for bridge_date in all_positive_reports.keys() {
-        // We could parse the fingerprint and date:
-        //let fingerprint: [u8; 20] = array_bytes::hex2array(&bridge_date[0..40]).unwrap();
-        //let date: u32 = &bridge_date[41..].parse().unwrap();
-        // but instead, let's just get it from the first report
-        let reports = all_positive_reports.get(bridge_date).unwrap();
+    // Key is [fingerprint]_[country]_[date]
+    for bridge_country_date in all_positive_reports.keys() {
+        let reports = all_positive_reports.get(bridge_country_date).unwrap();
         if !reports.is_empty() {
             let first_report = &reports[0];
             let fingerprint = first_report.fingerprint;
             let date = first_report.date;
             let country = first_report.country.clone();
             let count_valid = verify_positive_reports(&distributors, reports).await;
-            let mut count_per_country = BTreeMap::<String, u32>::new();
-            count_per_country.insert(country, count_valid).unwrap();
             let mut bridge_info = match db.get(&fingerprint).unwrap() {
                 Some(v) => bincode::deserialize(&v).unwrap(),
                 // It should already exist, unless the bridge hasn't published
@@ -458,17 +441,19 @@ pub async fn update_positive_reports(db: &Db, distributors: &BTreeMap<BridgeDist
                 None => BridgeInfo::new(fingerprint, &"".to_string()),
             };
             // Add the new report count to it
-            if bridge_info.info_by_day.contains_key(&date) {
-                let daily_bridge_info = bridge_info.info_by_day.get_mut(&date).unwrap();
-                daily_bridge_info.add_info(BridgeInfoType::PositiveReports, &count_per_country);
+            if bridge_info.info_by_country.contains_key(&country) {
+                let bridge_country_info = bridge_info.info_by_country.get_mut(&country).unwrap();
+                bridge_country_info.add_info(BridgeInfoType::PositiveReports, date, count_valid);
                 // Commit changes to database
                 db.insert(fingerprint, bincode::serialize(&bridge_info).unwrap())
                     .unwrap();
             } else {
                 // No existing entry; make a new one.
-                let mut daily_bridge_info = DailyBridgeInfo::new();
-                daily_bridge_info.add_info(BridgeInfoType::PositiveReports, &count_per_country);
-                bridge_info.info_by_day.insert(date, daily_bridge_info);
+                let mut bridge_country_info = BridgeCountryInfo::new();
+                bridge_country_info.add_info(BridgeInfoType::PositiveReports, date, count_valid);
+                bridge_info
+                    .info_by_country
+                    .insert(country, bridge_country_info);
                 // Commit changes to database
                 db.insert(fingerprint, bincode::serialize(&bridge_info).unwrap())
                     .unwrap();