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

saluki_io/net/
stream.rs
