libp2p_iroh/
connection.rs

1use std::{error::Error, fmt::Display, pin::Pin, task::Poll};
2
3use crate::{
4    TransportError,
5    stream::{Stream, StreamError},
6};
7use futures::{FutureExt, future::BoxFuture};
8use iroh::endpoint::{RecvStream, SendStream};
9use libp2p::core::StreamMuxer;
10use tokio::io::{AsyncReadExt, AsyncWriteExt};
11
12#[derive(Debug)]
13pub struct ConnectionError {
14    kind: ConnectionErrorKind,
15}
16
17#[derive(Debug)]
18pub enum ConnectionErrorKind {
19    Accept(String),
20    Open(String),
21    Stream(String),
22}
23
24impl Display for ConnectionError {
25    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
26        write!(f, "ConnectionError: {:?}", self.kind)
27    }
28}
29
30impl Error for ConnectionError {}
31
32impl From<iroh::endpoint::ConnectionError> for ConnectionError {
33    fn from(err: iroh::endpoint::ConnectionError) -> Self {
34        Self {
35            kind: ConnectionErrorKind::Accept(err.to_string()),
36        }
37    }
38}
39
40impl From<&str> for ConnectionError {
41    fn from(err: &str) -> Self {
42        Self {
43            kind: ConnectionErrorKind::Accept(err.to_string()),
44        }
45    }
46}
47
48impl From<StreamError> for ConnectionError {
49    fn from(err: StreamError) -> Self {
50        Self {
51            kind: ConnectionErrorKind::Stream(err.to_string()),
52        }
53    }
54}
55
56pub struct Connection {
57    connection: iroh::endpoint::Connection,
58    incoming: Option<BoxFuture<'static, Result<(SendStream, RecvStream), ConnectionError>>>,
59    outgoing: Option<BoxFuture<'static, Result<(SendStream, RecvStream), ConnectionError>>>,
60    closing: Option<BoxFuture<'static, ConnectionError>>,
61}
62
63pub struct Connecting {
64    pub connecting:
65        BoxFuture<'static, Result<(libp2p::PeerId, iroh::endpoint::Connection), TransportError>>,
66}
67
68impl Connection {
69    pub fn new(connection: iroh::endpoint::Connection) -> Self {
70        tracing::debug!("Connection::new - Creating new connection wrapper");
71        Self {
72            connection,
73            incoming: None,
74            outgoing: None,
75            closing: None,
76        }
77    }
78}
79
80impl StreamMuxer for Connection {
81    type Substream = Stream;
82    type Error = ConnectionError;
83
84    fn poll_inbound(
85        self: Pin<&mut Self>,
86        cx: &mut std::task::Context<'_>,
87    ) -> Poll<Result<Self::Substream, Self::Error>> {
88        let this = self.get_mut();
89
90        let incoming = this.incoming.get_or_insert_with(|| {
91            tracing::debug!("Connection::poll_inbound - Setting up incoming stream future");
92            let connection = this.connection.clone();
93            async move {
94                tracing::debug!("Connection::poll_inbound - Accepting bidirectional stream");
95                match connection.accept_bi().await {
96                    Ok((s, mut r)) => {
97                        tracing::debug!("Connection::poll_inbound - Bidirectional stream accepted, reading handshake byte");
98                        r.read_u8().await.map_err(|e| {
99                            tracing::error!("Connection::poll_inbound - Failed to read handshake byte: {}", e);
100                            ConnectionError::from("Failed to read from stream")
101                        })?;
102                        tracing::debug!("Connection::poll_inbound - Handshake byte read successfully");
103                        Ok((s, r))
104                    },
105                    Err(e) => {
106                        tracing::error!("Connection::poll_inbound - Failed to accept bidirectional stream: {}", e);
107                        Err(ConnectionError::from("Iroh handshake failed during accept"))
108                    }
109                }
110             }.boxed()
111        });
112
113        let (send, recv) = futures::ready!(incoming.poll_unpin(cx))?;
114        this.incoming.take();
115        tracing::debug!("Connection::poll_inbound - Inbound stream ready, creating Stream wrapper");
116        Poll::Ready(Stream::new(send, recv).map_err(Into::into))
117    }
118
119    fn poll_outbound(
120        self: Pin<&mut Self>,
121        cx: &mut std::task::Context<'_>,
122    ) -> Poll<Result<Self::Substream, Self::Error>> {
123        let this = self.get_mut();
124
125        let outgoing = this.outgoing.get_or_insert_with(|| {
126            tracing::debug!("Connection::poll_outbound - Setting up outgoing stream future");
127            let connection = this.connection.clone();
128            async move {
129                tracing::debug!("Connection::poll_outbound - Opening bidirectional stream");
130                match connection.open_bi().await {
131                    Ok((mut s, r)) => {
132                        tracing::debug!("Connection::poll_outbound - Bidirectional stream opened, writing handshake byte");
133                        // one byte iroh-handshake since accept only connects after open and write, not just open
134                        s.write_u8(0).await.map_err(|e| {
135                            tracing::error!("Connection::poll_outbound - Failed to write handshake byte: {}", e);
136                            ConnectionError::from("Failed to write to stream")
137                        })?;
138                        tracing::debug!("Connection::poll_outbound - Handshake byte written successfully");
139                        Ok((s, r))
140                    }
141                    Err(e) => {
142                        tracing::error!("Connection::poll_outbound - Failed to open bidirectional stream: {}", e);
143                        Err(ConnectionError::from("Iroh handshake failed during open"))
144                    }
145                }
146            }.boxed()
147        });
148
149        let (send, recv) = futures::ready!(outgoing.poll_unpin(cx))?;
150        this.outgoing.take();
151        tracing::debug!(
152            "Connection::poll_outbound - Outbound stream ready, creating Stream wrapper"
153        );
154        Poll::Ready(Stream::new(send, recv).map_err(Into::into))
155    }
156
157    fn poll_close(
158        self: Pin<&mut Self>,
159        cx: &mut std::task::Context<'_>,
160    ) -> Poll<Result<(), Self::Error>> {
161        let this = self.get_mut();
162
163        let closing = this.closing.get_or_insert_with(|| {
164            tracing::debug!("Connection::poll_close - Closing connection");
165            this.connection.close(From::from(0u32), &[]);
166            let connection = this.connection.clone();
167            async move {
168                tracing::debug!("Connection::poll_close - Waiting for connection to close");
169                connection.closed().await.into()
170            }
171            .boxed()
172        });
173
174        if matches!(
175            futures::ready!(closing.poll_unpin(cx)),
176            crate::ConnectionError { .. }
177        ) {
178            tracing::error!("Connection::poll_close - Failed to close connection");
179            return Poll::Ready(Err("failed to close connection".into()));
180        };
181
182        tracing::debug!("Connection::poll_close - Connection closed successfully");
183        Poll::Ready(Ok(()))
184    }
185
186    fn poll(
187        self: Pin<&mut Self>,
188        _cx: &mut std::task::Context<'_>,
189    ) -> Poll<Result<libp2p::core::muxing::StreamMuxerEvent, Self::Error>> {
190        Poll::Pending
191    }
192}
193
194impl Future for Connecting {
195    type Output = Result<(libp2p::PeerId, libp2p::core::muxing::StreamMuxerBox), TransportError>;
196
197    fn poll(mut self: Pin<&mut Self>, cx: &mut std::task::Context<'_>) -> Poll<Self::Output> {
198        tracing::debug!("Connecting::poll - Polling connection future");
199        let (peer_id, conn) = match self.connecting.poll_unpin(cx) {
200            Poll::Ready(Ok((peer_id, conn))) => {
201                tracing::debug!("Connecting::poll - Connection established");
202                (peer_id, conn)
203            }
204            Poll::Ready(Err(e)) => {
205                tracing::error!("Connecting::poll - Connection failed: {}", e);
206                return Poll::Ready(Err(e));
207            }
208            Poll::Pending => {
209                tracing::trace!("Connecting::poll - Connection still pending");
210                return Poll::Pending;
211            }
212        };
213
214        let muxer = Connection::new(conn);
215
216        tracing::debug!("Connecting::poll - Connection muxer created");
217        Poll::Ready(Ok((
218            peer_id,
219            libp2p::core::muxing::StreamMuxerBox::new(muxer),
220        )))
221    }
222}
libp2p_iroh/connection.rs

libp2p_iroh/
connection.rs
