Add a handy function for picking a random relay.

tor-netdir/Cargo.toml

@@ -14,6 +14,7 @@ signature = "*"
 tor-llcrypto = { path="../tor-llcrypto", version= "*" }
 hex = "0.4.2"
 log = "*"
+rand = "*"
 
 [dev-dependencies]
 simple-logging = "*"

tor-netdir/src/lib.rs

@@ -1,6 +1,5 @@
-#![allow(unused)]
-
 mod err;
+mod pick;
 
 use tor_checkable::{ExternallySigned, SelfSigned, Timebound};
 use tor_netdoc::authcert::AuthCert;
@@ -251,6 +250,16 @@ impl NetDir {
         }
         self.weight_fn.get().unwrap()
     }
+    pub fn pick_relay<'a, R, F>(&'a self, rng: &mut R, reweight: F) -> Option<Relay<'a>>
+    where
+        R: rand::Rng,
+        F: Fn(&Relay<'a>, u32) -> u32,
+    {
+        let weight_fn = self.get_weight_fn();
+        pick::pick_weighted(rng, self.relays(), |r| {
+            reweight(r, r.get_weight(weight_fn)) as u64
+        })
+    }
 }
 
 impl<'a> Relay<'a> {

tor-netdir/src/pick.rs

@@ -0,0 +1,34 @@
+// Performance note: this requires a fast RNG, but doesn't need much
+// storage.
+pub fn pick_weighted<R, I, F>(rng: &mut R, i: I, weightfn: F) -> Option<I::Item>
+where
+    I: Iterator,
+    F: Fn(&I::Item) -> u64,
+    R: rand::Rng,
+{
+    let mut result = None;
+    let mut weight_so_far: u64 = 0;
+
+    // Correctness argument: at the end of each iteration of the loop,
+    // `result` holds a value chosen with weighted probabability from
+    // all of the items yielded so far.  The loop body preserves this
+    // invariant.
+
+    for item in i {
+        let w = weightfn(&item);
+        if w == 0 {
+            continue;
+        }
+        // TODO: panics on overflow. Probably not best.
+        weight_so_far = weight_so_far.checked_add(w).unwrap();
+
+        let x = rng.gen_range(0, weight_so_far);
+        // TODO: we could probably do this in constant-time, if we are
+        // worried about a side-channel.
+        if x < w {
+            result = Some(item);
+        }
+    }
+
+    result
+}