Skip to main content

mpc_ristretto/
network.rs

1//! The `network` module defines abstractions of the transport used to
2//! communicate during the course of an MPC
3mod cert_verifier;
4mod config;
5pub mod dummy_network;
6
7use async_trait::async_trait;
8use bytes::{Bytes, BytesMut};
9use curve25519_dalek::{
10    ristretto::{CompressedRistretto, RistrettoPoint},
11    scalar::Scalar,
12};
13use quinn::{Endpoint, RecvStream, SendStream};
14use std::{convert::TryInto, net::SocketAddr};
15use tracing::log;
16
17use crate::error::{MpcNetworkError, SetupError};
18
19pub type PartyId = u64;
20
21const BYTES_PER_POINT: usize = 32;
22const BYTES_PER_SCALAR: usize = 32;
23const BYTES_PER_U64: usize = 8;
24
25/**
26 * Helpers
27 */
28
29/// Convert a vector of scalars to a byte buffer
30fn scalars_to_bytes(scalars: &[Scalar]) -> Bytes {
31    let mut payload = BytesMut::new();
32    scalars.iter().for_each(|scalar| {
33        let bytes = scalar.to_bytes();
34        payload.extend_from_slice(&bytes);
35    });
36
37    payload.freeze()
38}
39
40/// Convert a byte buffer back to a vector of scalars
41fn bytes_to_scalars(bytes: &[u8]) -> Result<Vec<Scalar>, MpcNetworkError> {
42    bytes
43        .chunks(BYTES_PER_SCALAR)
44        .map(|bytes_chunk| {
45            Scalar::from_canonical_bytes(
46                bytes_chunk
47                    .try_into()
48                    .expect("unexpected number of bytes per chunk"),
49            )
50            .ok_or(MpcNetworkError::SerializationError)
51        })
52        .collect::<Result<Vec<Scalar>, MpcNetworkError>>()
53}
54
55/// Convert a vector of Ristretto points to bytes
56fn points_to_bytes(points: &[RistrettoPoint]) -> Bytes {
57    // Map to bytes
58    let mut payload = BytesMut::new();
59    points.iter().for_each(|point| {
60        let bytes = point.compress().to_bytes();
61        payload.extend_from_slice(&bytes);
62    });
63
64    payload.freeze()
65}
66
67/// Convert a byte buffer back to a vector of points
68fn bytes_to_points(bytes: &[u8]) -> Result<Vec<RistrettoPoint>, MpcNetworkError> {
69    bytes
70        .chunks(BYTES_PER_POINT)
71        .map(|bytes_chunk| {
72            CompressedRistretto(
73                bytes_chunk
74                    .try_into()
75                    .expect("unexpected number of bytes per chunk"),
76            )
77            .decompress()
78            .ok_or(MpcNetworkError::SerializationError)
79        })
80        .collect::<Result<Vec<RistrettoPoint>, MpcNetworkError>>()
81}
82
83/// MpcNetwork represents the network functionality needed for 2PC execution
84/// Note that only two party computation is implemented here
85#[async_trait]
86pub trait MpcNetwork {
87    /// Returns the ID of the given party in the MPC computation
88    fn party_id(&self) -> u64;
89    /// Returns whether the local party is the king of the MPC (party 0)
90    fn am_king(&self) -> bool {
91        self.party_id() == 0
92    }
93    /// The local party sends a byte buffer to the peer with an additional u64
94    /// prepended specifying the length of the payload
95    async fn send_bytes(&mut self, bytes: &[u8]) -> Result<(), MpcNetworkError>;
96    /// The local party awaits bytes from a peer; receiver expects that the payload
97    /// is prepended with a length indicating the message size
98    async fn receive_bytes(&mut self) -> Result<Vec<u8>, MpcNetworkError>;
99    /// The local party sends a vector of scalars to the peer
100    async fn send_scalars(&mut self, scalars: &[Scalar]) -> Result<(), MpcNetworkError>;
101    /// The local party sends a single scalar to the peer
102    async fn send_single_scalar(&mut self, scalar: Scalar) -> Result<(), MpcNetworkError> {
103        self.send_scalars(&[scalar]).await
104    }
105    /// The local party receives exactly `n` scalars from the peer
106    async fn receive_scalars(
107        &mut self,
108        num_expected: usize,
109    ) -> Result<Vec<Scalar>, MpcNetworkError>;
110    /// The local party receives a single scalar from the peer
111    async fn receive_single_scalar(&mut self) -> Result<Scalar, MpcNetworkError> {
112        Ok(self.receive_scalars(1).await?[0])
113    }
114    /// Both parties broadcast a vector of scalars to one another
115    async fn broadcast_scalars(
116        &mut self,
117        scalars: &[Scalar],
118    ) -> Result<Vec<Scalar>, MpcNetworkError>;
119    /// Both parties broadcast a single scalar to one another
120    async fn broadcast_single_scalar(&mut self, scalar: Scalar) -> Result<Scalar, MpcNetworkError> {
121        Ok(self.broadcast_scalars(&[scalar]).await?[0])
122    }
123    /// The local party sends a vector of Ristretto points to the peer
124    async fn send_points(&mut self, points: &[RistrettoPoint]) -> Result<(), MpcNetworkError>;
125    /// The local party sends a single Ristretto point to the peer
126    async fn send_single_point(&mut self, point: RistrettoPoint) -> Result<(), MpcNetworkError> {
127        Ok(self.send_points(&[point]).await?)
128    }
129    /// The local party awaits a vector of Ristretto points from the peer
130    async fn receive_points(
131        &mut self,
132        num_expected: usize,
133    ) -> Result<Vec<RistrettoPoint>, MpcNetworkError>;
134    /// The local party awaits a single Ristretto point from the peer
135    async fn receive_single_point(&mut self) -> Result<RistrettoPoint, MpcNetworkError> {
136        Ok(self.receive_points(1).await?[0])
137    }
138    /// Both parties broadcast a vector of points to one another
139    async fn broadcast_points(
140        &mut self,
141        points: &[RistrettoPoint],
142    ) -> Result<Vec<RistrettoPoint>, MpcNetworkError>;
143    /// Both parties broadcast a single point to one another
144    async fn broadcast_single_point(
145        &mut self,
146        point: RistrettoPoint,
147    ) -> Result<RistrettoPoint, MpcNetworkError> {
148        Ok(self.broadcast_points(&[point]).await?[0])
149    }
150    /// Closes the connections opened in the handshake phase
151    async fn close(&mut self) -> Result<(), MpcNetworkError>;
152}
153
154/// The order in which the local party should read when exchanging values
155#[derive(Clone, Debug)]
156pub enum ReadWriteOrder {
157    ReadFirst,
158    WriteFirst,
159}
160
161/// Implements an MpcNetwork on top of QUIC
162#[derive(Debug)]
163pub struct QuicTwoPartyNet {
164    /// The index of the local party in the participants
165    party_id: PartyId,
166    /// Whether the network has been bootstrapped yet
167    connected: bool,
168    /// The address of the local peer
169    local_addr: SocketAddr,
170    /// Addresses of the counterparties in the MPC
171    peer_addr: SocketAddr,
172    /// The send side of the bidirectional stream
173    send_stream: Option<SendStream>,
174    /// The receive side of the bidirectional stream
175    recv_stream: Option<RecvStream>,
176}
177
178#[allow(clippy::redundant_closure)] // For readability of error handling
179impl<'a> QuicTwoPartyNet {
180    pub fn new(party_id: PartyId, local_addr: SocketAddr, peer_addr: SocketAddr) -> Self {
181        // Construct the QUIC net
182        Self {
183            party_id,
184            local_addr,
185            peer_addr,
186            connected: false,
187            send_stream: None,
188            recv_stream: None,
189        }
190    }
191
192    /// Returns the read order for the local peer; king is write first
193    fn read_order(&self) -> ReadWriteOrder {
194        if self.am_king() {
195            ReadWriteOrder::WriteFirst
196        } else {
197            ReadWriteOrder::ReadFirst
198        }
199    }
200
201    /// Returns an error if the network is not connected
202    fn assert_connected(&self) -> Result<(), MpcNetworkError> {
203        if self.connected {
204            Ok(())
205        } else {
206            Err(MpcNetworkError::NetworkUninitialized)
207        }
208    }
209    /// Establishes connections to the peer
210    pub async fn connect(&mut self) -> Result<(), MpcNetworkError> {
211        // Build the client and server configs
212        let (client_config, server_config) =
213            config::build_configs().map_err(|err| MpcNetworkError::ConnectionSetupError(err))?;
214
215        // Create a quinn server
216        let mut local_endpoint = Endpoint::server(server_config, self.local_addr).map_err(|e| {
217            log::error!("error setting up quinn server: {e:?}");
218            MpcNetworkError::ConnectionSetupError(SetupError::ServerSetupError)
219        })?;
220        local_endpoint.set_default_client_config(client_config);
221
222        // The king dials the peer who awaits connection
223        let connection = {
224            if self.am_king() {
225                local_endpoint
226                    .connect(self.peer_addr, config::SERVER_NAME)
227                    .map_err(|err| {
228                        log::error!("error setting up quic endpoint connection: {err}");
229                        MpcNetworkError::ConnectionSetupError(SetupError::ConnectError(err))
230                    })?
231                    .await
232                    .map_err(|err| {
233                        log::error!("error connecting to the remote quic endpoint: {err}");
234                        MpcNetworkError::ConnectionSetupError(SetupError::ConnectionError(err))
235                    })?
236            } else {
237                local_endpoint
238                    .accept()
239                    .await
240                    .ok_or_else(|| {
241                        log::error!("no incoming connection while awaiting quic endpoint");
242                        MpcNetworkError::ConnectionSetupError(SetupError::NoIncomingConnection)
243                    })?
244                    .await
245                    .map_err(|err| {
246                        log::error!("error while establishing remote connection as listener");
247                        MpcNetworkError::ConnectionSetupError(SetupError::ConnectionError(err))
248                    })?
249            }
250        };
251
252        // King opens a bidirectional stream on top of the connection
253        let (send, recv) = {
254            if self.am_king() {
255                connection.open_bi().await.map_err(|err| {
256                    log::error!("error opening bidirectional stream: {err}");
257                    MpcNetworkError::ConnectionSetupError(SetupError::ConnectionError(err))
258                })?
259            } else {
260                connection.accept_bi().await.map_err(|err| {
261                    log::error!("error accepting bidirectional stream: {err}");
262                    MpcNetworkError::ConnectionSetupError(SetupError::ConnectionError(err))
263                })?
264            }
265        };
266
267        // Update MpcNet state
268        self.connected = true;
269        self.send_stream = Some(send);
270        self.recv_stream = Some(recv);
271
272        Ok(())
273    }
274
275    /// Write a stream of bytes to the stream
276    async fn write_bytes(&mut self, payload: &[u8]) -> Result<(), MpcNetworkError> {
277        self.send_stream
278            .as_mut()
279            .unwrap()
280            .write_all(payload)
281            .await
282            .map_err(|_| MpcNetworkError::SendError)
283    }
284
285    /// Read exactly `n` bytes from the stream
286    async fn read_bytes(&mut self, num_bytes: usize) -> Result<Vec<u8>, MpcNetworkError> {
287        let mut read_buffer = vec![0u8; num_bytes];
288        self.recv_stream
289            .as_mut()
290            .unwrap()
291            .read_exact(&mut read_buffer)
292            .await
293            .map_err(|_| MpcNetworkError::RecvError)?;
294
295        Ok(read_buffer.to_vec())
296    }
297
298    /// Write a stream of bytes to the network, then expect the same back from the connected peer
299    async fn write_then_read_bytes(
300        &mut self,
301        order: ReadWriteOrder,
302        payload: &[u8],
303    ) -> Result<Vec<u8>, MpcNetworkError> {
304        let payload_length = payload.len();
305
306        Ok(match order {
307            ReadWriteOrder::ReadFirst => {
308                let bytes_read = self.read_bytes(payload_length).await?;
309                self.write_bytes(payload).await?;
310                bytes_read
311            }
312            ReadWriteOrder::WriteFirst => {
313                self.write_bytes(payload).await?;
314                self.read_bytes(payload_length).await?
315            }
316        })
317    }
318}
319
320#[async_trait]
321impl MpcNetwork for QuicTwoPartyNet {
322    fn party_id(&self) -> u64 {
323        self.party_id
324    }
325
326    async fn send_bytes(&mut self, bytes: &[u8]) -> Result<(), MpcNetworkError> {
327        self.assert_connected()?;
328
329        // Prepend the length of the payload as a little endian encoded u64
330        let length = (bytes.len() as u64).to_le_bytes();
331        self.write_bytes(&length).await?;
332        self.write_bytes(bytes).await
333    }
334
335    async fn receive_bytes(&mut self) -> Result<Vec<u8>, MpcNetworkError> {
336        self.assert_connected()?;
337
338        // Read a u64 indicating the payload length
339        let length = u64::from_le_bytes(self.read_bytes(BYTES_PER_U64).await?.try_into().unwrap());
340        self.read_bytes(length as usize).await
341    }
342
343    async fn send_scalars(&mut self, scalars: &[Scalar]) -> Result<(), MpcNetworkError> {
344        self.assert_connected()?;
345
346        // To byte buffer
347        let payload = scalars_to_bytes(scalars);
348        self.write_bytes(&payload).await?;
349
350        Ok(())
351    }
352
353    async fn receive_scalars(
354        &mut self,
355        num_scalars: usize,
356    ) -> Result<Vec<Scalar>, MpcNetworkError> {
357        self.assert_connected()?;
358        let bytes_read = self.read_bytes(num_scalars * BYTES_PER_SCALAR).await?;
359
360        bytes_to_scalars(&bytes_read)
361    }
362
363    async fn broadcast_scalars(
364        &mut self,
365        scalars: &[Scalar],
366    ) -> Result<Vec<Scalar>, MpcNetworkError> {
367        self.assert_connected()?;
368
369        // To byte buffer
370        let payload = scalars_to_bytes(scalars);
371
372        let read_buffer = self
373            .write_then_read_bytes(self.read_order(), &payload)
374            .await?;
375
376        bytes_to_scalars(&read_buffer)
377    }
378
379    async fn send_points(&mut self, points: &[RistrettoPoint]) -> Result<(), MpcNetworkError> {
380        let payload = points_to_bytes(points);
381        self.write_bytes(&payload).await
382    }
383
384    async fn receive_points(
385        &mut self,
386        num_points: usize,
387    ) -> Result<Vec<RistrettoPoint>, MpcNetworkError> {
388        let read_buffer = self.read_bytes(BYTES_PER_POINT * num_points).await?;
389        bytes_to_points(&read_buffer)
390    }
391
392    async fn broadcast_points(
393        &mut self,
394        points: &[RistrettoPoint],
395    ) -> Result<Vec<RistrettoPoint>, MpcNetworkError> {
396        self.assert_connected()?;
397
398        // To byte buffer
399        let payload = points_to_bytes(points);
400        let read_buffer = self
401            .write_then_read_bytes(self.read_order(), &payload)
402            .await?;
403
404        // Deserialize back to Ristretto points
405        bytes_to_points(&read_buffer)
406    }
407
408    async fn close(&mut self) -> Result<(), MpcNetworkError> {
409        self.assert_connected()?;
410
411        self.send_stream
412            .as_mut()
413            .unwrap()
414            .finish()
415            .await
416            .map_err(|_| MpcNetworkError::ConnectionTeardownError)
417    }
418}
419
420#[cfg(test)]
421mod test {
422    use std::net::SocketAddr;
423
424    use curve25519_dalek::ristretto::RistrettoPoint;
425    use rand_core::OsRng;
426    use tokio;
427
428    use super::{MpcNetwork, QuicTwoPartyNet};
429
430    #[tokio::test]
431    async fn test_errors() {
432        let socket_addr: SocketAddr = "127.0.0.1:8000".parse().unwrap();
433        let mut net = QuicTwoPartyNet::new(0, socket_addr, socket_addr);
434
435        assert!(net.broadcast_points(&[]).await.is_err());
436
437        let mut rng = OsRng {};
438        assert!(net
439            .broadcast_single_point(RistrettoPoint::random(&mut rng))
440            .await
441            .is_err())
442    }
443}