saluki_io/net/
stream.rs

1use std::{
2    io,
3    net::SocketAddr,
4    ops::Deref,
5    pin::Pin,
6    sync::Arc,
7    task::{Context, Poll},
8};
9
10use bytes::BufMut;
11use pin_project::pin_project;
12use saluki_core::runtime::state::Sublease;
13#[cfg(windows)]
14use tokio::net::windows::named_pipe::NamedPipeServer;
15#[cfg(unix)]
16use tokio::net::{UnixDatagram, UnixStream};
17use tokio::{
18    io::{AsyncRead, AsyncReadExt as _, AsyncWrite, ReadBuf},
19    net::{TcpStream, UdpSocket},
20};
21
22use super::addr::{ConnectionAddress, ProcessIdentity};
23#[cfg(unix)]
24use super::unix::{unix_recvmsg, unixgram_recvmsg};
25
26/// A connection-oriented socket.
27///
28/// This type wraps network sockets that operate in a connection-oriented manner, such as TCP or Unix domain sockets in
29/// stream mode.
30#[pin_project(project = ConnectionProjected)]
31pub enum Connection {
32    /// A TCP socket.
33    Tcp(#[pin] TcpStream, SocketAddr),
34
35    /// A Unix domain socket in stream mode (SOCK_STREAM).
36    #[cfg(unix)]
37    Unix(#[pin] tokio::net::UnixStream),
38
39    /// A Windows named pipe in byte stream mode.
40    #[cfg(windows)]
41    NamedPipe(#[pin] NamedPipeServer),
42}
43
44impl Connection {
45    async fn receive<B: BufMut>(&mut self, buf: &mut B) -> io::Result<(usize, ConnectionAddress)> {
46        match self {
47            Self::Tcp(inner, addr) => inner.read_buf(buf).await.map(|n| (n, (*addr).into())),
48            #[cfg(unix)]
49            Self::Unix(inner) => unix_recvmsg(inner, buf).await,
50            #[cfg(windows)]
51            Self::NamedPipe(inner) => inner
52                .read_buf(buf)
53                .await
54                .map(|n| (n, ConnectionAddress::ProcessLike(ProcessIdentity::Unavailable))),
55        }
56    }
57
58    pub(super) fn remote_addr(&self) -> ConnectionAddress {
59        match self {
60            Self::Tcp(_, addr) => ConnectionAddress::SocketLike(*addr),
61            #[cfg(unix)]
62            Self::Unix(_) => ConnectionAddress::ProcessLike(ProcessIdentity::Unavailable),
63            #[cfg(windows)]
64            Self::NamedPipe(_) => ConnectionAddress::ProcessLike(ProcessIdentity::Unavailable),
65        }
66    }
67}
68
69impl AsyncRead for Connection {
70    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
71        match self.project() {
72            ConnectionProjected::Tcp(inner, _) => inner.poll_read(cx, buf),
73            #[cfg(unix)]
74            ConnectionProjected::Unix(inner) => inner.poll_read(cx, buf),
75            #[cfg(windows)]
76            ConnectionProjected::NamedPipe(inner) => inner.poll_read(cx, buf),
77        }
78    }
79}
80
81impl AsyncWrite for Connection {
82    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
83        match self.project() {
84            ConnectionProjected::Tcp(inner, _) => inner.poll_write(cx, buf),
85            #[cfg(unix)]
86            ConnectionProjected::Unix(inner) => inner.poll_write(cx, buf),
87            #[cfg(windows)]
88            ConnectionProjected::NamedPipe(inner) => inner.poll_write(cx, buf),
89        }
90    }
91
92    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
93        match self.project() {
94            ConnectionProjected::Tcp(inner, _) => inner.poll_flush(cx),
95            #[cfg(unix)]
96            ConnectionProjected::Unix(inner) => inner.poll_flush(cx),
97            #[cfg(windows)]
98            ConnectionProjected::NamedPipe(inner) => inner.poll_flush(cx),
99        }
100    }
101
102    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
103        match self.project() {
104            ConnectionProjected::Tcp(inner, _) => inner.poll_shutdown(cx),
105            #[cfg(unix)]
106            ConnectionProjected::Unix(inner) => inner.poll_shutdown(cx),
107            #[cfg(windows)]
108            ConnectionProjected::NamedPipe(inner) => inner.poll_shutdown(cx),
109        }
110    }
111}
112
113/// A connectionless socket leased by a listener.
114///
115/// The socket is shared with the listener that yielded it rather than owned outright: for connectionless families the
116/// bound socket *is* the stream, so moving it out would take it away from the listener, and a listener owned by a
117/// [`ResourceRegistry`][saluki_core::runtime::state::ResourceRegistry] has to keep the sockets it was created with.
118///
119/// The sublease is what stops that sharing from becoming a hazard. While this value lives, the listener's lease is
120/// still active no matter what its holder does, so the registry won't hand the listener to another acquirer whose
121/// streams would then read from the same socket.
122pub(crate) struct SubleasedSocket<T> {
123    /// The bound socket, shared with the listener that yielded it.
124    ///
125    /// Declared before the sublease, and that order matters: fields drop in declaration order, so the socket closes
126    /// before the sublease is returned. Returning the sublease first would tell the registry the socket is gone while
127    /// it is still open, and a discarded listener is rebuilt the instant its last sublease comes back -- which would
128    /// bind the replacement alongside this socket rather than after it.
129    socket: Arc<T>,
130
131    /// The sublease held for as long as this socket is in use.
132    ///
133    /// `None` for a socket that didn't come from a registry-managed listener, which has no lease to sublet.
134    _sublease: Option<Sublease>,
135}
136
137impl<T> SubleasedSocket<T> {
138    pub(crate) fn new(socket: Arc<T>, sublease: Option<Sublease>) -> Self {
139        Self {
140            socket,
141            _sublease: sublease,
142        }
143    }
144}
145
146impl<T> Deref for SubleasedSocket<T> {
147    type Target = T;
148
149    fn deref(&self) -> &Self::Target {
150        &self.socket
151    }
152}
153
154/// A connectionless socket.
155///
156/// This type wraps network sockets that operate in a connectionless manner, such as UDP or Unix domain sockets in
157/// datagram mode.
158enum Connectionless {
159    /// A UDP socket.
160    Udp(SubleasedSocket<UdpSocket>),
161
162    /// A Unix domain socket in datagram mode (SOCK_DGRAM).
163    #[cfg(unix)]
164    Unixgram(SubleasedSocket<tokio::net::UnixDatagram>),
165}
166
167impl Connectionless {
168    async fn receive<B: BufMut>(&mut self, buf: &mut B) -> io::Result<(usize, ConnectionAddress)> {
169        match self {
170            Self::Udp(inner) => inner.recv_buf_from(buf).await.map(|(n, addr)| (n, addr.into())),
171            #[cfg(unix)]
172            Self::Unixgram(inner) => unixgram_recvmsg(inner, buf).await,
173        }
174    }
175}
176
177enum StreamInner {
178    Connection { socket: Connection },
179    Connectionless { socket: Connectionless },
180}
181
182/// A network stream.
183///
184/// `Stream` provides an abstraction over connectionless and connection-oriented network sockets. In many cases, it's
185/// not required to know the exact socket family (for example, TCP, UDP, Unix domain socket) that's being used, and it can be
186/// beneficial to allow abstracting over the differences to facilitate simpler code.
187///
188/// # Connection-oriented mode
189///
190/// In connection-oriented mode, the stream is backed by a socket that operates in a connection-oriented manner, which
191/// ensures a reliable, ordered stream of messages to and from the remote peer.
192///
193/// The connection address returned when receiving data _should_ be stable for the life of the `Stream`.
194///
195/// # Connectionless mode
196///
197/// In connectionless mode, the stream is backed by a socket that operates in a connectionless manner, which doesn't
198/// provide any assurances around reliability and ordering of messages to and from the remote peer. While a stream might
199/// be backed by a Unix domain socket in datagram mode, which _does_ provide reliability of messages, this can't and
200/// shouldn't be relied upon when using `Stream`.
201pub struct Stream {
202    inner: StreamInner,
203}
204
205impl Stream {
206    /// Returns `true` if the stream is connectionless.
207    pub fn is_connectionless(&self) -> bool {
208        matches!(self.inner, StreamInner::Connectionless { .. })
209    }
210
211    /// Receives data from the stream.
212    ///
213    /// On success, returns the number of bytes read and the address from whence the data came.
214    ///
215    /// # Errors
216    ///
217    /// If the underlying system call fails, an error is returned.
218    pub async fn receive<B: BufMut>(&mut self, buf: &mut B) -> io::Result<(usize, ConnectionAddress)> {
219        match &mut self.inner {
220            StreamInner::Connection { socket } => socket.receive(buf).await,
221            StreamInner::Connectionless { socket } => socket.receive(buf).await,
222        }
223    }
224
225    #[cfg(test)]
226    pub(crate) fn recv_buffer_size(&self) -> io::Result<usize> {
227        match &self.inner {
228            StreamInner::Connection { socket } => match socket {
229                Connection::Tcp(inner, _) => socket2::SockRef::from(inner).recv_buffer_size(),
230                #[cfg(unix)]
231                Connection::Unix(inner) => socket2::SockRef::from(inner).recv_buffer_size(),
232                #[cfg(windows)]
233                Connection::NamedPipe(_) => Ok(0),
234            },
235            StreamInner::Connectionless { socket } => match socket {
236                Connectionless::Udp(inner) => socket2::SockRef::from(&**inner).recv_buffer_size(),
237                #[cfg(unix)]
238                Connectionless::Unixgram(inner) => socket2::SockRef::from(&**inner).recv_buffer_size(),
239            },
240        }
241    }
242}
243
244impl From<(TcpStream, SocketAddr)> for Stream {
245    fn from((stream, remote_addr): (TcpStream, SocketAddr)) -> Self {
246        Self {
247            inner: StreamInner::Connection {
248                socket: Connection::Tcp(stream, remote_addr),
249            },
250        }
251    }
252}
253
254impl From<SubleasedSocket<UdpSocket>> for Stream {
255    fn from(socket: SubleasedSocket<UdpSocket>) -> Self {
256        Self {
257            inner: StreamInner::Connectionless {
258                socket: Connectionless::Udp(socket),
259            },
260        }
261    }
262}
263
264impl From<Arc<UdpSocket>> for Stream {
265    fn from(socket: Arc<UdpSocket>) -> Self {
266        Self::from(SubleasedSocket::new(socket, None))
267    }
268}
269
270impl From<UdpSocket> for Stream {
271    fn from(socket: UdpSocket) -> Self {
272        Self::from(Arc::new(socket))
273    }
274}
275
276#[cfg(unix)]
277impl From<SubleasedSocket<UnixDatagram>> for Stream {
278    fn from(socket: SubleasedSocket<UnixDatagram>) -> Self {
279        Self {
280            inner: StreamInner::Connectionless {
281                socket: Connectionless::Unixgram(socket),
282            },
283        }
284    }
285}
286
287#[cfg(unix)]
288impl From<Arc<UnixDatagram>> for Stream {
289    fn from(socket: Arc<UnixDatagram>) -> Self {
290        Self::from(SubleasedSocket::new(socket, None))
291    }
292}
293
294#[cfg(unix)]
295impl From<UnixDatagram> for Stream {
296    fn from(socket: UnixDatagram) -> Self {
297        Self::from(Arc::new(socket))
298    }
299}
300
301#[cfg(unix)]
302impl From<UnixStream> for Stream {
303    fn from(stream: UnixStream) -> Self {
304        Self {
305            inner: StreamInner::Connection {
306                socket: Connection::Unix(stream),
307            },
308        }
309    }
310}
311
312#[cfg(windows)]
313impl From<NamedPipeServer> for Stream {
314    fn from(stream: NamedPipeServer) -> Self {
315        Self {
316            inner: StreamInner::Connection {
317                socket: Connection::NamedPipe(stream),
318            },
319        }
320    }
321}
saluki_io/net/stream.rs

saluki_io/net/
stream.rs
