saluki_io/net/
stream.rs

1use std::{
2    io,
3    net::SocketAddr,
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use bytes::BufMut;
9use pin_project::pin_project;
10#[cfg(windows)]
11use tokio::net::windows::named_pipe::NamedPipeServer;
12use tokio::{
13    io::{AsyncRead, AsyncReadExt as _, AsyncWrite, ReadBuf},
14    net::{TcpStream, UdpSocket},
15};
16
17use super::addr::{ConnectionAddress, ProcessIdentity};
18#[cfg(unix)]
19use super::unix::{unix_recvmsg, unixgram_recvmsg};
20
21/// A connection-oriented socket.
22///
23/// This type wraps network sockets that operate in a connection-oriented manner, such as TCP or Unix domain sockets in
24/// stream mode.
25#[pin_project(project = ConnectionProjected)]
26pub enum Connection {
27    /// A TCP socket.
28    Tcp(#[pin] TcpStream, SocketAddr),
29
30    /// A Unix domain socket in stream mode (SOCK_STREAM).
31    #[cfg(unix)]
32    Unix(#[pin] tokio::net::UnixStream),
33
34    /// A Windows named pipe in byte stream mode.
35    #[cfg(windows)]
36    NamedPipe(#[pin] NamedPipeServer),
37}
38
39impl Connection {
40    async fn receive<B: BufMut>(&mut self, buf: &mut B) -> io::Result<(usize, ConnectionAddress)> {
41        match self {
42            Self::Tcp(inner, addr) => inner.read_buf(buf).await.map(|n| (n, (*addr).into())),
43            #[cfg(unix)]
44            Self::Unix(inner) => unix_recvmsg(inner, buf).await,
45            #[cfg(windows)]
46            Self::NamedPipe(inner) => inner
47                .read_buf(buf)
48                .await
49                .map(|n| (n, ConnectionAddress::ProcessLike(ProcessIdentity::Unavailable))),
50        }
51    }
52
53    pub(super) fn remote_addr(&self) -> ConnectionAddress {
54        match self {
55            Self::Tcp(_, addr) => ConnectionAddress::SocketLike(*addr),
56            #[cfg(unix)]
57            Self::Unix(_) => ConnectionAddress::ProcessLike(ProcessIdentity::Unavailable),
58            #[cfg(windows)]
59            Self::NamedPipe(_) => ConnectionAddress::ProcessLike(ProcessIdentity::Unavailable),
60        }
61    }
62}
63
64impl AsyncRead for Connection {
65    fn poll_read(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<io::Result<()>> {
66        match self.project() {
67            ConnectionProjected::Tcp(inner, _) => inner.poll_read(cx, buf),
68            #[cfg(unix)]
69            ConnectionProjected::Unix(inner) => inner.poll_read(cx, buf),
70            #[cfg(windows)]
71            ConnectionProjected::NamedPipe(inner) => inner.poll_read(cx, buf),
72        }
73    }
74}
75
76impl AsyncWrite for Connection {
77    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<io::Result<usize>> {
78        match self.project() {
79            ConnectionProjected::Tcp(inner, _) => inner.poll_write(cx, buf),
80            #[cfg(unix)]
81            ConnectionProjected::Unix(inner) => inner.poll_write(cx, buf),
82            #[cfg(windows)]
83            ConnectionProjected::NamedPipe(inner) => inner.poll_write(cx, buf),
84        }
85    }
86
87    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
88        match self.project() {
89            ConnectionProjected::Tcp(inner, _) => inner.poll_flush(cx),
90            #[cfg(unix)]
91            ConnectionProjected::Unix(inner) => inner.poll_flush(cx),
92            #[cfg(windows)]
93            ConnectionProjected::NamedPipe(inner) => inner.poll_flush(cx),
94        }
95    }
96
97    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
98        match self.project() {
99            ConnectionProjected::Tcp(inner, _) => inner.poll_shutdown(cx),
100            #[cfg(unix)]
101            ConnectionProjected::Unix(inner) => inner.poll_shutdown(cx),
102            #[cfg(windows)]
103            ConnectionProjected::NamedPipe(inner) => inner.poll_shutdown(cx),
104        }
105    }
106}
107
108/// A connectionless socket.
109///
110/// This type wraps network sockets that operate in a connectionless manner, such as UDP or Unix domain sockets in
111/// datagram mode.
112enum Connectionless {
113    /// A UDP socket.
114    Udp(UdpSocket),
115
116    /// A Unix domain socket in datagram mode (SOCK_DGRAM).
117    #[cfg(unix)]
118    Unixgram(tokio::net::UnixDatagram),
119}
120
121impl Connectionless {
122    async fn receive<B: BufMut>(&mut self, buf: &mut B) -> io::Result<(usize, ConnectionAddress)> {
123        match self {
124            Self::Udp(inner) => inner.recv_buf_from(buf).await.map(|(n, addr)| (n, addr.into())),
125            #[cfg(unix)]
126            Self::Unixgram(inner) => unixgram_recvmsg(inner, buf).await,
127        }
128    }
129}
130
131enum StreamInner {
132    Connection { socket: Connection },
133    Connectionless { socket: Connectionless },
134}
135
136/// A network stream.
137///
138/// `Stream` provides an abstraction over connectionless and connection-oriented network sockets. In many cases, it's
139/// not required to know the exact socket family (for example, TCP, UDP, Unix domain socket) that's being used, and it can be
140/// beneficial to allow abstracting over the differences to facilitate simpler code.
141///
142/// ## Connection-oriented mode
143///
144/// In connection-oriented mode, the stream is backed by a socket that operates in a connection-oriented manner, which
145/// ensures a reliable, ordered stream of messages to and from the remote peer.
146///
147/// The connection address returned when receiving data _should_ be stable for the life of the `Stream`.
148///
149/// ## Connectionless mode
150///
151/// In connectionless mode, the stream is backed by a socket that operates in a connectionless manner, which doesn't
152/// provide any assurances around reliability and ordering of messages to and from the remote peer. While a stream might
153/// be backed by a Unix domain socket in datagram mode, which _does_ provide reliability of messages, this can't and
154/// shouldn't be relied upon when using `Stream`.
155pub struct Stream {
156    inner: StreamInner,
157}
158
159impl Stream {
160    /// Returns `true` if the stream is connectionless.
161    pub fn is_connectionless(&self) -> bool {
162        matches!(self.inner, StreamInner::Connectionless { .. })
163    }
164
165    /// Receives data from the stream.
166    ///
167    /// On success, returns the number of bytes read and the address from whence the data came.
168    ///
169    /// ## Errors
170    ///
171    /// If the underlying system call fails, an error is returned.
172    pub async fn receive<B: BufMut>(&mut self, buf: &mut B) -> io::Result<(usize, ConnectionAddress)> {
173        match &mut self.inner {
174            StreamInner::Connection { socket } => socket.receive(buf).await,
175            StreamInner::Connectionless { socket } => socket.receive(buf).await,
176        }
177    }
178
179    #[cfg(test)]
180    pub(crate) fn recv_buffer_size(&self) -> io::Result<usize> {
181        match &self.inner {
182            StreamInner::Connection { socket } => match socket {
183                Connection::Tcp(inner, _) => socket2::SockRef::from(inner).recv_buffer_size(),
184                #[cfg(unix)]
185                Connection::Unix(inner) => socket2::SockRef::from(inner).recv_buffer_size(),
186                #[cfg(windows)]
187                Connection::NamedPipe(_) => Ok(0),
188            },
189            StreamInner::Connectionless { socket } => match socket {
190                Connectionless::Udp(inner) => socket2::SockRef::from(inner).recv_buffer_size(),
191                #[cfg(unix)]
192                Connectionless::Unixgram(inner) => socket2::SockRef::from(inner).recv_buffer_size(),
193            },
194        }
195    }
196}
197
198impl From<(TcpStream, SocketAddr)> for Stream {
199    fn from((stream, remote_addr): (TcpStream, SocketAddr)) -> Self {
200        Self {
201            inner: StreamInner::Connection {
202                socket: Connection::Tcp(stream, remote_addr),
203            },
204        }
205    }
206}
207
208impl From<UdpSocket> for Stream {
209    fn from(socket: UdpSocket) -> Self {
210        Self {
211            inner: StreamInner::Connectionless {
212                socket: Connectionless::Udp(socket),
213            },
214        }
215    }
216}
217
218#[cfg(unix)]
219impl From<tokio::net::UnixDatagram> for Stream {
220    fn from(socket: tokio::net::UnixDatagram) -> Self {
221        Self {
222            inner: StreamInner::Connectionless {
223                socket: Connectionless::Unixgram(socket),
224            },
225        }
226    }
227}
228
229#[cfg(unix)]
230impl From<tokio::net::UnixStream> for Stream {
231    fn from(stream: tokio::net::UnixStream) -> Self {
232        Self {
233            inner: StreamInner::Connection {
234                socket: Connection::Unix(stream),
235            },
236        }
237    }
238}
239
240#[cfg(windows)]
241impl From<NamedPipeServer> for Stream {
242    fn from(stream: NamedPipeServer) -> Self {
243        Self {
244            inner: StreamInner::Connection {
245                socket: Connection::NamedPipe(stream),
246            },
247        }
248    }
249}
saluki_io/net/stream.rs

saluki_io/net/
stream.rs
