Add 5-party networking and PRF key setup to ampc-actor-utils

ampc-actor-utils/src/execution/local.rs

@@ -28,6 +28,19 @@ pub fn generate_local_identities() -> Vec<Identity> {
     ]
 }
 
+/// 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<Identity> {
+    vec![
+        Identity::from("alice"),
+        Identity::from("bob"),
+        Identity::from("charlie"),
+        Identity::from("dave"),
+        Identity::from("erin"),
+    ]
+}
+
 static USED_PORTS: LazyLock<Mutex<HashSet<u16>>> = LazyLock::new(|| Mutex::new(HashSet::new()));
 
 pub async fn get_free_local_addresses(num_ports: usize) -> Result<Vec<String>> {

ampc-actor-utils/src/execution/player.rs

@@ -30,7 +30,7 @@ impl From<String> 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 {

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<NetworkValue> {
+        let identity = self.identity(role)?.clone();
+        self.receive(&identity).await
+    }
 }
 
 // Helper methods for sending and receiving VecRingElement<T>.

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<T: NetworkConnection> ControlChannel for TcpControlChannel<T> {
         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<NetworkValue>;
+
+    /// 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<T: NetworkConnection> {
+    streams: BTreeMap<Role, T>,
+    shutdown_ct: CancellationToken,
+}
+
+impl<T: NetworkConnection> TcpMeshControlChannel<T> {
+    pub(super) fn new(streams: BTreeMap<Role, T>, shutdown_ct: CancellationToken) -> Self {
+        Self {
+            streams,
+            shutdown_ct,
+        }
+    }
+}
+
+#[async_trait]
+impl<T: NetworkConnection> MeshControlChannel for TcpMeshControlChannel<T> {
+    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<NetworkValue> {
+        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<Role> = 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(())
+    }
+}

ampc-actor-utils/src/network/mpc/handle/data.rs

@@ -14,14 +14,19 @@ pub type OutStream = mpsc::UnboundedSender<OutboundMsg>;
 pub type InStream = mpsc::UnboundedReceiver<NetworkValue>;
 
 pub struct PeerConnections<T: NetworkConnection + 'static> {
-    peers: [Arc<Peer>; 2],
-    c0: Vec<T>,
-    c1: Vec<T>,
+    peers: Vec<Arc<Peer>>,
+    // conns[i] holds the connections established with peers[i]
+    conns: Vec<Vec<T>>,
 }
 
 impl<T: NetworkConnection + 'static> PeerConnections<T> {
-    pub fn new(peers: [Arc<Peer>; 2], c0: Vec<T>, c1: Vec<T>) -> Self {
-        Self { peers, c0, c1 }
+    pub fn new(peers: Vec<Arc<Peer>>, conns: Vec<Vec<T>>) -> Self {
+        assert_eq!(
+            peers.len(),
+            conns.len(),
+            "expected one connection group per peer"
+        );
+        Self { peers, conns }
     }
 
     pub fn peer_ids(&self) -> Vec<Identity> {
@@ -34,10 +39,11 @@ impl<T: NetworkConnection + 'static> IntoIterator for PeerConnections<T> {
     type IntoIter = std::vec::IntoIter<(Identity, Vec<T>)>;
 
     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::<Vec<_>>()
+            .into_iter()
     }
 }

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;
+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;
+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;
@@ -40,6 +41,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<Box<dyn ControlChannel>>;
+
+    /// 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<Box<dyn MeshControlChannel>> {
+        Err(eyre::eyre!(
+            "mesh_control_channel is not implemented for this NetworkHandle"
+        ))
+    }
 }
 
 pub struct NetworkHandleArgs {
@@ -62,7 +79,11 @@ pub async fn build_network_handle(
 ) -> Result<Box<dyn NetworkHandle>> {
     tcp::init_rustls_crypto_provider();
 
-    let identities = generate_local_identities();
+    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()
@@ -210,7 +231,10 @@ pub mod testing {
         connection_parallelism: usize,
         request_parallelism: usize,
     ) -> Result<Vec<MpcNetworkHandle<TcpStreamConn, TcpClient>>> {
-        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 +244,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 +326,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_orbit5};
     use crate::execution::player::{Identity, Role};
     use crate::execution::session::NetworkSession;
     use crate::network::mpc::NetworkValue;
@@ -320,11 +343,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 +460,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_orbit5();
+        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<NetworkSession> = sessions.iter_mut().map(|s| s.remove(0)).collect();
+        all_parties_talk(identities, players).await;
+
+        Ok(())
+    }
 }