Skip to main content

nuclei/net/
split.rs

1//! A stream's two halves, to read on one task and write on another:
2//! borrowed ([`ReadHalf`], [`WriteHalf`], from `split`, allocating nothing)
3//! or owned ([`OwnedReadHalf`], [`OwnedWriteHalf`], from `into_split`, which
4//! share the stream in one `Arc`), and owned halves made one stream again
5//! with [`OwnedReadHalf::reunite`].
6//!
7//! A stream is read and written through a shared reference already, so the
8//! halves only keep each side to its task. Dropping the write half does not
9//! end the write side, as bytes may still wait in the write-behind:
10//! `close` (`AsyncWriteExt::close`) sends them, then ends it.
11
12use core::pin::Pin;
13use core::task::{Context, Poll};
14use std::fmt;
15use std::io;
16use std::sync::Arc;
17
18use futures_io::{AsyncRead, AsyncWrite};
19
20use crate::io::Handle;
21
22/// The read half of a stream, borrowed from it.
23pub struct ReadHalf<'a, S>(pub(super) &'a Handle<S>);
24
25/// The write half of a stream, borrowed from it.
26pub struct WriteHalf<'a, S>(pub(super) &'a Handle<S>);
27
28/// The read half of a stream, owning it with the write half.
29pub struct OwnedReadHalf<S>(pub(super) Arc<Handle<S>>);
30
31/// The write half of a stream, owning it with the read half.
32pub struct OwnedWriteHalf<S>(pub(super) Arc<Handle<S>>);
33
34/// Two halves of different streams, given back by
35/// [`OwnedReadHalf::reunite`].
36pub struct ReuniteError<S>(pub OwnedReadHalf<S>, pub OwnedWriteHalf<S>);
37
38impl<S> OwnedReadHalf<S> {
39    /// The stream again, from its two halves.
40    ///
41    /// # Errors
42    ///
43    /// When `other` is a half of another stream; both are given back.
44    pub fn reunite(self, other: OwnedWriteHalf<S>) -> Result<Handle<S>, ReuniteError<S>> {
45        if !Arc::ptr_eq(&self.0, &other.0) {
46            return Err(ReuniteError(self, other));
47        }
48        drop(other);
49        // The two were the only holders (halves are not cloned): the stream
50        // is this one's alone now.
51        Arc::try_unwrap(self.0)
52            .map_err(|both| ReuniteError(OwnedReadHalf(Arc::clone(&both)), OwnedWriteHalf(both)))
53    }
54}
55
56impl<S> OwnedWriteHalf<S> {
57    /// The stream again, from its two halves.
58    ///
59    /// # Errors
60    ///
61    /// When `other` is a half of another stream; both are given back.
62    pub fn reunite(self, other: OwnedReadHalf<S>) -> Result<Handle<S>, ReuniteError<S>> {
63        other.reunite(self)
64    }
65}
66
67impl<S> AsRef<Handle<S>> for ReadHalf<'_, S> {
68    fn as_ref(&self) -> &Handle<S> {
69        self.0
70    }
71}
72
73impl<S> AsRef<Handle<S>> for WriteHalf<'_, S> {
74    fn as_ref(&self) -> &Handle<S> {
75        self.0
76    }
77}
78
79impl<S> AsRef<Handle<S>> for OwnedReadHalf<S> {
80    fn as_ref(&self) -> &Handle<S> {
81        &self.0
82    }
83}
84
85impl<S> AsRef<Handle<S>> for OwnedWriteHalf<S> {
86    fn as_ref(&self) -> &Handle<S> {
87        &self.0
88    }
89}
90
91macro_rules! debug_half {
92    ($($half:ident $(<$lt:lifetime>)?),*) => {$(
93        impl<$($lt,)? S: fmt::Debug> fmt::Debug for $half<$($lt,)? S> {
94            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
95                f.debug_tuple(stringify!($half)).field(self.as_ref()).finish()
96            }
97        }
98    )*};
99}
100
101debug_half!(ReadHalf<'a>, WriteHalf<'a>, OwnedReadHalf, OwnedWriteHalf);
102
103impl<S> fmt::Debug for ReuniteError<S> {
104    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
105        f.write_str("ReuniteError(..)")
106    }
107}
108
109impl<S> fmt::Display for ReuniteError<S> {
110    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
111        f.write_str("the two halves are of different streams")
112    }
113}
114
115impl<S> std::error::Error for ReuniteError<S> {}
116
117macro_rules! read_half {
118    ($($half:ident $(<$lt:lifetime>)?),*) => {$(
119        impl<$($lt,)? S> AsyncRead for $half<$($lt,)? S>
120        where
121            for<'x> &'x Handle<S>: AsyncRead,
122        {
123            fn poll_read(
124                self: Pin<&mut Self>,
125                cx: &mut Context<'_>,
126                buf: &mut [u8],
127            ) -> Poll<io::Result<usize>> {
128                Pin::new(&mut self.get_mut().as_ref()).poll_read(cx, buf)
129            }
130
131            fn poll_read_vectored(
132                self: Pin<&mut Self>,
133                cx: &mut Context<'_>,
134                bufs: &mut [io::IoSliceMut<'_>],
135            ) -> Poll<io::Result<usize>> {
136                Pin::new(&mut self.get_mut().as_ref()).poll_read_vectored(cx, bufs)
137            }
138        }
139    )*};
140}
141
142macro_rules! write_half {
143    ($($half:ident $(<$lt:lifetime>)?),*) => {$(
144        impl<$($lt,)? S> AsyncWrite for $half<$($lt,)? S>
145        where
146            for<'x> &'x Handle<S>: AsyncWrite,
147        {
148            fn poll_write(
149                self: Pin<&mut Self>,
150                cx: &mut Context<'_>,
151                buf: &[u8],
152            ) -> Poll<io::Result<usize>> {
153                Pin::new(&mut self.get_mut().as_ref()).poll_write(cx, buf)
154            }
155
156            fn poll_write_vectored(
157                self: Pin<&mut Self>,
158                cx: &mut Context<'_>,
159                bufs: &[io::IoSlice<'_>],
160            ) -> Poll<io::Result<usize>> {
161                Pin::new(&mut self.get_mut().as_ref()).poll_write_vectored(cx, bufs)
162            }
163
164            fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
165                Pin::new(&mut self.get_mut().as_ref()).poll_flush(cx)
166            }
167
168            fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
169                Pin::new(&mut self.get_mut().as_ref()).poll_close(cx)
170            }
171        }
172    )*};
173}
174
175read_half!(ReadHalf<'a>, OwnedReadHalf);
176write_half!(WriteHalf<'a>, OwnedWriteHalf);