fedimint_server/net/
p2p_connection.rs

1#[cfg(test)]
2mod tests;
3
4use std::io::Cursor;
5use std::pin::Pin;
6use std::time::Duration;
7
8use anyhow::Context;
9use async_trait::async_trait;
10use bytes::{Bytes, BytesMut};
11use fedimint_core::encoding::{Decodable, Encodable};
12use fedimint_core::module::registry::ModuleDecoderRegistry;
13use fedimint_server_core::dashboard_ui::ConnectionType;
14use futures::{SinkExt, Stream, StreamExt};
15use iroh_next::endpoint::{Connection as IrohV1Connection, RecvStream as IrohV1RecvStream};
16use serde::Serialize;
17use serde::de::DeserializeOwned;
18use tokio::net::TcpStream;
19use tokio_rustls::TlsStream;
20use tokio_util::codec::{Framed, LengthDelimitedCodec};
21
22/// Maximum size of a p2p message in bytes. The largest message we expect to
23/// receive is a signed session outcome.
24const MAX_P2P_MESSAGE_SIZE: usize = 10_000_000;
25
26pub type DynP2PConnection<M> = Box<dyn IP2PConnection<M>>;
27
28pub type DynIP2PFrame<M> = Box<dyn IP2PFrame<M>>;
29
30/// Type-erased stream notifying the P2P state machine that connection metadata
31/// may have changed.
32pub type DynConnectionStatusUpdates = Pin<Box<dyn Stream<Item = ()> + Send + 'static>>;
33
34#[async_trait]
35pub trait IP2PFrame<M>: Send + 'static {
36    /// Read the entire frame from the connection and deserialize it into a
37    /// message. This is *not* required to be cancel-safe.
38    async fn read_to_end(&mut self) -> anyhow::Result<M>;
39
40    fn into_dyn(self) -> DynIP2PFrame<M>
41    where
42        Self: Sized,
43    {
44        Box::new(self)
45    }
46}
47
48#[async_trait]
49pub trait IP2PConnection<M>: Send + 'static {
50    /// Send a message over the connection. This is *not* required to be
51    /// cancel-safe.
52    async fn send(&mut self, message: M) -> anyhow::Result<()>;
53
54    /// Receive a p2p frame from the connection. This is *required* to be
55    /// cancel-safe.
56    async fn receive(&mut self) -> anyhow::Result<DynIP2PFrame<M>>;
57
58    /// Get the round-trip time of the connection.
59    fn rtt(&self) -> Option<Duration>;
60
61    /// Get the transport type currently backing this live connection.
62    fn connection_type(&self) -> Option<ConnectionType> {
63        None
64    }
65
66    /// Subscribe to notifications that the live connection metadata may have
67    /// changed.
68    ///
69    /// Implementations should treat these as wake-ups only. The state machine
70    /// reads a fresh [`Self::connection_type`] and [`Self::rtt`] snapshot after
71    /// each notification. `None` means notifications are unsupported, while the
72    /// end of a returned stream means there will be no more notifications.
73    fn connection_status_updates(&self) -> Option<DynConnectionStatusUpdates> {
74        None
75    }
76
77    fn into_dyn(self) -> DynP2PConnection<M>
78    where
79        Self: Sized,
80    {
81        Box::new(self)
82    }
83}
84
85/// Implementations of the IP2PFrame and IP2PConnection traits for TLS
86
87#[async_trait]
88impl<M> IP2PFrame<M> for BytesMut
89where
90    M: Decodable + DeserializeOwned + Send + 'static,
91{
92    async fn read_to_end(&mut self) -> anyhow::Result<M> {
93        if let Ok(message) = M::consensus_decode_whole(self, &ModuleDecoderRegistry::default()) {
94            return Ok(message);
95        }
96
97        Ok(bincode::deserialize_from(Cursor::new(&**self))?)
98    }
99}
100
101#[async_trait]
102impl<M> IP2PConnection<M> for Framed<TlsStream<TcpStream>, LengthDelimitedCodec>
103where
104    M: Encodable + Decodable + Serialize + DeserializeOwned + Send + 'static,
105{
106    async fn send(&mut self, message: M) -> anyhow::Result<()> {
107        let mut bytes = Vec::new();
108
109        bincode::serialize_into(&mut bytes, &message)?;
110
111        SinkExt::send(self, Bytes::from_owner(bytes)).await?;
112
113        Ok(())
114    }
115
116    async fn receive(&mut self) -> anyhow::Result<DynIP2PFrame<M>> {
117        let message = self
118            .next()
119            .await
120            .context("Framed stream is closed")??
121            .into_dyn();
122
123        Ok(message)
124    }
125
126    fn rtt(&self) -> Option<Duration> {
127        None
128    }
129
130    fn connection_type(&self) -> Option<ConnectionType> {
131        Some(ConnectionType::Direct)
132    }
133}
134
135/// Compatibility implementations for the public Iroh 0.35 connection types.
136
137#[async_trait]
138impl<M> IP2PFrame<M> for iroh::endpoint::RecvStream
139where
140    M: Decodable + DeserializeOwned + Send + 'static,
141{
142    async fn read_to_end(&mut self) -> anyhow::Result<M> {
143        let bytes = self.read_to_end(MAX_P2P_MESSAGE_SIZE).await?;
144
145        if let Ok(message) = M::consensus_decode_whole(&bytes, &ModuleDecoderRegistry::default()) {
146            return Ok(message);
147        }
148
149        Ok(bincode::deserialize_from(Cursor::new(&bytes))?)
150    }
151}
152
153#[async_trait]
154impl<M> IP2PConnection<M> for iroh::endpoint::Connection
155where
156    M: Encodable + Decodable + Serialize + DeserializeOwned + Send + 'static,
157{
158    async fn send(&mut self, message: M) -> anyhow::Result<()> {
159        let mut bytes = Vec::new();
160
161        bincode::serialize_into(&mut bytes, &message)?;
162
163        let mut sink = self.open_uni().await?;
164
165        sink.write_all(&bytes).await?;
166        sink.finish()?;
167
168        Ok(())
169    }
170
171    async fn receive(&mut self) -> anyhow::Result<DynIP2PFrame<M>> {
172        Ok(self.accept_uni().await?.into_dyn())
173    }
174
175    fn rtt(&self) -> Option<Duration> {
176        Some(iroh::endpoint::Connection::rtt(self))
177    }
178}
179
180/// Implementations of the P2P traits for Iroh 1.0.
181
182#[async_trait]
183impl<M> IP2PFrame<M> for IrohV1RecvStream
184where
185    M: Decodable + DeserializeOwned + Send + 'static,
186{
187    async fn read_to_end(&mut self) -> anyhow::Result<M> {
188        let bytes = self.read_to_end(MAX_P2P_MESSAGE_SIZE).await?;
189
190        if let Ok(message) = M::consensus_decode_whole(&bytes, &ModuleDecoderRegistry::default()) {
191            return Ok(message);
192        }
193
194        Ok(bincode::deserialize_from(Cursor::new(&bytes))?)
195    }
196}
197
198#[async_trait]
199impl<M> IP2PConnection<M> for IrohV1Connection
200where
201    M: Encodable + Decodable + Serialize + DeserializeOwned + Send + 'static,
202{
203    async fn send(&mut self, message: M) -> anyhow::Result<()> {
204        let mut bytes = Vec::new();
205
206        bincode::serialize_into(&mut bytes, &message)?;
207
208        let mut sink = self.open_uni().await?;
209
210        sink.write_all(&bytes).await?;
211
212        sink.finish()?;
213
214        Ok(())
215    }
216
217    async fn receive(&mut self) -> anyhow::Result<DynIP2PFrame<M>> {
218        Ok(self.accept_uni().await?.into_dyn())
219    }
220
221    fn rtt(&self) -> Option<Duration> {
222        self.paths()
223            .iter()
224            .find(iroh_next::endpoint::Path::is_selected)
225            .and_then(|path| IrohV1Connection::rtt(self, path.id()))
226    }
227
228    fn connection_type(&self) -> Option<ConnectionType> {
229        connection_type_from_paths(self.paths().iter().map(|path| IrohPath {
230            selected: path.is_selected(),
231            kind: if path.is_ip() {
232                IrohPathKind::Direct
233            } else if path.is_relay() {
234                IrohPathKind::Relay
235            } else {
236                IrohPathKind::Unknown
237            },
238        }))
239    }
240
241    fn connection_status_updates(&self) -> Option<DynConnectionStatusUpdates> {
242        Some(Box::pin(self.path_events().map(|_| ())))
243    }
244}
245
246#[derive(Clone, Copy)]
247enum IrohPathKind {
248    Direct,
249    Relay,
250    Unknown,
251}
252
253#[derive(Clone, Copy)]
254struct IrohPath {
255    selected: bool,
256    kind: IrohPathKind,
257}
258
259fn connection_type_from_paths(paths: impl IntoIterator<Item = IrohPath>) -> Option<ConnectionType> {
260    let mut direct = false;
261    let mut relay = false;
262    for path in paths {
263        if path.selected {
264            direct |= matches!(path.kind, IrohPathKind::Direct);
265            relay |= matches!(path.kind, IrohPathKind::Relay);
266        }
267    }
268
269    match (direct, relay) {
270        (true, true) => Some(ConnectionType::Mixed),
271        (true, false) => Some(ConnectionType::Direct),
272        (false, true) => Some(ConnectionType::Relay),
273        (false, false) => None,
274    }
275}
fedimint_server/net/p2p_connection.rs

fedimint_server/net/
p2p_connection.rs
