Turn IdMap into an extension trait.

Also, use the "right" (cyclic, potentially expensive) algorithm for
stream IDs.

tor-proto/src/channel/circmap.rs

@@ -8,12 +8,15 @@ use crate::Result;
 
 use futures::channel::{mpsc, oneshot};
 
+use rand::distributions::Distribution;
 use rand::Rng;
+use std::collections::HashMap;
 
 /// Which group of circuit IDs are we allowed to allocate in this map?
 ///
 /// If we initiated the channel, we use High circuit ids.  If we're the
 /// responder, we use low circuit ids.
+#[derive(Copy, Clone)]
 pub(super) enum CircIDRange {
     /// Only use circuit IDs with the MSB cleared.
     Low,
@@ -66,14 +69,16 @@ pub(super) enum CircEnt {
 
 /// A map from circuit IDs to circuit entries. Each channel has one.
 pub(super) struct CircMap {
-    m: IdMap<CircID, CircIDRange, CircEnt>,
+    m: HashMap<CircID, CircEnt>,
+    range: CircIDRange,
 }
 
 impl CircMap {
     /// Make a new empty CircMap
     pub(super) fn new(idrange: CircIDRange) -> Self {
         CircMap {
-            m: IdMap::new(idrange),
+            m: HashMap::new(),
+            range: idrange,
         }
     }
 
@@ -87,8 +92,9 @@ impl CircMap {
         createdsink: oneshot::Sender<ChanMsg>,
         sink: mpsc::Sender<ChanMsg>,
     ) -> Result<CircID> {
+        let mut iter = (&mut self.range).sample_iter(rng).take(16);
         let ent = CircEnt::Opening(createdsink, sink);
-        self.m.add_ent(rng, ent)
+        self.m.add_ent(&mut iter, ent)
     }
 
     /// Return the entry for `id` in this map, if any.
@@ -104,7 +110,7 @@ impl CircMap {
         let ok = matches!(self.m.get(&id), Some(CircEnt::Opening(_, _)));
         if ok {
             if let Some(CircEnt::Opening(oneshot, sink)) = self.m.remove(&id) {
-                self.m.put_ent(id, CircEnt::Open(sink));
+                self.m.insert(id, CircEnt::Open(sink));
                 Some(oneshot)
             } else {
                 panic!("internal error: inconsistent circuit state");

tor-proto/src/circuit.rs

@@ -270,11 +270,9 @@ impl ClientCirc {
         // XXXX Both a bound and a lack of bound are scary here :/
         let (sender, receiver) = mpsc::channel(128);
 
-        let mut rng = rand::thread_rng();
-
         let mut c = self.c.lock().await;
         let hopnum = c.hops.len() - 1;
-        let id = c.hops[hopnum].map.add_ent(&mut rng, sender)?;
+        let id = c.hops[hopnum].map.add_ent(sender)?;
         let relaycell = RelayCell::new(id, begin_msg);
         let hopnum = (hopnum as u8).into();
         c.send_relay_cell(hopnum, false, relaycell).await?;

tor-proto/src/circuit/streammap.rs

@@ -5,10 +5,8 @@ use crate::relaycell::{msg::RelayMsg, StreamID};
 use crate::util::idmap::IdMap;
 use crate::Result;
 
-use rand::distributions::Distribution;
-use rand::Rng;
-
 use futures::channel::mpsc;
+use std::collections::HashMap;
 
 /// The entry for a stream.
 pub(super) enum StreamEnt {
@@ -17,42 +15,28 @@ pub(super) enum StreamEnt {
     Open(mpsc::Sender<RelayMsg>),
 }
 
-/// A distribution to construct (nonzero) stream IDs
-struct StreamIDDist;
-impl Distribution<StreamID> for StreamIDDist {
-    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> StreamID {
-        loop {
-            let val: u16 = rng.gen();
-            if val != 0 {
-                return val.into();
-            }
-        }
-    }
-}
-
 /// A map from stream IDs to stream entries. Each circuit has one for each
 /// hop.
 pub(super) struct StreamMap {
-    m: IdMap<StreamID, StreamIDDist, StreamEnt>,
+    m: HashMap<StreamID, StreamEnt>,
+    i: std::iter::Cycle<std::ops::RangeInclusive<u16>>,
 }
 
 impl StreamMap {
     /// Make a new empty StreamMap.
     pub(super) fn new() -> Self {
+        let iter = (1_u16..=65535_u16).cycle();
         StreamMap {
-            m: IdMap::new(StreamIDDist),
+            m: HashMap::new(),
+            i: iter,
         }
     }
 
     /// Add an entry to this map; return the newly allocated StreamID.
-    pub(super) fn add_ent<R: Rng>(
-        &mut self,
-        rng: &mut R,
-        sink: mpsc::Sender<RelayMsg>,
-    ) -> Result<StreamID> {
+    pub(super) fn add_ent(&mut self, sink: mpsc::Sender<RelayMsg>) -> Result<StreamID> {
         let ent = StreamEnt::Open(sink);
-        let id = self.m.add_ent(rng, ent)?;
-        Ok(id)
+        let mut iter = (&mut self.i).map(|x| x.into()).take(65536);
+        self.m.add_ent(&mut iter, ent)
     }
 
     /// Return the entry for `id` in this map, if any.

tor-proto/src/util/idmap.rs

@@ -1,82 +1,35 @@
 // NOTE: This is a work in progress and I bet I'll refactor it a lot;
 // it needs to stay opaque!
 
-// TODO: I bet we could turn this into an extension trait.
-
 use crate::{Error, Result};
 
-use rand::distributions::Distribution;
-use rand::Rng;
 use std::collections::HashMap;
-use std::hash::Hash;
+use std::hash::{BuildHasher, Hash};
 
-/// An IdMap is map from identifiers to keys, along with a distribution
-/// for allocating new identifiers.
-///
-/// We use it to implement maps for circuit IDs and stream IDs.
-pub struct IdMap<ID, DST, VAL>
+/// Extension trait for hashmap that can add an allocate a new key as
+/// needed.
+pub trait IdMap<K, V>
 where
-    ID: Hash + Eq + Clone,
-    DST: Distribution<ID>,
+    K: Hash + Eq + Clone,
 {
-    d: DST,
-    m: HashMap<ID, VAL>,
+    /// Insert a new entry into this map, allocating an identifier for it.
+    ///
+    /// Keep trying until the iterator is done.
+    fn add_ent<I: Iterator<Item = K>>(&mut self, iter: &mut I, val: V) -> Result<K>;
 }
 
-impl<ID, DST, VAL> IdMap<ID, DST, VAL>
+impl<K, V, S> IdMap<K, V> for HashMap<K, V, S>
 where
-    ID: Hash + Eq + Clone,
-    DST: Distribution<ID>,
+    K: Hash + Eq + Clone,
+    S: BuildHasher,
 {
-    /// Make a new empty map
-    pub fn new(dist: DST) -> Self {
-        Self {
-            d: dist,
-            m: HashMap::new(),
-        }
-    }
-
-    /// Construct a new random identifier for an owned entry in this map.
-    /// This can fail if the map is too full.
-    fn gen_id<R: Rng>(&self, rng: &mut R) -> Option<ID> {
-        // How many times to we try before giving up?
-        const MAX_ATTEMPTS: usize = 16;
-        for _ in 0..MAX_ATTEMPTS {
-            let id = self.d.sample(rng);
-            if !self.m.contains_key(&id) {
-                return Some(id);
+    fn add_ent<I: Iterator<Item = K>>(&mut self, iter: &mut I, val: V) -> Result<K> {
+        for i in iter {
+            if !self.contains_key(&i) {
+                self.insert(i.clone(), val);
+                return Ok(i);
             }
         }
-        None
+        Err(Error::IDRangeFull)
     }
-
-    /// Insert a new entry into this map, allocating an identifier for it.
-    pub fn add_ent<R: Rng>(&mut self, rng: &mut R, val: VAL) -> Result<ID> {
-        let id = self.gen_id(rng).ok_or(Error::IDRangeFull)?;
-        self.m.insert(id.clone(), val);
-        Ok(id)
-    }
-
-    /// Replace the current entry at 'id' with 'val'.
-    pub fn put_ent(&mut self, id: ID, val: VAL) {
-        self.m.insert(id, val);
-    }
-
-    /// Return a reference to the value at 'id'
-    pub fn get(&self, id: &ID) -> Option<&VAL> {
-        self.m.get(id)
-    }
-
-    /// Remove the entry for `id` on this map, if any.
-    pub fn remove(&mut self, id: &ID) -> Option<VAL> {
-        self.m.remove(id)
-    }
-
-    /// Return the entry for `id` in this map, if any.
-    pub fn get_mut(&mut self, id: &ID) -> Option<&mut VAL> {
-        self.m.get_mut(&id)
-    }
-
-    // TODO: Eventually if we want relay support, we'll need to support
-    // IDs chosen by somebody else. But for now, we don't need those.
 }