From 6492334affd361bd521acb4839e18f57e1522093 Mon Sep 17 00:00:00 2001 From: Nick Mathewson Date: Wed, 23 Sep 2020 17:53:58 -0400 Subject: [PATCH] Turn IdMap into an extension trait. Also, use the "right" (cyclic, potentially expensive) algorithm for stream IDs. --- tor-proto/src/channel/circmap.rs | 14 +++-- tor-proto/src/circuit.rs | 4 +- tor-proto/src/circuit/streammap.rs | 34 ++++-------- tor-proto/src/util/idmap.rs | 83 +++++++----------------------- 4 files changed, 38 insertions(+), 97 deletions(-) diff --git a/tor-proto/src/channel/circmap.rs b/tor-proto/src/channel/circmap.rs index b84fe79aa..bfd939e2c 100644 --- a/tor-proto/src/channel/circmap.rs +++ b/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, + m: HashMap, + 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, sink: mpsc::Sender, ) -> Result { + 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"); diff --git a/tor-proto/src/circuit.rs b/tor-proto/src/circuit.rs index eb4d5af20..2cfdcc057 100644 --- a/tor-proto/src/circuit.rs +++ b/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?; diff --git a/tor-proto/src/circuit/streammap.rs b/tor-proto/src/circuit/streammap.rs index 32e01445e..35731ed45 100644 --- a/tor-proto/src/circuit/streammap.rs +++ b/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), } -/// A distribution to construct (nonzero) stream IDs -struct StreamIDDist; -impl Distribution for StreamIDDist { - fn sample(&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, + m: HashMap, + i: std::iter::Cycle>, } 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( - &mut self, - rng: &mut R, - sink: mpsc::Sender, - ) -> Result { + pub(super) fn add_ent(&mut self, sink: mpsc::Sender) -> Result { 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. diff --git a/tor-proto/src/util/idmap.rs b/tor-proto/src/util/idmap.rs index 744ab4b51..8946f8e60 100644 --- a/tor-proto/src/util/idmap.rs +++ b/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 +/// Extension trait for hashmap that can add an allocate a new key as +/// needed. +pub trait IdMap where - ID: Hash + Eq + Clone, - DST: Distribution, + K: Hash + Eq + Clone, { - d: DST, - m: HashMap, + /// Insert a new entry into this map, allocating an identifier for it. + /// + /// Keep trying until the iterator is done. + fn add_ent>(&mut self, iter: &mut I, val: V) -> Result; } -impl IdMap +impl IdMap for HashMap where - ID: Hash + Eq + Clone, - DST: Distribution, + 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(&self, rng: &mut R) -> Option { - // 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>(&mut self, iter: &mut I, val: V) -> Result { + 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(&mut self, rng: &mut R, val: VAL) -> Result { - 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 { - 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. }