From 138ae1f5c966829760d7716aa2b4352714eb3895 Mon Sep 17 00:00:00 2001 From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com> Date: Sun, 25 Aug 2024 00:45:26 +0900 Subject: [PATCH] refactor: AllSenderPeers and ReceiverPeers --- mixnet-rs/ordering/src/iteration.rs | 14 ++--- mixnet-rs/ordering/src/message.rs | 6 +- mixnet-rs/ordering/src/outputs.rs | 22 +++++--- mixnet-rs/ordering/src/topology.rs | 87 +++++++++++++++++++++-------- 4 files changed, 88 insertions(+), 41 deletions(-) diff --git a/mixnet-rs/ordering/src/iteration.rs b/mixnet-rs/ordering/src/iteration.rs index 39e765d..1a8c2dd 100644 --- a/mixnet-rs/ordering/src/iteration.rs +++ b/mixnet-rs/ordering/src/iteration.rs @@ -15,7 +15,7 @@ use crate::{ }; pub fn run_iteration(paramset: ParamSet, seed: u64, outputs: &mut Outputs) -> f32 { - let (mut mixnodes, sender_peers_list, receiver_peer_conn_idx) = if paramset.random_topology { + let (mut mixnodes, all_sender_peers, receiver_peers) = if paramset.random_topology { build_random_network(¶mset, seed, outputs) } else { build_striped_network(¶mset, seed) @@ -56,9 +56,9 @@ pub fn run_iteration(paramset: ParamSet, seed: u64, outputs: &mut Outputs) -> f3 // All senders emit a message (data or noise) to all of their own adjacent peers. if all_sent_count < all_sent_count_target { // For each sender - for (sender_idx, sender_peers) in sender_peers_list.iter().enumerate() { + for (sender_idx, sender_peers) in all_sender_peers.iter() { if try_probability(&mut data_msg_rng, paramset.sender_data_msg_prob) { - let msg = data_msg_gen.next(sender_idx.try_into().unwrap()); + let msg = data_msg_gen.next(sender_idx); sender_peers.iter().for_each(|peer_id| { mixnodes .get_mut(*peer_id as usize) @@ -107,13 +107,13 @@ pub fn run_iteration(paramset: ParamSet, seed: u64, outputs: &mut Outputs) -> f3 } } // Record msg to the sequence - let conn_idx = receiver_peer_conn_idx.get(&relayer_id).unwrap(); - outputs.add_recv_msg(&msg, *conn_idx as usize); + let conn_idx = receiver_peers.conn_idx(&relayer_id).unwrap(); + outputs.add_recv_msg(&msg, conn_idx); } Message::Noise => { // Record noise to the sequence - let conn_idx = receiver_peer_conn_idx.get(&relayer_id).unwrap(); - outputs.add_recv_noise(*conn_idx as usize); + let conn_idx = receiver_peers.conn_idx(&relayer_id).unwrap(); + outputs.add_recv_noise(conn_idx); } } } else if let Message::Data(msg) = msg { diff --git a/mixnet-rs/ordering/src/message.rs b/mixnet-rs/ordering/src/message.rs index 2eecacf..f62f8a4 100644 --- a/mixnet-rs/ordering/src/message.rs +++ b/mixnet-rs/ordering/src/message.rs @@ -1,8 +1,10 @@ use std::fmt::Display; +pub type SenderIdx = u8; + #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub struct DataMessage { - pub sender: u8, + pub sender: SenderIdx, pub msg_id: u32, } @@ -23,7 +25,7 @@ impl DataMessageGenerator { } } - pub fn next(&mut self, sender: u8) -> DataMessage { + pub fn next(&mut self, sender: SenderIdx) -> DataMessage { let msg_id = self.next_msg_ids[sender as usize]; self.next_msg_ids[sender as usize] += 1; DataMessage { sender, msg_id } diff --git a/mixnet-rs/ordering/src/outputs.rs b/mixnet-rs/ordering/src/outputs.rs index e43b895..9f1dccf 100644 --- a/mixnet-rs/ordering/src/outputs.rs +++ b/mixnet-rs/ordering/src/outputs.rs @@ -5,7 +5,11 @@ use protocol::{ topology::Topology, }; -use crate::{message::DataMessage, ordercoeff::SequenceWriter}; +use crate::{ + message::{DataMessage, SenderIdx}, + ordercoeff::SequenceWriter, + topology::AllSenderPeers, +}; pub struct Outputs { closed: bool, @@ -103,18 +107,18 @@ impl Outputs { writer.add_message(msg); } - pub fn add_sent_noise(&mut self, sender_idx: usize) { - let writer = &mut self.sent_sequence_writers[sender_idx]; + pub fn add_sent_noise(&mut self, sender_idx: SenderIdx) { + let writer = &mut self.sent_sequence_writers[sender_idx as usize]; writer.add_noise(); } - pub fn add_recv_msg(&mut self, msg: &DataMessage, conn_idx: usize) { - let writer = &mut self.recv_sequence_writers[conn_idx]; + pub fn add_recv_msg(&mut self, msg: &DataMessage, conn_idx: u16) { + let writer = &mut self.recv_sequence_writers[conn_idx as usize]; writer.add_message(msg); } - pub fn add_recv_noise(&mut self, conn_idx: usize) { - let writer = &mut self.recv_sequence_writers[conn_idx]; + pub fn add_recv_noise(&mut self, conn_idx: u16) { + let writer = &mut self.recv_sequence_writers[conn_idx as usize]; writer.add_noise(); } @@ -151,7 +155,7 @@ impl Outputs { pub fn write_topology( &self, topology: &Topology, - sender_peers_list: &[Vec], + all_sender_peers: &AllSenderPeers, receiver_peers: &[NodeId], ) { let mut writer = csv::Writer::from_path(&self.topology_path).unwrap(); @@ -176,7 +180,7 @@ impl Outputs { } // Write peers of senders - for (sender_idx, peers) in sender_peers_list.iter().enumerate() { + for (sender_idx, peers) in all_sender_peers.iter() { writer .write_record(&[ format!("sender-{}", sender_idx), diff --git a/mixnet-rs/ordering/src/topology.rs b/mixnet-rs/ordering/src/topology.rs index bde0455..7d81e96 100644 --- a/mixnet-rs/ordering/src/topology.rs +++ b/mixnet-rs/ordering/src/topology.rs @@ -8,14 +8,14 @@ use protocol::{ use rand::{rngs::StdRng, seq::SliceRandom, RngCore, SeedableRng}; use rustc_hash::FxHashMap; -use crate::{outputs::Outputs, paramset::ParamSet}; +use crate::{message::SenderIdx, outputs::Outputs, paramset::ParamSet}; pub const RECEIVER_NODE_ID: NodeId = NodeId::MAX; pub fn build_striped_network( paramset: &ParamSet, seed: u64, -) -> (Vec>, Vec>, FxHashMap) { +) -> (Vec>, AllSenderPeers, ReceiverPeers) { assert!(!paramset.random_topology); let mut next_node_id: NodeId = 0; let mut queue_seed_rng = StdRng::seed_from_u64(seed); @@ -43,7 +43,7 @@ pub fn build_striped_network = FxHashMap::default(); + let mut receiver_peers = ReceiverPeers::new(); for path in paths.iter() { for (i, id) in path.iter().enumerate() { if i != path.len() - 1 { @@ -56,24 +56,28 @@ pub fn build_striped_network> = - vec![ - paths.iter().map(|path| *path.first().unwrap()).collect(); - paramset.num_senders as usize - ]; - (mixnodes, sender_peers_list, receiver_peer_conn_idx) + + let mut all_sender_peers = AllSenderPeers::new(paramset.num_senders); + let sender_peers = paths + .iter() + .map(|path| *path.first().unwrap()) + .collect::>(); + (0..paramset.num_senders).for_each(|_| { + all_sender_peers.add(sender_peers.clone()); + }); + + (mixnodes, all_sender_peers, receiver_peers) } pub fn build_random_network( paramset: &ParamSet, seed: u64, outputs: &mut Outputs, -) -> (Vec>, Vec>, FxHashMap) { +) -> (Vec>, AllSenderPeers, ReceiverPeers) { assert!(paramset.random_topology); // Init mix nodes let mut queue_seed_rng = StdRng::seed_from_u64(seed); @@ -95,7 +99,7 @@ pub fn build_random_network = mixnodes.iter().map(|mixnode| mixnode.id).collect(); assert!(candidates.len() >= paramset.peering_degree as usize); - let mut sender_peers_list: Vec> = Vec::with_capacity(paramset.num_senders as usize); + let mut all_sender_peers = AllSenderPeers::new(paramset.num_senders); for _ in 0..paramset.num_senders { candidates.as_mut_slice().shuffle(&mut peers_rng); let mut peers: Vec = candidates @@ -104,15 +108,15 @@ pub fn build_random_network = candidates + let mut receiver_peer_ids: Vec = candidates .iter() .cloned() .take(paramset.peering_degree as usize) .collect(); - receiver_peers.sort(); + receiver_peer_ids.sort(); // Connect mix nodes let topology = build_topology( @@ -129,18 +133,55 @@ pub fn build_random_network conn_idx - let mut receiver_peer_conn_idx: FxHashMap = FxHashMap::default(); - for (conn_idx, mixnode_id) in receiver_peers.iter().enumerate() { + let mut receiver_peers = ReceiverPeers::new(); + for (conn_idx, mixnode_id) in receiver_peer_ids.iter().enumerate() { let mixnode = mixnodes.get_mut(*mixnode_id as usize).unwrap(); assert_eq!(mixnode.id, *mixnode_id); mixnode.connect(RECEIVER_NODE_ID); - receiver_peer_conn_idx.insert(*mixnode_id, conn_idx.try_into().unwrap()); + receiver_peers.add(*mixnode_id, conn_idx.try_into().unwrap()); } - outputs.write_topology(&topology, &sender_peers_list, &receiver_peers); + outputs.write_topology(&topology, &all_sender_peers, &receiver_peer_ids); - (mixnodes, sender_peers_list, receiver_peer_conn_idx) + (mixnodes, all_sender_peers, receiver_peers) +} + +pub struct AllSenderPeers(Vec>); + +impl AllSenderPeers { + fn new(num_senders: u8) -> Self { + Self(Vec::with_capacity(num_senders as usize)) + } + + fn add(&mut self, peers: Vec) { + self.0.push(peers) + } + + pub fn iter(&self) -> impl Iterator)> { + self.0 + .iter() + .enumerate() + .map(|(idx, v)| (idx.try_into().unwrap(), v)) + } +} + +pub struct ReceiverPeers(FxHashMap); + +impl ReceiverPeers { + fn new() -> Self { + ReceiverPeers(FxHashMap::default()) + } + + fn add(&mut self, peer_id: NodeId, conn_idx: u16) { + self.0.insert(peer_id, conn_idx); + } + + pub fn conn_idx(&self, node_id: &NodeId) -> Option { + self.0.get(node_id).cloned() + } + + fn len(&self) -> usize { + self.0.len() + } }