From 9381ec8160c5ebd9bdbae5f8a398f25aeeff2666 Mon Sep 17 00:00:00 2001 From: Makis Arsenis Date: Tue, 25 Aug 2026 21:39:08 -0700 Subject: [PATCH 1/3] Add 5-party networking and PRF key setup to ampc-actor-utils Generalizes the node-setup layer beyond the 3-party ring: NetworkSession gains Role-addressed send_to/receive_from, MpcNetworkHandle/PeerConnections now establish connections to any number of peers instead of exactly 2, and a new MeshControlChannel provides a Role-addressed control-plane channel (with an all-to-all sync barrier) alongside the untouched 3-party ControlChannel. Also fixes build_network_handle deriving only 3 local identities regardless of party count. Adds two new shared-PRF key configurations for a 5-party protocol: ThresholdPrfKeys (a (3,5) config where k_{i,j} is known only to the 3 parties outside {i,j}) and PairwisePrfKeys (lambda_{i,j} known only to i and j), each with a symmetric XOR-combine network handshake in protocol::ops. All existing 3PC structs, signatures, and behavior are unchanged. Co-Authored-By: Claude Sonnet 5 --- ampc-actor-utils/src/execution/local.rs | 22 +- ampc-actor-utils/src/execution/player.rs | 2 +- ampc-actor-utils/src/execution/session.rs | 17 ++ .../src/network/mpc/handle/control_channel.rs | 84 +++++- .../src/network/mpc/handle/data.rs | 26 +- .../src/network/mpc/handle/mod.rs | 54 +++- .../src/network/mpc/handle/network_handle.rs | 208 +++++++++----- ampc-actor-utils/src/protocol/ops.rs | 184 ++++++++++++- ampc-actor-utils/src/protocol/prf.rs | 254 +++++++++++++++++- .../tests/mesh_control_channel.rs | 141 ++++++++++ 10 files changed, 887 insertions(+), 105 deletions(-) create mode 100644 ampc-actor-utils/tests/mesh_control_channel.rs diff --git a/ampc-actor-utils/src/execution/local.rs b/ampc-actor-utils/src/execution/local.rs index 8a7bb8c1..fb9f7028 100644 --- a/ampc-actor-utils/src/execution/local.rs +++ b/ampc-actor-utils/src/execution/local.rs @@ -20,12 +20,24 @@ use std::{ }; use tokio::{sync::Mutex, task::JoinHandle}; +const LOCAL_IDENTITY_NAMES: [&str; 5] = ["alice", "bob", "charlie", "dave", "erin"]; + pub fn generate_local_identities() -> Vec { - vec![ - Identity::from("alice"), - Identity::from("bob"), - Identity::from("charlie"), - ] + generate_local_identities_n(3) +} + +/// Generate `n` deterministic, fixed local identities, for use both in tests +/// and to derive a consistent per-index self-identity across parties in +/// production (see `build_network_handle`). Supports up to 5 parties. +pub fn generate_local_identities_n(n: usize) -> Vec { + assert!( + n <= LOCAL_IDENTITY_NAMES.len(), + "not enough predefined local identities for {n} parties" + ); + LOCAL_IDENTITY_NAMES[..n] + .iter() + .map(|name| Identity::from(*name)) + .collect() } static USED_PORTS: LazyLock>> = LazyLock::new(|| Mutex::new(HashSet::new())); diff --git a/ampc-actor-utils/src/execution/player.rs b/ampc-actor-utils/src/execution/player.rs index 9b5e8ccd..d256dd0c 100644 --- a/ampc-actor-utils/src/execution/player.rs +++ b/ampc-actor-utils/src/execution/player.rs @@ -30,7 +30,7 @@ impl From for Identity { } /// Struct that keeps the player id (role), zero indexed; -#[derive(Debug, Eq, Hash, PartialEq, Copy, Clone)] +#[derive(Debug, Eq, Hash, PartialEq, PartialOrd, Ord, Copy, Clone)] pub struct Role(u8); impl Role { diff --git a/ampc-actor-utils/src/execution/session.rs b/ampc-actor-utils/src/execution/session.rs index 51f4d9fa..e58f8768 100644 --- a/ampc-actor-utils/src/execution/session.rs +++ b/ampc-actor-utils/src/execution/session.rs @@ -69,6 +69,23 @@ impl NetworkSession { let prev_identity = self.prev_identity()?; self.receive(&prev_identity).await } + + /// Send `value` to an arbitrary party identified by `role`. Unlike + /// `send_next`/`send_prev`, this works regardless of ring adjacency and + /// is the preferred API for protocols with more than 3 parties. + pub async fn send_to(&mut self, value: NetworkValue, role: &Role) -> Result<()> { + let identity = self.identity(role)?.clone(); + self.send(value, &identity).await + } + + /// Receive a value from an arbitrary party identified by `role`. Unlike + /// `receive_next`/`receive_prev`, this works regardless of ring + /// adjacency and is the preferred API for protocols with more than 3 + /// parties. + pub async fn receive_from(&mut self, role: &Role) -> Result { + let identity = self.identity(role)?.clone(); + self.receive(&identity).await + } } // Helper methods for sending and receiving VecRingElement. diff --git a/ampc-actor-utils/src/network/mpc/handle/control_channel.rs b/ampc-actor-utils/src/network/mpc/handle/control_channel.rs index c3263f8f..ac6bf611 100644 --- a/ampc-actor-utils/src/network/mpc/handle/control_channel.rs +++ b/ampc-actor-utils/src/network/mpc/handle/control_channel.rs @@ -1,8 +1,10 @@ use async_trait::async_trait; -use eyre::{bail, Result}; +use eyre::{bail, eyre, Result}; +use std::collections::BTreeMap; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio_util::sync::CancellationToken; +use crate::execution::player::Role; use crate::network::mpc::NetworkValue; use crate::network::tcp::NetworkConnection; @@ -169,3 +171,83 @@ impl ControlChannel for TcpControlChannel { Ok(()) } } + +/// A synchronization-safe control-plane channel to every other party, +/// addressed by [`Role`] rather than ring position. Generalizes +/// [`ControlChannel`] beyond the 3-party ring case: works for any number of +/// parties, not just 3. +/// +/// Like [`ControlChannel`], sends block until flushed and there is no +/// automatic retry on error; call +/// [`crate::network::mpc::NetworkHandle::mesh_control_channel`] again to +/// reconnect. +#[async_trait] +pub trait MeshControlChannel: Send { + /// Send a value to `to`. Blocks until flushed. + async fn send(&mut self, to: Role, value: NetworkValue) -> Result<()>; + + /// Receive a value from `from`. Blocks until a full message arrives. + async fn recv(&mut self, from: Role) -> Result; + + /// All-to-all barrier: send a sync token to every other party, then + /// receive one from every other party. + /// + /// Every party must call `sync()` concurrently. Sends are issued before + /// receives to avoid deadlock, as in [`ControlChannel::sync`]. + async fn sync(&mut self) -> Result<()>; +} + +/// [`MeshControlChannel`] implementation over a generic [`NetworkConnection`] +/// stream. Holds one dedicated stream per other party. Constructed by +/// [`crate::network::mpc::NetworkHandle::mesh_control_channel`]. +pub(crate) struct TcpMeshControlChannel { + streams: BTreeMap, + shutdown_ct: CancellationToken, +} + +impl TcpMeshControlChannel { + pub(super) fn new(streams: BTreeMap, shutdown_ct: CancellationToken) -> Self { + Self { + streams, + shutdown_ct, + } + } +} + +#[async_trait] +impl MeshControlChannel for TcpMeshControlChannel { + async fn send(&mut self, to: Role, value: NetworkValue) -> Result<()> { + let stream = self + .streams + .get_mut(&to) + .ok_or_else(|| eyre!("no control channel to role {to:?}"))?; + write_value(stream, value, &self.shutdown_ct).await + } + + async fn recv(&mut self, from: Role) -> Result { + let stream = self + .streams + .get_mut(&from) + .ok_or_else(|| eyre!("no control channel to role {from:?}"))?; + read_value(stream, &self.shutdown_ct).await + } + + async fn sync(&mut self) -> Result<()> { + let token = NetworkValue::Bytes(SYNC_TOKEN_BYTES.to_vec().into()); + // BTreeMap iteration is deterministic (unlike a HashMap), so every + // party visits the same roles in the same order here. + let roles: Vec = self.streams.keys().copied().collect(); + + for role in &roles { + self.send(*role, token.clone()).await?; + } + for role in &roles { + match self.recv(*role).await? { + NetworkValue::Bytes(ref bytes) if &bytes[..] == SYNC_TOKEN_BYTES => {} + _ => bail!("invalid sync token received from role {role:?}"), + } + } + + Ok(()) + } +} diff --git a/ampc-actor-utils/src/network/mpc/handle/data.rs b/ampc-actor-utils/src/network/mpc/handle/data.rs index fbd52f64..abcfa943 100644 --- a/ampc-actor-utils/src/network/mpc/handle/data.rs +++ b/ampc-actor-utils/src/network/mpc/handle/data.rs @@ -14,14 +14,19 @@ pub type OutStream = mpsc::UnboundedSender; pub type InStream = mpsc::UnboundedReceiver; pub struct PeerConnections { - peers: [Arc; 2], - c0: Vec, - c1: Vec, + peers: Vec>, + // conns[i] holds the connections established with peers[i] + conns: Vec>, } impl PeerConnections { - pub fn new(peers: [Arc; 2], c0: Vec, c1: Vec) -> Self { - Self { peers, c0, c1 } + pub fn new(peers: Vec>, conns: Vec>) -> Self { + assert_eq!( + peers.len(), + conns.len(), + "expected one connection group per peer" + ); + Self { peers, conns } } pub fn peer_ids(&self) -> Vec { @@ -34,10 +39,11 @@ impl IntoIterator for PeerConnections { type IntoIter = std::vec::IntoIter<(Identity, Vec)>; fn into_iter(self) -> Self::IntoIter { - vec![ - (self.peers[0].id().clone(), self.c0), - (self.peers[1].id().clone(), self.c1), - ] - .into_iter() + self.peers + .into_iter() + .map(|peer| peer.id().clone()) + .zip(self.conns) + .collect::>() + .into_iter() } } diff --git a/ampc-actor-utils/src/network/mpc/handle/mod.rs b/ampc-actor-utils/src/network/mpc/handle/mod.rs index b65df28f..168a6dfc 100644 --- a/ampc-actor-utils/src/network/mpc/handle/mod.rs +++ b/ampc-actor-utils/src/network/mpc/handle/mod.rs @@ -10,10 +10,10 @@ use std::time::Duration; use self::config::MpcConfig; use self::network_handle::MpcNetworkHandle; -use crate::execution::local::generate_local_identities; +use crate::execution::local::generate_local_identities_n; use crate::execution::player::{Role, RoleAssignment}; use crate::execution::session::{NetworkSession, Session}; -use crate::network::mpc::handle::control_channel::ControlChannel; +use crate::network::mpc::handle::control_channel::{ControlChannel, MeshControlChannel}; use crate::network::tcp::connection::client::{BoxTcpClient, TcpClient, TlsClient}; use crate::network::tcp::connection::server::{BoxTcpServer, TcpServer, TlsServer}; use crate::network::tcp::{self, TcpStreamConn, TlsClientConfig, TlsConfig, TlsServerConfig}; @@ -40,6 +40,22 @@ pub trait NetworkHandle: Send + Sync { /// channel return an error immediately — there is no retry; call this method /// again to reconnect. async fn control_channel(&mut self) -> Result>; + + /// Establish a dedicated control-plane channel to every other party, + /// addressed by [`Role`] rather than ring position. + /// + /// Generalizes `control_channel` beyond the 3-party ring case: use this + /// for any protocol with more than 3 parties. Same blocking-send, + /// no-retry semantics as `control_channel`. + /// + /// Defaults to an error so existing `NetworkHandle` implementations + /// outside this crate keep compiling without change; only the + /// TCP-backed `MpcNetworkHandle` in this crate currently overrides it. + async fn mesh_control_channel(&mut self) -> Result> { + Err(eyre::eyre!( + "mesh_control_channel is not implemented for this NetworkHandle" + )) + } } pub struct NetworkHandleArgs { @@ -62,7 +78,7 @@ pub async fn build_network_handle( ) -> Result> { tcp::init_rustls_crypto_provider(); - let identities = generate_local_identities(); + let identities = generate_local_identities_n(args.addresses.len()); let role_assignments: RoleAssignment = identities .iter() .enumerate() @@ -210,7 +226,10 @@ pub mod testing { connection_parallelism: usize, request_parallelism: usize, ) -> Result>> { - assert_eq!(parties.len(), 3); + assert!( + parties.len() >= 2, + "MPC networking requires at least 2 parties" + ); let config = MpcConfig::new( Duration::from_secs(30), @@ -220,8 +239,7 @@ pub mod testing { let addresses = get_free_local_addresses(parties.len()).await?; let shutdown_ct = CancellationToken::new(); - let identities = generate_local_identities(); - let role_assignments: RoleAssignment = identities + let role_assignments: RoleAssignment = parties .iter() .enumerate() .map(|(index, id)| (Role::new(index), id.clone())) @@ -303,7 +321,7 @@ mod tests { use tokio::time::sleep; use tracing_test::traced_test; - use crate::execution::local::generate_local_identities; + use crate::execution::local::{generate_local_identities, generate_local_identities_n}; use crate::execution::player::{Identity, Role}; use crate::execution::session::NetworkSession; use crate::network::mpc::NetworkValue; @@ -320,11 +338,12 @@ mod tests { let mut tasks = JoinSet::new(); let message_to_next = get_prf(); let message_to_prev = get_prf(); + let num_parties = identities.len() as u8; for (player_id, session) in sessions.into_iter().enumerate() { let role = Role::new(player_id); - let next = role.next(3).index(); - let prev = role.prev(3).index(); + let next = role.next(num_parties).index(); + let prev = role.prev(num_parties).index(); let next_id = identities[next].clone(); let prev_id = identities[prev].clone(); @@ -436,4 +455,21 @@ mod tests { Ok(()) } + + #[tokio::test(flavor = "multi_thread")] + #[traced_test] + async fn test_mpc_comms_correct_five_parties() -> Result<()> { + let identities = generate_local_identities_n(5); + let (_managers, mut sessions) = + setup_local_mpc_networking(identities.clone(), 1, 1).await?; + sleep(Duration::from_millis(500)).await; + + assert_eq!(sessions.len(), 5); + assert_eq!(sessions[0].len(), 1); + + let players: Vec = sessions.iter_mut().map(|s| s.remove(0)).collect(); + all_parties_talk(identities, players).await; + + Ok(()) + } } diff --git a/ampc-actor-utils/src/network/mpc/handle/network_handle.rs b/ampc-actor-utils/src/network/mpc/handle/network_handle.rs index 2ab7409f..aa91bdda 100644 --- a/ampc-actor-utils/src/network/mpc/handle/network_handle.rs +++ b/ampc-actor-utils/src/network/mpc/handle/network_handle.rs @@ -1,4 +1,4 @@ -use std::collections::HashMap; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use crate::execution::scheduler::parallelize; @@ -12,7 +12,9 @@ use crate::{ mpc::{ handle::{ config::MpcConfig, - control_channel::{ControlChannel, TcpControlChannel}, + control_channel::{ + ControlChannel, MeshControlChannel, TcpControlChannel, TcpMeshControlChannel, + }, data::PeerConnections, session::TcpSession, }, @@ -37,7 +39,7 @@ use tokio_rustls::rustls::pki_types::CertificateDer; use tokio_util::sync::CancellationToken; pub struct MpcNetworkHandle + 'static> { - peers: [Arc; 2], + peers: Vec>, my_id: Arc, connector: Arc, conn_cmd_tx: UnboundedSender>, @@ -144,20 +146,31 @@ impl + 'static> NetworkHan } async fn control_channel(&mut self) -> Result> { + // This is a ring (next/prev) API: it only has a well-defined meaning + // for a 3-party configuration, where each party has exactly two + // peers. For other party counts, use `mesh_control_channel` instead. + if self.peers.len() != 2 { + bail!( + "control_channel() requires exactly 2 peers (a 3-party ring), found {}; \ + use mesh_control_channel() instead", + self.peers.len() + ); + } + let err_ct = CancellationToken::new(); let drop_guard = err_ct.clone().drop_guard(); let connection_state = ConnectionState::new(self.shutdown_ct.clone(), err_ct); - // One connection per peer; peers[0] → c0[0], peers[1] → c1[0]. - let (c0, c1) = self.make_connections(1, connection_state).await?; - let stream_to_peer0 = c0 - .into_iter() - .next() - .ok_or_else(|| eyre!("no connection to peer 0"))?; - let stream_to_peer1 = c1 - .into_iter() - .next() + // One connection per peer; peers[0] → conns[0][0], peers[1] → conns[1][0]. + let mut conns = self.make_connections(1, connection_state).await?; + let stream_to_peer1 = conns + .pop() + .and_then(|mut c| c.pop()) .ok_or_else(|| eyre!("no connection to peer 1"))?; + let stream_to_peer0 = conns + .pop() + .and_then(|mut c| c.pop()) + .ok_or_else(|| eyre!("no connection to peer 0"))?; // Resolve which stream is "next" and which is "prev" using role assignments. let my_role = Role::new(self.party_index); @@ -179,6 +192,57 @@ impl + 'static> NetworkHan self.shutdown_ct.clone(), ))) } + + async fn mesh_control_channel(&mut self) -> Result> { + let err_ct = CancellationToken::new(); + let drop_guard = err_ct.clone().drop_guard(); + let connection_state = ConnectionState::new(self.shutdown_ct.clone(), err_ct); + + // One connection per peer. + let conns = self.make_connections(1, connection_state).await?; + let mut peer_streams: Vec<(Identity, T)> = self + .peers + .iter() + .zip(conns) + .map(|(peer, mut conn)| { + let stream = conn + .pop() + .ok_or_else(|| eyre!("no connection to peer {:?}", peer.id()))?; + Ok((peer.id().clone(), stream)) + }) + .collect::>()?; + + // Resolve each peer's Role via role_assignments (forward lookup by + // .get(), not a scan over the HashMap, since iterating a HashMap's + // order is not guaranteed to match across parties). + let num_parties = self.role_assignments.len() as u8; + let own_role = Role::new(self.party_index); + let mut streams = BTreeMap::new(); + for idx in 0..num_parties { + let role = Role::new(idx as usize); + if role == own_role { + continue; + } + let identity = self + .role_assignments + .get(&role) + .ok_or_else(|| eyre!("role {role:?} not found in role_assignments"))?; + let pos = peer_streams + .iter() + .position(|(id, _)| id == identity) + .ok_or_else(|| { + eyre!("no connection established for role {role:?} ({identity:?})") + })?; + let (_, stream) = peer_streams.remove(pos); + streams.insert(role, stream); + } + + drop_guard.disarm(); + Ok(Box::new(TcpMeshControlChannel::new( + streams, + self.shutdown_ct.clone(), + ))) + } } impl + 'static> MpcNetworkHandle { @@ -200,14 +264,14 @@ impl + 'static> MpcNetwork { let peers: Vec<_> = peers.collect(); let identities: Vec<_> = peers.iter().map(|(id, _address)| id.0.clone()).collect(); - let mut peers = peers.into_iter(); let tls = build_runtime_tls_config(tls_cfg, &identities, party_index)?; let my_id = Arc::new(my_id); - let peers: [Arc; 2] = [ - Arc::new(peers.next().expect("expected at least 2 identities").into()), - Arc::new(peers.next().expect("expected at least 2 identities").into()), - ]; + assert!( + !peers.is_empty(), + "expected at least 1 peer identity to build a NetworkHandle" + ); + let peers: Vec> = peers.into_iter().map(|p| Arc::new(p.into())).collect(); // use the shutdown_ct to cancel anything spawned by the NetworkHandle. But don't want this to affect the calling code. // Hence the child_token() @@ -236,24 +300,23 @@ impl + 'static> MpcNetwork conns_per_peer: u32, connection_state: ConnectionState, ) -> Result> { - let (c0, c1) = self + let conns = self .make_connections(conns_per_peer, connection_state) .await?; - Ok(PeerConnections::new(self.peers.clone(), c0, c1)) + Ok(PeerConnections::new(self.peers.clone(), conns)) } - // returns the connections for each peer + // returns the connections for each peer, grouped in the same order as + // `self.peers` (i.e. result[i] holds the connections to peers[i]). // when returned, the handshakes have successfully completed async fn make_connections( &self, conns_per_peer: u32, connection_state: ConnectionState, - ) -> Result<(Vec, Vec)> { - assert_eq!(self.peers.len(), 2); + ) -> Result>> { + assert!(!self.peers.is_empty()); let mut connect_futures = Vec::with_capacity(conns_per_peer as usize * self.peers.len()); - // peers[0] will be associated with connections c0 - // peers[1] will be associated with connections c1 for peer in self.peers.iter() { for idx in 0..conns_per_peer { let connection_id = ConnectionId::new(idx); @@ -275,12 +338,15 @@ impl + 'static> MpcNetwork .into_iter() .collect::, _>>()?; - let mut c1 = results; - let mut c0 = vec![]; - c0.extend(c1.drain(0..c1.len() / 2)); - assert_eq!(c1.len(), c0.len()); + // results is ordered [peer0_conn0..peer0_connN, peer1_conn0..peer1_connN, ...]; + // regroup it into one Vec per peer, in peer order. + let conns_per_peer = conns_per_peer as usize; + let mut results = results.into_iter(); + let per_peer: Vec> = (0..self.peers.len()) + .map(|_| results.by_ref().take(conns_per_peer).collect()) + .collect(); - Ok((c0, c1)) + Ok(per_peer) } async fn validate_sessions(&self, sessions: &mut [TcpSession]) -> Result<()> { @@ -380,20 +446,24 @@ mod tests { use tokio_util::sync::CancellationToken; use tracing_test::traced_test; - use crate::execution::local::generate_local_identities; - use crate::network::tcp::{ConnectionState, TcpStreamConn}; + use crate::execution::local::generate_local_identities_n; + use crate::network::tcp::ConnectionState; - #[tokio::test(flavor = "multi_thread")] - #[traced_test] - async fn test_tcp_network_handle() -> Result<()> { + /// Establishes `num_parties` TCP handles (each connected to every other + /// party) and has every party exchange raw bytes with every peer over + /// every connection, to verify the N-peer connection-establishment path + /// (not just the 3-party ring case). + async fn run_tcp_network_handle_test(num_parties: u8) -> Result<()> { const CONNECTIONS_PER_PEER: u32 = 2; - let identities = generate_local_identities(); + let identities = generate_local_identities_n(num_parties as usize); let handles = get_local_mpc_handles(identities, CONNECTIONS_PER_PEER as usize, 1).await?; let cs = ConnectionState::new(CancellationToken::new(), CancellationToken::new()); - // for each peer, a vec of the connections to other peers - let connections: Vec<(Vec, Vec)> = futures::future::join_all( + // for each handle (party), one Vec per peer, in the + // same order as that party's `self.peers` (i.e. global index order, + // skipping its own index). + let connections = futures::future::join_all( handles .iter() .map(|h| h.make_connections(CONNECTIONS_PER_PEER, cs.clone())), @@ -403,43 +473,49 @@ mod tests { .collect::, _>>()?; let mut jobs = JoinSet::new(); - let peer_ids: [u8; 3] = [0, 1, 2]; + let peer_ids: Vec = (0..num_parties).collect(); tracing::debug!("connections created. sending data"); - for (peer_idx, (p0, p1)) in connections.into_iter().enumerate() { - let my_peer_ids = peer_ids + for (peer_idx, per_peer_conns) in connections.into_iter().enumerate() { + let other_ids: Vec = peer_ids .iter() - .enumerate() - .filter(|(idx, _)| *idx != peer_idx) - .map(|(_, id)| *id) - .collect::>(); - - for (conn_idx, (mut c0, mut c1)) in p0.into_iter().zip(p1.into_iter()).enumerate() { - let p0_data = [my_peer_ids[0], conn_idx as u8]; - let p1_data = [my_peer_ids[1], conn_idx as u8]; - - jobs.spawn(async move { - c0.write_all(&p0_data).await.unwrap(); - c0.flush().await.unwrap(); - - let mut recv_data = [0u8; 2]; - c0.read_exact(&mut recv_data).await.unwrap(); - assert_eq!(&recv_data, &[peer_idx as u8, conn_idx as u8]); - }); - - jobs.spawn(async move { - c1.write_all(&p1_data).await.unwrap(); - c1.flush().await.unwrap(); - - let mut recv_data = [0u8; 2]; - c1.read_exact(&mut recv_data).await.unwrap(); - assert_eq!(&recv_data, &[peer_idx as u8, conn_idx as u8]); - }); + .copied() + .filter(|id| *id != peer_idx as u8) + .collect(); + assert_eq!(per_peer_conns.len(), other_ids.len()); + + for (other_id, conns) in other_ids.into_iter().zip(per_peer_conns) { + for (conn_idx, mut c) in conns.into_iter().enumerate() { + // send the neighbour's own id, marked with the + // connection index; the neighbour does the same in + // reverse, so what comes back is this party's own id. + let data = [other_id, conn_idx as u8]; + jobs.spawn(async move { + c.write_all(&data).await.unwrap(); + c.flush().await.unwrap(); + + let mut recv_data = [0u8; 2]; + c.read_exact(&mut recv_data).await.unwrap(); + assert_eq!(&recv_data, &[peer_idx as u8, conn_idx as u8]); + }); + } } } jobs.join_all().await; Ok(()) } + + #[tokio::test(flavor = "multi_thread")] + #[traced_test] + async fn test_tcp_network_handle() -> Result<()> { + run_tcp_network_handle_test(3).await + } + + #[tokio::test(flavor = "multi_thread")] + #[traced_test] + async fn test_tcp_network_handle_five_parties() -> Result<()> { + run_tcp_network_handle_test(5).await + } } diff --git a/ampc-actor-utils/src/protocol/ops.rs b/ampc-actor-utils/src/protocol/ops.rs index 30e1d7c8..7d3adae9 100644 --- a/ampc-actor-utils/src/protocol/ops.rs +++ b/ampc-actor-utils/src/protocol/ops.rs @@ -1,10 +1,13 @@ // Protocol operations for MPC // This file contains only the non-iris-specific protocol operations +use crate::execution::player::Role; use crate::execution::session::{NetworkSession, Session, SessionHandles}; use crate::network::mpc::{NetworkInt, NetworkValue}; use crate::protocol::binary::{bit_inject, extract_msb_batch, lift, lift_to_ring48, open_bin}; -use crate::protocol::prf::{Prf, PrfSeed}; +use crate::protocol::prf::{ + PairwisePrfKeys, PartyPair, Prf, PrfSeed, ThresholdPrfKeys, FIVE_PARTY_COUNT, +}; use ampc_secret_sharing::shares::bit::Bit; use ampc_secret_sharing::shares::share::DistanceShare; use ampc_secret_sharing::shares::RingRandFillable; @@ -14,6 +17,7 @@ use ampc_secret_sharing::shares::{ use eyre::{bail, eyre, Result}; use itertools::{izip, Itertools}; use rand_distr::{Distribution, Standard}; +use std::collections::BTreeMap; use tracing::instrument; pub type DistancePair = (DistanceShare, DistanceShare); @@ -115,6 +119,112 @@ pub async fn setup_shared_seed(session: &mut NetworkSession, my_seed: PrfSeed) - Ok(shared_seed) } +fn xor_seeds(a: PrfSeed, b: PrfSeed) -> PrfSeed { + std::array::from_fn(|i| a[i] ^ b[i]) +} + +fn decode_prf_seed(msg: Result, from: Role) -> Result { + match msg { + Ok(NetworkValue::PrfKey(seed)) => Ok(seed), + Ok(_) => Err(eyre!("expected a PrfKey message from {from:?}")), + Err(e) => Err(e), + } +} + +/// Establishes the `(3, 5)` threshold PRF key configuration described by +/// [`ThresholdPrfKeys`]. Requires exactly [`FIVE_PARTY_COUNT`] parties. +/// +/// For each of the six excluded pairs `{i, j}` not containing this party's +/// role, this party and the other two owners each contribute a random seed; +/// the three contributions are exchanged and XORed together, so no single +/// owner controls the resulting key. +/// +/// Must run to completion, on every party, before any other message is sent +/// on `session` — and before [`setup_pairwise_prf_keys`] is called on the +/// same session — since both exchange bare, uncorrelated `PrfKey` messages. +#[instrument( + level = "trace", + target = "mpc::network", + fields(party = ?session.own_role), + skip_all +)] +pub async fn setup_threshold_prf_keys(session: &mut NetworkSession) -> Result { + let own_role = session.own_role(); + let num_parties = session.role_assignments.len(); + if num_parties != FIVE_PARTY_COUNT as usize { + bail!( + "threshold PRF key setup requires exactly {FIVE_PARTY_COUNT} parties, found {num_parties}" + ); + } + + let mut seeds = BTreeMap::new(); + for pair in PartyPair::excluding(own_role) { + let (a, b) = pair.parties(); + let co_owners: Vec = (0..FIVE_PARTY_COUNT) + .map(|i| Role::new(i as usize)) + .filter(|role| *role != a && *role != b && *role != own_role) + .collect(); + debug_assert_eq!(co_owners.len(), 2); + + let my_seed = Prf::gen_seed(); + for co_owner in &co_owners { + session + .send_to(NetworkValue::PrfKey(my_seed), co_owner) + .await?; + } + + let mut combined = my_seed; + for co_owner in &co_owners { + let seed = decode_prf_seed(session.receive_from(co_owner).await, *co_owner)?; + combined = xor_seeds(combined, seed); + } + seeds.insert(pair, combined); + } + + ThresholdPrfKeys::from_seeds(own_role, seeds) +} + +/// Establishes one pairwise PRF key per other party, as described by +/// [`PairwisePrfKeys`]. Requires exactly [`FIVE_PARTY_COUNT`] parties. +/// +/// This party and each other party contribute a random seed for their +/// shared key; the two contributions are XORed together so neither party +/// alone controls the result. +/// +/// Must run to completion, on every party, before any other message is sent +/// on `session` — and before [`setup_threshold_prf_keys`] is called on the +/// same session — since both exchange bare, uncorrelated `PrfKey` messages. +#[instrument( + level = "trace", + target = "mpc::network", + fields(party = ?session.own_role), + skip_all +)] +pub async fn setup_pairwise_prf_keys(session: &mut NetworkSession) -> Result { + let own_role = session.own_role(); + let num_parties = session.role_assignments.len(); + if num_parties != FIVE_PARTY_COUNT as usize { + bail!( + "pairwise PRF key setup requires exactly {FIVE_PARTY_COUNT} parties, found {num_parties}" + ); + } + + let mut seeds = BTreeMap::new(); + for other in (0..FIVE_PARTY_COUNT) + .map(|i| Role::new(i as usize)) + .filter(|role| *role != own_role) + { + let my_seed = Prf::gen_seed(); + session + .send_to(NetworkValue::PrfKey(my_seed), &other) + .await?; + let other_seed = decode_prf_seed(session.receive_from(&other).await, other)?; + seeds.insert(other, xor_seeds(my_seed, other_seed)); + } + + PairwisePrfKeys::from_seeds(own_role, seeds) +} + /// Convert Galois Ring elements to replicated secret shares (Rep3) /// This takes a vector of ring elements and converts them to replicated shares #[instrument(level = "trace", target = "searcher::network", skip_all)] @@ -670,7 +780,9 @@ pub async fn batch_signed_lift_vec_ring48( #[cfg(test)] mod tests { use super::*; - use crate::execution::local::{generate_local_identities, LocalRuntime}; + use crate::execution::local::{ + generate_local_identities, generate_local_identities_n, LocalRuntime, + }; use crate::protocol::test_utils::create_array_sharing; use aes_prng::AesRng; use rand::RngCore; @@ -730,6 +842,74 @@ mod tests { ); } + #[tokio::test] + async fn test_setup_threshold_and_pairwise_prf_keys_five_parties() { + let identities = generate_local_identities_n(FIVE_PARTY_COUNT as usize); + let mut seeds = Vec::new(); + for i in 0..FIVE_PARTY_COUNT { + let mut seed = [0_u8; 16]; + seed[0] = i; + seeds.push(seed); + } + let runtime = LocalRuntime::new(identities.clone(), seeds.clone()) + .await + .unwrap(); + + let mut jobs = JoinSet::new(); + for session in runtime.sessions { + jobs.spawn(async move { + let mut network_session = session.network_session; + let threshold = setup_threshold_prf_keys(&mut network_session) + .await + .unwrap(); + let pairwise = setup_pairwise_prf_keys(&mut network_session).await.unwrap(); + (threshold, pairwise) + }); + } + let mut by_role: Vec<(ThresholdPrfKeys, PairwisePrfKeys)> = jobs.join_all().await; + by_role.sort_by_key(|(t, _)| t.own_role().index()); + + for (threshold, pairwise) in &by_role { + let role = threshold.own_role(); + assert_eq!(threshold.owned_pairs().count(), 6); + assert!(threshold.owned_pairs().all(|pair| !pair.contains(role))); + assert_eq!(pairwise.parties().count(), 4); + assert!(pairwise.parties().all(|other| other != role)); + } + + // Every threshold key must be agreed identically by all three owners. + for a in 0..FIVE_PARTY_COUNT { + for b in (a + 1)..FIVE_PARTY_COUNT { + let (role_a, role_b) = (Role::new(a as usize), Role::new(b as usize)); + let mut agreed_value: Option = None; + let mut num_owners = 0; + for (threshold, _) in by_role.iter_mut() { + let Some(rng) = threshold.get_mut(role_a, role_b) else { + continue; + }; + num_owners += 1; + let value = rng.next_u64(); + match agreed_value { + None => agreed_value = Some(value), + Some(expected) => assert_eq!(value, expected), + } + } + assert_eq!(num_owners, 3, "pair ({a},{b}) should have 3 owners"); + } + } + + // Every pairwise key must be agreed identically by both parties. + for a in 0..FIVE_PARTY_COUNT { + for b in (a + 1)..FIVE_PARTY_COUNT { + let role_a = Role::new(a as usize); + let role_b = Role::new(b as usize); + let a_value = by_role[a as usize].1.get_mut(role_b).unwrap().next_u64(); + let b_value = by_role[b as usize].1.get_mut(role_a).unwrap().next_u64(); + assert_eq!(a_value, b_value); + } + } + } + #[tokio::test] async fn test_setup_shared_seed() { let mut seeds = Vec::new(); diff --git a/ampc-actor-utils/src/protocol/prf.rs b/ampc-actor-utils/src/protocol/prf.rs index 5d117810..017abe07 100644 --- a/ampc-actor-utils/src/protocol/prf.rs +++ b/ampc-actor-utils/src/protocol/prf.rs @@ -1,3 +1,4 @@ +use crate::execution::player::Role; use crate::protocol::shuffle::Permutation; use ampc_secret_sharing::shares::{ int_ring::IntRing2k, @@ -5,6 +6,7 @@ use ampc_secret_sharing::shares::{ }; use eyre::{bail, Result}; use rand::{distributions::Standard, prelude::Distribution, Rng, SeedableRng}; +use std::collections::{BTreeMap, BTreeSet}; /// Generate a uniformly random u32 in [0, modulus) fn gen_u32_mod(rng: &mut PrfRng, modulus: u32) -> Result { @@ -48,11 +50,10 @@ impl Default for Prf { } impl Prf { - #[cfg(not(feature = "aes_rng_prf"))] pub fn new(my_key: PrfSeed, prev_key: PrfSeed) -> Self { Self { - my_prf: PrfRng::from_seed(Self::expand_seed(my_key)), - prev_prf: PrfRng::from_seed(Self::expand_seed(prev_key)), + my_prf: seed_to_rng(my_key), + prev_prf: seed_to_rng(prev_key), } } @@ -67,14 +68,6 @@ impl Prf { out } - #[cfg(feature = "aes_rng_prf")] - pub fn new(my_key: PrfSeed, prev_key: PrfSeed) -> Self { - Self { - my_prf: PrfRng::from_seed(my_key), - prev_prf: PrfRng::from_seed(prev_key), - } - } - #[inline(always)] pub fn get_my_prf(&mut self) -> &mut PrfRng { &mut self.my_prf @@ -165,6 +158,186 @@ impl Prf { } } +#[inline] +fn seed_to_rng(seed: PrfSeed) -> PrfRng { + #[cfg(not(feature = "aes_rng_prf"))] + { + PrfRng::from_seed(Prf::expand_seed(seed)) + } + #[cfg(feature = "aes_rng_prf")] + { + PrfRng::from_seed(seed) + } +} + +/// Number of parties in the 5-party protocol configuration used by +/// [`ThresholdPrfKeys`] and [`PairwisePrfKeys`]. +pub const FIVE_PARTY_COUNT: u8 = 5; + +fn five_party_roles() -> [Role; FIVE_PARTY_COUNT as usize] { + std::array::from_fn(Role::new) +} + +/// An unordered, canonically-ordered pair of two distinct parties. +/// +/// `PartyPair::new(a, b) == PartyPair::new(b, a)`, so it can be used as a +/// map key representing the pair `{a, b}` regardless of which order the two +/// roles are known in. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct PartyPair(Role, Role); + +impl PartyPair { + pub fn new(a: Role, b: Role) -> Self { + assert_ne!(a, b, "a PartyPair requires two distinct roles"); + if a.index() <= b.index() { + Self(a, b) + } else { + Self(b, a) + } + } + + pub fn contains(&self, role: Role) -> bool { + self.0 == role || self.1 == role + } + + pub fn parties(&self) -> (Role, Role) { + (self.0, self.1) + } + + /// All unordered pairs among the five parties that do not contain + /// `own_role`. There are `C(4, 2) = 6` of them. + pub fn excluding(own_role: Role) -> Vec { + let roles = five_party_roles(); + let mut pairs = Vec::with_capacity(6); + for (idx, &a) in roles.iter().enumerate() { + if a == own_role { + continue; + } + for &b in &roles[idx + 1..] { + if b == own_role { + continue; + } + pairs.push(PartyPair::new(a, b)); + } + } + pairs + } +} + +/// A `(3, 5)` shared-PRF key configuration for a 5-party protocol. +/// +/// For every unordered pair of parties `{i, j}` there is one key `k_{i,j}`, +/// known only to the three parties *not* in `{i, j}`. Each party therefore +/// owns keys for the `C(4, 2) = 6` pairs that exclude it. This is a +/// generalization, to five parties, of the same trick that the replicated +/// (3-party) [`Prf`] above is built on: a value only the "other" parties can +/// predict is exactly what is needed to mask/rerandomize shares that those +/// parties, and not the excluded pair, must be able to jointly verify or +/// reconstruct. +/// +/// Construct one with [`from_seeds`](Self::from_seeds) once the underlying +/// seeds have been agreed with the other owners of each key (see +/// `setup_threshold_prf_keys` in `protocol::ops`). +#[derive(Debug)] +pub struct ThresholdPrfKeys { + own_role: Role, + keys: BTreeMap, +} + +impl ThresholdPrfKeys { + /// Build the key set from one already-agreed seed per owned pair. + /// + /// Fails unless `seeds` contains exactly the six pairs returned by + /// [`PartyPair::excluding(own_role)`](PartyPair::excluding). + pub fn from_seeds(own_role: Role, seeds: BTreeMap) -> Result { + let expected: BTreeSet = PartyPair::excluding(own_role).into_iter().collect(); + let actual: BTreeSet = seeds.keys().copied().collect(); + if actual != expected { + bail!( + "threshold PRF key set for role {own_role:?} does not match the expected \ + 3-of-5 key set: expected {expected:?}, got {actual:?}" + ); + } + let keys = seeds + .into_iter() + .map(|(pair, seed)| (pair, seed_to_rng(seed))) + .collect(); + Ok(Self { own_role, keys }) + } + + pub fn own_role(&self) -> Role { + self.own_role + } + + /// The PRF for key `k_{i,j}`. Returns `None` if this party does not own + /// that key (i.e. `own_role` is `i` or `j`) or if `i == j`. + pub fn get_mut(&mut self, i: Role, j: Role) -> Option<&mut PrfRng> { + if i == j { + return None; + } + self.keys.get_mut(&PartyPair::new(i, j)) + } + + /// The excluded pairs whose key this party holds. + pub fn owned_pairs(&self) -> impl Iterator + '_ { + self.keys.keys().copied() + } +} + +/// Pairwise shared-PRF keys: for every unordered pair of parties `{i, j}`, +/// a key `\lambda_{i,j}` known only to `i` and `j` themselves. Each party +/// owns one such key per other party. +/// +/// Construct one with [`from_seeds`](Self::from_seeds) once the underlying +/// seed has been agreed with the other party of each pair (see +/// `setup_pairwise_prf_keys` in `protocol::ops`). +#[derive(Debug)] +pub struct PairwisePrfKeys { + own_role: Role, + keys: BTreeMap, +} + +impl PairwisePrfKeys { + /// Build the key set from one already-agreed seed per other party. + /// + /// Fails unless `seeds` contains exactly one entry per other party in + /// the 5-party configuration (i.e. `0..FIVE_PARTY_COUNT`, excluding + /// `own_role`). + pub fn from_seeds(own_role: Role, seeds: BTreeMap) -> Result { + let expected: BTreeSet = five_party_roles() + .into_iter() + .filter(|role| *role != own_role) + .collect(); + let actual: BTreeSet = seeds.keys().copied().collect(); + if actual != expected { + bail!( + "pairwise PRF key set for role {own_role:?} does not match the expected \ + key set: expected {expected:?}, got {actual:?}" + ); + } + let keys = seeds + .into_iter() + .map(|(role, seed)| (role, seed_to_rng(seed))) + .collect(); + Ok(Self { own_role, keys }) + } + + pub fn own_role(&self) -> Role { + self.own_role + } + + /// The PRF for key `\lambda_{own_role, other}`. Returns `None` if + /// `other` is not a party this key set was built for. + pub fn get_mut(&mut self, other: Role) -> Option<&mut PrfRng> { + self.keys.get_mut(&other) + } + + /// The other parties this party shares a pairwise key with. + pub fn parties(&self) -> impl Iterator + '_ { + self.keys.keys().copied() + } +} + #[cfg(test)] mod tests { use std::collections::HashMap; @@ -256,4 +429,63 @@ mod tests { Ok(()) } + + #[test] + fn test_party_pair_excluding_gives_six_disjoint_pairs() { + for own in 0..FIVE_PARTY_COUNT { + let own_role = Role::new(own as usize); + let pairs = PartyPair::excluding(own_role); + assert_eq!(pairs.len(), 6); + assert!(pairs.iter().all(|pair| !pair.contains(own_role))); + + let unique: std::collections::HashSet<_> = pairs.iter().collect(); + assert_eq!(unique.len(), 6, "pairs must be pairwise distinct"); + } + + // Every one of the 10 possible pairs is owned by exactly 3 of the 5 parties. + let mut owner_counts: HashMap = HashMap::new(); + for own in 0..FIVE_PARTY_COUNT { + for pair in PartyPair::excluding(Role::new(own as usize)) { + *owner_counts.entry(pair).or_insert(0) += 1; + } + } + assert_eq!(owner_counts.len(), 10); + assert!(owner_counts.values().all(|&count| count == 3)); + } + + #[test] + fn test_threshold_prf_keys_from_seeds_rejects_wrong_key_set() { + let own_role = Role::new(0); + // Missing key set (empty) should be rejected. + assert!(ThresholdPrfKeys::from_seeds(own_role, BTreeMap::new()).is_err()); + + // A key set that includes a pair containing own_role should be rejected. + let mut seeds: BTreeMap = PartyPair::excluding(Role::new(1)) + .into_iter() + .map(|pair| (pair, [0u8; 16])) + .collect(); + assert!(ThresholdPrfKeys::from_seeds(own_role, seeds.clone()).is_err()); + + // The correct key set is accepted. + seeds = PartyPair::excluding(own_role) + .into_iter() + .map(|pair| (pair, [0u8; 16])) + .collect(); + assert!(ThresholdPrfKeys::from_seeds(own_role, seeds).is_ok()); + } + + #[test] + fn test_pairwise_prf_keys_from_seeds_rejects_wrong_key_set() { + let own_role = Role::new(0); + assert!(PairwisePrfKeys::from_seeds(own_role, BTreeMap::new()).is_err()); + + // Including a key for `own_role` itself should be rejected. + let mut seeds: BTreeMap = (0..FIVE_PARTY_COUNT) + .map(|i| (Role::new(i as usize), [0u8; 16])) + .collect(); + assert!(PairwisePrfKeys::from_seeds(own_role, seeds.clone()).is_err()); + + seeds.remove(&own_role); + assert!(PairwisePrfKeys::from_seeds(own_role, seeds).is_ok()); + } } diff --git a/ampc-actor-utils/tests/mesh_control_channel.rs b/ampc-actor-utils/tests/mesh_control_channel.rs new file mode 100644 index 00000000..a87b0d07 --- /dev/null +++ b/ampc-actor-utils/tests/mesh_control_channel.rs @@ -0,0 +1,141 @@ +//! Integration tests for the N-party mesh control-plane channel. +//! +//! Builds five `NetworkHandle`s over plain TCP (no TLS) and exercises +//! `MeshControlChannel`'s Role-addressed and barrier APIs. Unlike +//! `ControlChannel` (a 3-party ring: next/prev), `MeshControlChannel` +//! addresses every other party directly by `Role`, so it works for any +//! number of parties. + +use ampc_actor_utils::execution::player::Role; +use ampc_actor_utils::network::mpc::{ + build_network_handle, NetworkHandle, NetworkHandleArgs, NetworkValue, +}; +use futures::future::join_all; +use tokio_util::sync::CancellationToken; +use tracing_test::traced_test; + +const NUM_PARTIES: usize = 5; + +/// Bind `n` listeners simultaneously so the OS assigns `n` distinct free ports, +/// then return those ports. Holding all listeners alive until all ports are +/// collected prevents the OS from handing out the same port twice. +fn find_free_ports(n: usize) -> Vec { + let listeners: Vec = (0..n) + .map(|_| std::net::TcpListener::bind("127.0.0.1:0").unwrap()) + .collect(); + listeners + .iter() + .map(|l| l.local_addr().unwrap().port()) + .collect() +} + +async fn build_handle( + party_index: usize, + addresses: Vec, + shutdown_ct: CancellationToken, +) -> Box { + build_network_handle( + NetworkHandleArgs { + party_index, + addresses: addresses.clone(), + outbound_addresses: addresses, + connection_parallelism: 1, + request_parallelism: 1, + sessions_per_request: 1, + tls: None, + }, + shutdown_ct, + ) + .await + .expect("build_network_handle failed") +} + +#[tokio::test(flavor = "multi_thread")] +#[traced_test] +async fn test_mesh_control_channel_sync() { + let ports = find_free_ports(NUM_PARTIES); + let addresses: Vec = ports.iter().map(|p| format!("127.0.0.1:{p}")).collect(); + + let shutdown_ct = CancellationToken::new(); + + let party_tasks = (0..NUM_PARTIES).map(|party_index| { + let addresses = addresses.clone(); + let shutdown_ct = shutdown_ct.clone(); + + tokio::spawn(async move { + let mut handle = build_handle(party_index, addresses, shutdown_ct).await; + let mut mc = handle + .mesh_control_channel() + .await + .expect("mesh_control_channel() failed"); + + mc.sync().await.expect("sync() failed"); + }) + }); + + let results = join_all(party_tasks).await; + for result in results { + result.expect("party task panicked"); + } +} + +/// Every party sends every other party a payload tagged with its own index, +/// addressed directly by `Role` (not ring position), and confirms it +/// receives back the sender's own tag from each of them. +#[tokio::test(flavor = "multi_thread")] +#[traced_test] +async fn test_mesh_control_channel_send_recv_by_role() { + let ports = find_free_ports(NUM_PARTIES); + let addresses: Vec = ports.iter().map(|p| format!("127.0.0.1:{p}")).collect(); + + let shutdown_ct = CancellationToken::new(); + + let party_tasks = (0..NUM_PARTIES).map(|party_index| { + let addresses = addresses.clone(); + let shutdown_ct = shutdown_ct.clone(); + + tokio::spawn(async move { + let mut handle = build_handle(party_index, addresses, shutdown_ct).await; + let mut mc = handle + .mesh_control_channel() + .await + .expect("mesh_control_channel() failed"); + + for other in 0..NUM_PARTIES { + if other == party_index { + continue; + } + mc.send( + Role::new(other), + NetworkValue::Bytes(vec![party_index as u8].into()), + ) + .await + .unwrap_or_else(|e| panic!("party {party_index}: send to {other} failed: {e}")); + } + + for other in 0..NUM_PARTIES { + if other == party_index { + continue; + } + let received = mc.recv(Role::new(other)).await.unwrap_or_else(|e| { + panic!("party {party_index}: recv from {other} failed: {e}") + }); + match received { + NetworkValue::Bytes(b) => assert_eq!( + &b[..], + &[other as u8], + "party {party_index}: recv from {other} got wrong payload" + ), + other_variant => panic!( + "party {party_index}: unexpected variant {other_variant:?} from {other}" + ), + } + } + }) + }); + + let results = join_all(party_tasks).await; + for result in results { + result.expect("party task panicked"); + } +} From e2aafa310ccd5dedee1dad9cddc91d90ac7e66fc Mon Sep 17 00:00:00 2001 From: Makis Arsenis Date: Tue, 1 Sep 2026 15:06:15 -0700 Subject: [PATCH 2/3] Rename FIVE_PARTY_COUNT to ORBIT5_PARTY_COUNT and reuse orbit5_roles() ORBIT5 is the internal name of the 5-party system. Also makes the roles helper public and replaces the manual (0..N).map(Role::new) constructions in the PRF key setup functions and tests with it. Co-Authored-By: Claude Fable 5 --- ampc-actor-utils/src/protocol/ops.rs | 35 +++++++++++++--------------- ampc-actor-utils/src/protocol/prf.rs | 31 ++++++++++++------------ 2 files changed, 32 insertions(+), 34 deletions(-) diff --git a/ampc-actor-utils/src/protocol/ops.rs b/ampc-actor-utils/src/protocol/ops.rs index 7d3adae9..8fb76243 100644 --- a/ampc-actor-utils/src/protocol/ops.rs +++ b/ampc-actor-utils/src/protocol/ops.rs @@ -6,7 +6,7 @@ use crate::execution::session::{NetworkSession, Session, SessionHandles}; use crate::network::mpc::{NetworkInt, NetworkValue}; use crate::protocol::binary::{bit_inject, extract_msb_batch, lift, lift_to_ring48, open_bin}; use crate::protocol::prf::{ - PairwisePrfKeys, PartyPair, Prf, PrfSeed, ThresholdPrfKeys, FIVE_PARTY_COUNT, + orbit5_roles, PairwisePrfKeys, PartyPair, Prf, PrfSeed, ThresholdPrfKeys, ORBIT5_PARTY_COUNT, }; use ampc_secret_sharing::shares::bit::Bit; use ampc_secret_sharing::shares::share::DistanceShare; @@ -132,7 +132,7 @@ fn decode_prf_seed(msg: Result, from: Role) -> Result { } /// Establishes the `(3, 5)` threshold PRF key configuration described by -/// [`ThresholdPrfKeys`]. Requires exactly [`FIVE_PARTY_COUNT`] parties. +/// [`ThresholdPrfKeys`]. Requires exactly [`ORBIT5_PARTY_COUNT`] parties. /// /// For each of the six excluded pairs `{i, j}` not containing this party's /// role, this party and the other two owners each contribute a random seed; @@ -151,17 +151,17 @@ fn decode_prf_seed(msg: Result, from: Role) -> Result { pub async fn setup_threshold_prf_keys(session: &mut NetworkSession) -> Result { let own_role = session.own_role(); let num_parties = session.role_assignments.len(); - if num_parties != FIVE_PARTY_COUNT as usize { + if num_parties != ORBIT5_PARTY_COUNT as usize { bail!( - "threshold PRF key setup requires exactly {FIVE_PARTY_COUNT} parties, found {num_parties}" + "threshold PRF key setup requires exactly {ORBIT5_PARTY_COUNT} parties, found {num_parties}" ); } let mut seeds = BTreeMap::new(); for pair in PartyPair::excluding(own_role) { let (a, b) = pair.parties(); - let co_owners: Vec = (0..FIVE_PARTY_COUNT) - .map(|i| Role::new(i as usize)) + let co_owners: Vec = orbit5_roles() + .into_iter() .filter(|role| *role != a && *role != b && *role != own_role) .collect(); debug_assert_eq!(co_owners.len(), 2); @@ -185,7 +185,7 @@ pub async fn setup_threshold_prf_keys(session: &mut NetworkSession) -> Result Result Result { let own_role = session.own_role(); let num_parties = session.role_assignments.len(); - if num_parties != FIVE_PARTY_COUNT as usize { + if num_parties != ORBIT5_PARTY_COUNT as usize { bail!( - "pairwise PRF key setup requires exactly {FIVE_PARTY_COUNT} parties, found {num_parties}" + "pairwise PRF key setup requires exactly {ORBIT5_PARTY_COUNT} parties, found {num_parties}" ); } let mut seeds = BTreeMap::new(); - for other in (0..FIVE_PARTY_COUNT) - .map(|i| Role::new(i as usize)) - .filter(|role| *role != own_role) - { + for other in orbit5_roles().into_iter().filter(|role| *role != own_role) { let my_seed = Prf::gen_seed(); session .send_to(NetworkValue::PrfKey(my_seed), &other) @@ -844,9 +841,9 @@ mod tests { #[tokio::test] async fn test_setup_threshold_and_pairwise_prf_keys_five_parties() { - let identities = generate_local_identities_n(FIVE_PARTY_COUNT as usize); + let identities = generate_local_identities_n(ORBIT5_PARTY_COUNT as usize); let mut seeds = Vec::new(); - for i in 0..FIVE_PARTY_COUNT { + for i in 0..ORBIT5_PARTY_COUNT { let mut seed = [0_u8; 16]; seed[0] = i; seeds.push(seed); @@ -878,8 +875,8 @@ mod tests { } // Every threshold key must be agreed identically by all three owners. - for a in 0..FIVE_PARTY_COUNT { - for b in (a + 1)..FIVE_PARTY_COUNT { + for a in 0..ORBIT5_PARTY_COUNT { + for b in (a + 1)..ORBIT5_PARTY_COUNT { let (role_a, role_b) = (Role::new(a as usize), Role::new(b as usize)); let mut agreed_value: Option = None; let mut num_owners = 0; @@ -899,8 +896,8 @@ mod tests { } // Every pairwise key must be agreed identically by both parties. - for a in 0..FIVE_PARTY_COUNT { - for b in (a + 1)..FIVE_PARTY_COUNT { + for a in 0..ORBIT5_PARTY_COUNT { + for b in (a + 1)..ORBIT5_PARTY_COUNT { let role_a = Role::new(a as usize); let role_b = Role::new(b as usize); let a_value = by_role[a as usize].1.get_mut(role_b).unwrap().next_u64(); diff --git a/ampc-actor-utils/src/protocol/prf.rs b/ampc-actor-utils/src/protocol/prf.rs index 017abe07..1f0835c4 100644 --- a/ampc-actor-utils/src/protocol/prf.rs +++ b/ampc-actor-utils/src/protocol/prf.rs @@ -170,11 +170,12 @@ fn seed_to_rng(seed: PrfSeed) -> PrfRng { } } -/// Number of parties in the 5-party protocol configuration used by +/// Number of parties in the ORBIT5 (5-party) protocol configuration used by /// [`ThresholdPrfKeys`] and [`PairwisePrfKeys`]. -pub const FIVE_PARTY_COUNT: u8 = 5; +pub const ORBIT5_PARTY_COUNT: u8 = 5; -fn five_party_roles() -> [Role; FIVE_PARTY_COUNT as usize] { +/// The roles of all ORBIT5 parties, in index order. +pub fn orbit5_roles() -> [Role; ORBIT5_PARTY_COUNT as usize] { std::array::from_fn(Role::new) } @@ -204,10 +205,10 @@ impl PartyPair { (self.0, self.1) } - /// All unordered pairs among the five parties that do not contain - /// `own_role`. There are `C(4, 2) = 6` of them. + /// All unordered pairs among the five ORBIT5 parties that do not + /// contain `own_role`. There are `C(4, 2) = 6` of them. pub fn excluding(own_role: Role) -> Vec { - let roles = five_party_roles(); + let roles = orbit5_roles(); let mut pairs = Vec::with_capacity(6); for (idx, &a) in roles.iter().enumerate() { if a == own_role { @@ -224,7 +225,7 @@ impl PartyPair { } } -/// A `(3, 5)` shared-PRF key configuration for a 5-party protocol. +/// A `(3, 5)` shared-PRF key configuration for the ORBIT5 (5-party) protocol. /// /// For every unordered pair of parties `{i, j}` there is one key `k_{i,j}`, /// known only to the three parties *not* in `{i, j}`. Each party therefore @@ -301,10 +302,10 @@ impl PairwisePrfKeys { /// Build the key set from one already-agreed seed per other party. /// /// Fails unless `seeds` contains exactly one entry per other party in - /// the 5-party configuration (i.e. `0..FIVE_PARTY_COUNT`, excluding + /// the ORBIT5 configuration (i.e. `0..ORBIT5_PARTY_COUNT`, excluding /// `own_role`). pub fn from_seeds(own_role: Role, seeds: BTreeMap) -> Result { - let expected: BTreeSet = five_party_roles() + let expected: BTreeSet = orbit5_roles() .into_iter() .filter(|role| *role != own_role) .collect(); @@ -432,8 +433,7 @@ mod tests { #[test] fn test_party_pair_excluding_gives_six_disjoint_pairs() { - for own in 0..FIVE_PARTY_COUNT { - let own_role = Role::new(own as usize); + for own_role in orbit5_roles() { let pairs = PartyPair::excluding(own_role); assert_eq!(pairs.len(), 6); assert!(pairs.iter().all(|pair| !pair.contains(own_role))); @@ -444,8 +444,8 @@ mod tests { // Every one of the 10 possible pairs is owned by exactly 3 of the 5 parties. let mut owner_counts: HashMap = HashMap::new(); - for own in 0..FIVE_PARTY_COUNT { - for pair in PartyPair::excluding(Role::new(own as usize)) { + for own_role in orbit5_roles() { + for pair in PartyPair::excluding(own_role) { *owner_counts.entry(pair).or_insert(0) += 1; } } @@ -480,8 +480,9 @@ mod tests { assert!(PairwisePrfKeys::from_seeds(own_role, BTreeMap::new()).is_err()); // Including a key for `own_role` itself should be rejected. - let mut seeds: BTreeMap = (0..FIVE_PARTY_COUNT) - .map(|i| (Role::new(i as usize), [0u8; 16])) + let mut seeds: BTreeMap = orbit5_roles() + .into_iter() + .map(|role| (role, [0u8; 16])) .collect(); assert!(PairwisePrfKeys::from_seeds(own_role, seeds.clone()).is_err()); From dbff6102f15336c53d1a78dc4174d46669cd0434 Mon Sep 17 00:00:00 2001 From: Makis Arsenis Date: Tue, 1 Sep 2026 16:01:23 -0700 Subject: [PATCH 3/3] Replace generate_local_identities_n with a hardcoded ORBIT5 variant Restores the original 3-party generate_local_identities() and adds generate_local_identities_orbit5() with the five party names hardcoded, instead of slicing a shared list by count. build_network_handle now dispatches on the address count and rejects anything other than 3 or 5 parties explicitly. Co-Authored-By: Claude Fable 5 --- ampc-actor-utils/src/execution/local.rs | 31 ++++++++++--------- .../src/network/mpc/handle/mod.rs | 13 +++++--- .../src/network/mpc/handle/network_handle.rs | 13 ++++---- ampc-actor-utils/src/protocol/ops.rs | 4 +-- 4 files changed, 34 insertions(+), 27 deletions(-) diff --git a/ampc-actor-utils/src/execution/local.rs b/ampc-actor-utils/src/execution/local.rs index fb9f7028..700e6cec 100644 --- a/ampc-actor-utils/src/execution/local.rs +++ b/ampc-actor-utils/src/execution/local.rs @@ -20,24 +20,25 @@ use std::{ }; use tokio::{sync::Mutex, task::JoinHandle}; -const LOCAL_IDENTITY_NAMES: [&str; 5] = ["alice", "bob", "charlie", "dave", "erin"]; - pub fn generate_local_identities() -> Vec { - generate_local_identities_n(3) + vec![ + Identity::from("alice"), + Identity::from("bob"), + Identity::from("charlie"), + ] } -/// Generate `n` deterministic, fixed local identities, for use both in tests -/// and to derive a consistent per-index self-identity across parties in -/// production (see `build_network_handle`). Supports up to 5 parties. -pub fn generate_local_identities_n(n: usize) -> Vec { - assert!( - n <= LOCAL_IDENTITY_NAMES.len(), - "not enough predefined local identities for {n} parties" - ); - LOCAL_IDENTITY_NAMES[..n] - .iter() - .map(|name| Identity::from(*name)) - .collect() +/// The fixed identities of the five ORBIT5 parties, in role order. Used both +/// in tests and to derive a consistent per-index self-identity across +/// parties in production (see `build_network_handle`). +pub fn generate_local_identities_orbit5() -> Vec { + vec![ + Identity::from("alice"), + Identity::from("bob"), + Identity::from("charlie"), + Identity::from("dave"), + Identity::from("erin"), + ] } static USED_PORTS: LazyLock>> = LazyLock::new(|| Mutex::new(HashSet::new())); diff --git a/ampc-actor-utils/src/network/mpc/handle/mod.rs b/ampc-actor-utils/src/network/mpc/handle/mod.rs index 168a6dfc..cfab0072 100644 --- a/ampc-actor-utils/src/network/mpc/handle/mod.rs +++ b/ampc-actor-utils/src/network/mpc/handle/mod.rs @@ -10,13 +10,14 @@ use std::time::Duration; use self::config::MpcConfig; use self::network_handle::MpcNetworkHandle; -use crate::execution::local::generate_local_identities_n; +use crate::execution::local::{generate_local_identities, generate_local_identities_orbit5}; use crate::execution::player::{Role, RoleAssignment}; use crate::execution::session::{NetworkSession, Session}; use crate::network::mpc::handle::control_channel::{ControlChannel, MeshControlChannel}; use crate::network::tcp::connection::client::{BoxTcpClient, TcpClient, TlsClient}; use crate::network::tcp::connection::server::{BoxTcpServer, TcpServer, TlsServer}; use crate::network::tcp::{self, TcpStreamConn, TlsClientConfig, TlsConfig, TlsServerConfig}; +use crate::protocol::prf::ORBIT5_PARTY_COUNT; use async_trait::async_trait; use eyre::Result; use itertools::izip; @@ -78,7 +79,11 @@ pub async fn build_network_handle( ) -> Result> { tcp::init_rustls_crypto_provider(); - let identities = generate_local_identities_n(args.addresses.len()); + let identities = match args.addresses.len() { + 3 => generate_local_identities(), + n if n == ORBIT5_PARTY_COUNT as usize => generate_local_identities_orbit5(), + n => eyre::bail!("unsupported party count {n}: expected 3 or {ORBIT5_PARTY_COUNT}"), + }; let role_assignments: RoleAssignment = identities .iter() .enumerate() @@ -321,7 +326,7 @@ mod tests { use tokio::time::sleep; use tracing_test::traced_test; - use crate::execution::local::{generate_local_identities, generate_local_identities_n}; + use crate::execution::local::{generate_local_identities, generate_local_identities_orbit5}; use crate::execution::player::{Identity, Role}; use crate::execution::session::NetworkSession; use crate::network::mpc::NetworkValue; @@ -459,7 +464,7 @@ mod tests { #[tokio::test(flavor = "multi_thread")] #[traced_test] async fn test_mpc_comms_correct_five_parties() -> Result<()> { - let identities = generate_local_identities_n(5); + let identities = generate_local_identities_orbit5(); let (_managers, mut sessions) = setup_local_mpc_networking(identities.clone(), 1, 1).await?; sleep(Duration::from_millis(500)).await; diff --git a/ampc-actor-utils/src/network/mpc/handle/network_handle.rs b/ampc-actor-utils/src/network/mpc/handle/network_handle.rs index aa91bdda..e40aec3c 100644 --- a/ampc-actor-utils/src/network/mpc/handle/network_handle.rs +++ b/ampc-actor-utils/src/network/mpc/handle/network_handle.rs @@ -446,17 +446,18 @@ mod tests { use tokio_util::sync::CancellationToken; use tracing_test::traced_test; - use crate::execution::local::generate_local_identities_n; + use crate::execution::local::{generate_local_identities, generate_local_identities_orbit5}; + use crate::execution::player::Identity; use crate::network::tcp::ConnectionState; - /// Establishes `num_parties` TCP handles (each connected to every other + /// Establishes one TCP handle per identity (each connected to every other /// party) and has every party exchange raw bytes with every peer over /// every connection, to verify the N-peer connection-establishment path /// (not just the 3-party ring case). - async fn run_tcp_network_handle_test(num_parties: u8) -> Result<()> { + async fn run_tcp_network_handle_test(identities: Vec) -> Result<()> { const CONNECTIONS_PER_PEER: u32 = 2; - let identities = generate_local_identities_n(num_parties as usize); + let num_parties = identities.len() as u8; let handles = get_local_mpc_handles(identities, CONNECTIONS_PER_PEER as usize, 1).await?; let cs = ConnectionState::new(CancellationToken::new(), CancellationToken::new()); @@ -510,12 +511,12 @@ mod tests { #[tokio::test(flavor = "multi_thread")] #[traced_test] async fn test_tcp_network_handle() -> Result<()> { - run_tcp_network_handle_test(3).await + run_tcp_network_handle_test(generate_local_identities()).await } #[tokio::test(flavor = "multi_thread")] #[traced_test] async fn test_tcp_network_handle_five_parties() -> Result<()> { - run_tcp_network_handle_test(5).await + run_tcp_network_handle_test(generate_local_identities_orbit5()).await } } diff --git a/ampc-actor-utils/src/protocol/ops.rs b/ampc-actor-utils/src/protocol/ops.rs index 8fb76243..c2169f0d 100644 --- a/ampc-actor-utils/src/protocol/ops.rs +++ b/ampc-actor-utils/src/protocol/ops.rs @@ -778,7 +778,7 @@ pub async fn batch_signed_lift_vec_ring48( mod tests { use super::*; use crate::execution::local::{ - generate_local_identities, generate_local_identities_n, LocalRuntime, + generate_local_identities, generate_local_identities_orbit5, LocalRuntime, }; use crate::protocol::test_utils::create_array_sharing; use aes_prng::AesRng; @@ -841,7 +841,7 @@ mod tests { #[tokio::test] async fn test_setup_threshold_and_pairwise_prf_keys_five_parties() { - let identities = generate_local_identities_n(ORBIT5_PARTY_COUNT as usize); + let identities = generate_local_identities_orbit5(); let mut seeds = Vec::new(); for i in 0..ORBIT5_PARTY_COUNT { let mut seed = [0_u8; 16];