Skip to main content

libp2p_iroh/
stream.rs

1use std::{fmt::Display, pin::Pin};
2
3use tokio::io::AsyncWrite;
4
5// IrohStream error:
6#[derive(Debug, Clone)]
7pub struct StreamError {
8    kind: StreamErrorKind,
9}
10
11#[derive(Debug, Clone)]
12pub enum StreamErrorKind {
13    Read(String),
14    Write(String),
15    Connection(String),
16}
17
18impl From<std::io::Error> for StreamError {
19    fn from(err: std::io::Error) -> Self {
20        Self {
21            kind: StreamErrorKind::Read(err.to_string()),
22        }
23    }
24}
25
26impl From<iroh::endpoint::ConnectionError> for StreamError {
27    fn from(err: iroh::endpoint::ConnectionError) -> Self {
28        Self {
29            kind: StreamErrorKind::Connection(err.to_string()),
30        }
31    }
32}
33
34impl From<iroh::endpoint::WriteError> for StreamError {
35    fn from(err: iroh::endpoint::WriteError) -> Self {
36        Self {
37            kind: StreamErrorKind::Write(err.to_string()),
38        }
39    }
40}
41
42impl From<iroh::endpoint::ReadError> for StreamError {
43    fn from(err: iroh::endpoint::ReadError) -> Self {
44        Self {
45            kind: StreamErrorKind::Read(err.to_string()),
46        }
47    }
48}
49
50impl From<&str> for StreamError {
51    fn from(err: &str) -> Self {
52        Self {
53            kind: StreamErrorKind::Connection(err.to_string()),
54        }
55    }
56}
57
58impl Display for StreamError {
59    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
60        match &self.kind {
61            StreamErrorKind::Read(msg) => write!(f, "IrohStream Read Error: {msg}"),
62            StreamErrorKind::Write(msg) => write!(f, "IrohStream Write Error: {msg}"),
63            StreamErrorKind::Connection(msg) => {
64                write!(f, "IrohStream Connection Error: {msg}")
65            }
66        }
67    }
68}
69
70impl std::error::Error for StreamError {}
71
72#[derive(Debug)]
73pub struct Stream {
74    sender: Option<iroh::endpoint::SendStream>,
75    receiver: Option<iroh::endpoint::RecvStream>,
76    closing: bool,
77}
78
79impl Stream {
80    pub fn new(
81        sender: iroh::endpoint::SendStream,
82        receiver: iroh::endpoint::RecvStream,
83    ) -> Result<Self, StreamError> {
84        tracing::debug!("Stream::new - Creating new stream wrapper");
85        Ok(Self {
86            sender: Some(sender),
87            receiver: Some(receiver),
88            closing: false,
89        })
90    }
91}
92
93impl futures::AsyncRead for Stream {
94    fn poll_read(
95        mut self: std::pin::Pin<&mut Self>,
96        cx: &mut std::task::Context<'_>,
97        buf: &mut [u8],
98    ) -> std::task::Poll<std::io::Result<usize>> {
99        if let Some(receiver) = &mut self.receiver {
100            match Pin::new(receiver).poll_read(cx, buf) {
101                std::task::Poll::Ready(Ok(n)) => {
102                    if n == 0 {
103                        tracing::debug!("Stream::poll_read - EOF reached (0 bytes)");
104                    } else {
105                        tracing::trace!("Stream::poll_read - Read {} bytes", n);
106                    }
107                    std::task::Poll::Ready(Ok(n))
108                }
109                std::task::Poll::Ready(Err(e)) => {
110                    tracing::debug!("Stream::poll_read - Read error: {}", e);
111                    std::task::Poll::Ready(Err(std::io::Error::other(e)))
112                }
113                std::task::Poll::Pending => std::task::Poll::Pending,
114            }
115        } else {
116            tracing::debug!("Stream::poll_read - Stream receiver already closed locally");
117            std::task::Poll::Ready(Err(std::io::Error::new(
118                std::io::ErrorKind::BrokenPipe,
119                "stream receiver closed",
120            )))
121        }
122    }
123}
124
125impl futures::AsyncWrite for Stream {
126    fn poll_write(
127        mut self: Pin<&mut Self>,
128        cx: &mut std::task::Context<'_>,
129        buf: &[u8],
130    ) -> std::task::Poll<std::io::Result<usize>> {
131        if let Some(sender) = &mut self.sender {
132            match Pin::new(sender).poll_write(cx, buf) {
133                std::task::Poll::Ready(Ok(n)) => {
134                    tracing::trace!("Stream::poll_write - Wrote {} bytes", n);
135                    std::task::Poll::Ready(Ok(n))
136                }
137                std::task::Poll::Ready(Err(e)) => {
138                    // Check if this is a "stopped" error (remote side closed)
139                    let err_str = e.to_string();
140                    if err_str.contains("stopped") || err_str.contains("error 0") {
141                        tracing::debug!("Stream::poll_write - Remote peer closed stream: {}", e);
142                    } else {
143                        tracing::error!("Stream::poll_write - Write error: {}", e);
144                    }
145                    std::task::Poll::Ready(Err(std::io::Error::other(e)))
146                }
147                std::task::Poll::Pending => std::task::Poll::Pending,
148            }
149        } else {
150            tracing::debug!("Stream::poll_write - Stream sender already closed locally");
151            std::task::Poll::Ready(Err(std::io::Error::new(
152                std::io::ErrorKind::BrokenPipe,
153                "stream sender closed",
154            )))
155        }
156    }
157
158    fn poll_flush(
159        mut self: Pin<&mut Self>,
160        cx: &mut std::task::Context<'_>,
161    ) -> std::task::Poll<std::io::Result<()>> {
162        if let Some(sender) = &mut self.sender {
163            match Pin::new(sender).poll_flush(cx) {
164                std::task::Poll::Ready(Ok(())) => {
165                    tracing::trace!("Stream::poll_flush - Flush successful");
166                    std::task::Poll::Ready(Ok(()))
167                }
168                std::task::Poll::Ready(Err(e)) => {
169                    tracing::debug!("Stream::poll_flush - Flush error: {}", e);
170                    std::task::Poll::Ready(Err(std::io::Error::other(e)))
171                }
172                std::task::Poll::Pending => std::task::Poll::Pending,
173            }
174        } else {
175            tracing::debug!("Stream::poll_flush - Stream sender already closed locally");
176            std::task::Poll::Ready(Err(std::io::Error::new(
177                std::io::ErrorKind::BrokenPipe,
178                "stream sender closed",
179            )))
180        }
181    }
182
183    fn poll_close(
184        mut self: Pin<&mut Self>,
185        _cx: &mut std::task::Context<'_>,
186    ) -> std::task::Poll<std::io::Result<()>> {
187        if !self.closing {
188            tracing::debug!("Stream::poll_close - Starting to close stream (write side)");
189            self.closing = true;
190
191            // Finish the sender to signal we're done writing
192            if let Some(mut sender) = self.sender.take() {
193                if let Err(e) = sender.finish() {
194                    tracing::warn!("Stream::poll_close - Error finishing sender: {}", e);
195                } else {
196                    tracing::debug!("Stream::poll_close - Sender finished successfully");
197                }
198            }
199        }
200        tracing::debug!("Stream::poll_close - Write side closed");
201        std::task::Poll::Ready(Ok(()))
202    }
203}