diff --git a/src/bin/server.rs b/src/bin/server.rs index ab36705..9a0cc04 100644 --- a/src/bin/server.rs +++ b/src/bin/server.rs @@ -40,6 +40,7 @@ pub struct Config { pub db: DbConfig, // map of distributor name to IP:port to contact it pub distributors: BTreeMap, + extra_infos_base_url: String, //require_bridge_token: bool, port: u16, } @@ -58,8 +59,12 @@ impl Default for DbConfig { } } -async fn update_daily_info(db: &Db, distributors: &BTreeMap) { - update_extra_infos(&db).await; +async fn update_daily_info( + db: &Db, + distributors: &BTreeMap, + extra_infos_base_url: &str, +) { + update_extra_infos(&db, &extra_infos_base_url).await; update_negative_reports(&db, &distributors).await; update_positive_reports(&db, &distributors).await; let new_blockages = guess_blockages(&db, &analyzer::ExampleAnalyzer {}); @@ -73,11 +78,12 @@ async fn run_updater(updater_tx: mpsc::Sender) { async fn create_context_manager( db_config: DbConfig, distributors: BTreeMap, + extra_infos_base_url: &str, context_rx: mpsc::Receiver, mut kill: broadcast::Receiver<()>, ) { tokio::select! { - create_context = context_manager(db_config, distributors, context_rx) => create_context, + create_context = context_manager(db_config, distributors, extra_infos_base_url, context_rx) => create_context, _ = kill.recv() => {println!("Shut down manager");}, } } @@ -85,6 +91,7 @@ async fn create_context_manager( async fn context_manager( db_config: DbConfig, distributors: BTreeMap, + extra_infos_base_url: &str, mut context_rx: mpsc::Receiver, ) { let db: Db = sled::open(&db_config.db_path).unwrap(); @@ -105,7 +112,7 @@ async fn context_manager( println!("Shutdown Sent."); } Update {} => { - update_daily_info(&db, &distributors).await; + update_daily_info(&db, &distributors, &extra_infos_base_url).await; } } } @@ -175,7 +182,14 @@ async fn main() { }); let context_manager = spawn(async move { - create_context_manager(config.db, config.distributors, request_rx, kill).await + create_context_manager( + config.db, + config.distributors, + &config.extra_infos_base_url, + request_rx, + kill, + ) + .await }); let make_service = make_service_fn(move |_conn: &AddrStream| { diff --git a/src/extra_info.rs b/src/extra_info.rs index fbd5bd0..b33f975 100644 --- a/src/extra_info.rs +++ b/src/extra_info.rs @@ -148,9 +148,9 @@ pub fn add_extra_infos<'a>(filename: &str, set: &mut HashSet) { /// Download new extra-infos files and save them in DIRECTORY. This function /// returns the set of newly downloaded filenames. pub async fn download_extra_infos( + base_url: &str, ) -> Result, Box> { // Download directory of recent extra-infos - let base_url = "https://collector.torproject.org/recent/bridge-descriptors/extra-infos/"; let url = base_url.parse().unwrap(); let https = hyper_rustls::HttpsConnectorBuilder::new() .with_native_roots() // TODO: Pin certificate? Is this data signed/verifiable? diff --git a/src/lib.rs b/src/lib.rs index 59cd6ed..3623492 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -165,6 +165,18 @@ impl fmt::Display for BridgeCountryInfo { } } +/// We store a set of all known bridges so that we can later iterate over them. +/// This function just adds a bridge fingerprint to that set. +pub fn add_bridge_to_db(db: &Db, fingerprint: [u8; 20]) { + let mut bridges = match db.get("bridges").unwrap() { + Some(v) => bincode::deserialize(&v).unwrap(), + None => HashSet::<[u8; 20]>::new(), + }; + bridges.insert(fingerprint); + db.insert("bridges", bincode::serialize(&bridges).unwrap()) + .unwrap(); +} + // Process extra-infos /// Adds the extra-info data for a single bridge to the database. If the @@ -172,18 +184,14 @@ impl fmt::Display for BridgeCountryInfo { /// but this extra-info contains different data for some reason, use the /// greater count of connections from each country. pub fn add_extra_info_to_db(db: &Db, extra_info: ExtraInfo) { - let mut bridges = match db.get("bridges").unwrap() { - Some(v) => bincode::deserialize(&v).unwrap(), - None => BTreeMap::<[u8; 20], BridgeInfo>::new(), - }; let fingerprint = extra_info.fingerprint; - if !bridges.contains_key(&fingerprint) { - bridges.insert( - fingerprint, - BridgeInfo::new(fingerprint, &extra_info.nickname), - ); - } - let bridge_info = bridges.get_mut(&fingerprint).unwrap(); + let mut bridge_info = match db.get(fingerprint).unwrap() { + Some(v) => bincode::deserialize(&v).unwrap(), + None => { + add_bridge_to_db(&db, fingerprint); + BridgeInfo::new(fingerprint, &extra_info.nickname) + } + }; for country in extra_info.bridge_ips.keys() { if bridge_info.info_by_country.contains_key::(country) { bridge_info @@ -209,12 +217,12 @@ pub fn add_extra_info_to_db(db: &Db, extra_info: ExtraInfo) { } } // Commit changes to database - db.insert("bridges", bincode::serialize(&bridges).unwrap()) + db.insert(fingerprint, bincode::serialize(&bridge_info).unwrap()) .unwrap(); } /// Download new extra-infos files and add their data to the database -pub async fn update_extra_infos(db: &Db) { +pub async fn update_extra_infos(db: &Db, base_url: &str) { // Track which files have been processed. This is slightly redundant // because we're only downloading files we don't already have, but it // might be a good idea to check in case we downloaded a file but didn't @@ -224,7 +232,7 @@ pub async fn update_extra_infos(db: &Db) { None => HashSet::::new(), }; - let new_files = extra_info::download_extra_infos().await.unwrap(); + let new_files = extra_info::download_extra_infos(base_url).await.unwrap(); let mut new_extra_infos = HashSet::::new(); @@ -331,18 +339,16 @@ pub async fn update_negative_reports(db: &Db, distributors: &BTreeMap bincode::deserialize(&v).unwrap(), - None => BTreeMap::<[u8; 20], BridgeInfo>::new(), - }; - // Get bridge info or make new one - if !bridges.contains_key(&fingerprint) { - // This case shouldn't happen unless the bridge hasn't published - // any bridge stats. - bridges.insert(fingerprint, BridgeInfo::new(fingerprint, &"".to_string())); - } - let bridge_info = bridges.get_mut(&fingerprint).unwrap(); + let mut bridge_info = match db.get(fingerprint).unwrap() { + Some(v) => bincode::deserialize(&v).unwrap(), + None => { + // This case shouldn't happen unless the bridge hasn't + // published any bridge stats. + add_bridge_to_db(&db, fingerprint); + BridgeInfo::new(fingerprint, &String::default()) + } + }; // Add the new report count to it if bridge_info.info_by_country.contains_key(&country) { @@ -358,7 +364,7 @@ pub async fn update_negative_reports(db: &Db, distributors: &BTreeMap bincode::deserialize(&v).unwrap(), - None => BTreeMap::<[u8; 20], BridgeInfo>::new(), - }; - // Get bridge info or make new one - if !bridges.contains_key(&fingerprint) { - // This case shouldn't happen unless the bridge hasn't published - // any bridge stats. - bridges.insert(fingerprint, BridgeInfo::new(fingerprint, &"".to_string())); - } - let bridge_info = bridges.get_mut(&fingerprint).unwrap(); + let mut bridge_info = match db.get(fingerprint).unwrap() { + Some(v) => bincode::deserialize(&v).unwrap(), + None => { + // This case shouldn't happen unless the bridge hasn't + // published any bridge stats. + add_bridge_to_db(&db, fingerprint); + BridgeInfo::new(fingerprint, &String::default()) + } + }; // Add the new report count to it if bridge_info.info_by_country.contains_key(&country) { @@ -476,7 +479,7 @@ pub async fn update_positive_reports(db: &Db, distributors: &BTreeMap HashMap<[u8; 20], Ha // Map of bridge fingerprint to set of countries which newly block it let mut blockages = HashMap::<[u8; 20], HashSet>::new(); - // Get bridge data from database - let mut bridges = match db.get("bridges").unwrap() { + // Get list of bridges from database + let bridges = match db.get("bridges").unwrap() { Some(v) => bincode::deserialize(&v).unwrap(), - None => BTreeMap::<[u8; 20], BridgeInfo>::new(), + None => HashSet::<[u8; 20]>::new(), }; // Guess for each bridge - for (fingerprint, bridge_info) in &mut bridges { + for fingerprint in bridges { + let mut bridge_info: BridgeInfo = + bincode::deserialize(&db.get(fingerprint).unwrap().unwrap()).unwrap(); let mut new_blockages = HashSet::::new(); let blocked_in = analyzer.blocked_in(&bridge_info); for country in blocked_in { @@ -516,12 +521,12 @@ pub fn guess_blockages(db: &Db, analyzer: &dyn Analyzer) -> HashMap<[u8; 20], Ha bridge_country_info.blocked = true; } } - blockages.insert(*fingerprint, new_blockages); - } + blockages.insert(fingerprint, new_blockages); - // Commit changes to database - db.insert("bridges", bincode::serialize(&bridges).unwrap()) - .unwrap(); + // Commit changes to database + db.insert(fingerprint, bincode::serialize(&bridge_info).unwrap()) + .unwrap(); + } // Return map of new blockages blockages diff --git a/src/tests.rs b/src/tests.rs index 1c2fe27..abc8f59 100644 --- a/src/tests.rs +++ b/src/tests.rs @@ -144,6 +144,36 @@ pub fn random() -> BridgeLine { res } +#[tokio::test] +async fn test_extra_infos() { + let bridge_to_test = + array_bytes::hex2array("72E12B89136B45BBC81D1EF0AC7DDDBB91B148DB").unwrap(); + + // Open test database + let db: Db = sled::open("test_db").unwrap(); + + // Delete all data in test DB + db.clear().unwrap(); + assert!(!db.contains_key("bridges").unwrap()); + assert!(!db.contains_key(bridge_to_test).unwrap()); + + // Download and process recent extra-infos files + update_extra_infos( + &db, + "https://collector.torproject.org/recent/bridge-descriptors/extra-infos/", + ) + .await; + + // Check that DB contains information on a bridge with high uptime + assert!(db.contains_key("bridges").unwrap()); + let bridges: HashSet<[u8; 20]> = + bincode::deserialize(&db.get("bridges").unwrap().unwrap()).unwrap(); + assert!(bridges.contains(&bridge_to_test)); + assert!(db.contains_key(bridge_to_test).unwrap()); + let bridge_info: BridgeInfo = + bincode::deserialize(&db.get(bridge_to_test).unwrap().unwrap()).unwrap(); +} + #[test] fn test_negative_reports() { let mut th = TestHarness::new();