Skip to main content

zlink_smol/unix/
stream.rs

1use crate::{
2    Result,
3    connection::socket::{self, Socket},
4};
5use async_io::Async;
6use std::{
7    os::{
8        fd::{AsFd, BorrowedFd},
9        unix::net::UnixStream as StdUnixStream,
10    },
11    sync::Arc,
12};
13use zlink_core::connection::socket::ReadResult;
14
15/// The connection type that uses Unix Domain Sockets for transport.
16pub type Connection = crate::Connection<Stream>;
17
18/// Connect to Unix Domain Socket at the given path.
19pub async fn connect<P>(path: P) -> Result<Connection>
20where
21    P: AsRef<std::path::Path>,
22{
23    Async::<StdUnixStream>::connect(path)
24        .await
25        .map_err(Into::into)
26        .and_then(TryInto::try_into)
27        .map(Connection::new)
28}
29
30/// The [`Socket`] implementation using Unix Domain Sockets.
31#[derive(Debug)]
32pub struct Stream(Async<StdUnixStream>);
33
34impl Socket for Stream {
35    type ReadHalf = ReadHalf;
36    type WriteHalf = WriteHalf;
37
38    const CAN_TRANSFER_FDS: bool = true;
39
40    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
41        let stream = Arc::new(self.0);
42
43        (ReadHalf(Arc::clone(&stream)), WriteHalf(stream))
44    }
45}
46
47impl TryFrom<Async<StdUnixStream>> for Stream {
48    type Error = crate::Error;
49
50    fn try_from(stream: Async<StdUnixStream>) -> Result<Self> {
51        #[cfg(target_os = "linux")]
52        zlink_core::unix_utils::enable_passcred(&stream)?;
53        Ok(Self(stream))
54    }
55}
56
57impl TryFrom<StdUnixStream> for Stream {
58    type Error = crate::Error;
59
60    fn try_from(stream: StdUnixStream) -> Result<Self> {
61        stream.set_nonblocking(true)?;
62        Async::new(stream)?.try_into()
63    }
64}
65
66impl AsFd for Stream {
67    fn as_fd(&self) -> BorrowedFd<'_> {
68        self.0.as_fd()
69    }
70}
71
72impl socket::UnixSocket for Stream {}
73
74/// The [`ReadHalf`] implementation using Unix Domain Sockets.
75#[derive(Debug)]
76pub struct ReadHalf(Arc<Async<StdUnixStream>>);
77
78impl socket::ReadHalf for ReadHalf {
79    async fn read(&mut self, buf: &mut [u8]) -> Result<ReadResult> {
80        use std::{future::poll_fn, task::Poll};
81
82        poll_fn(|cx| {
83            loop {
84                match crate::unix_utils::recvmsg(self.0.as_ref(), buf) {
85                    Ok(result) => return Poll::Ready(Ok(result)),
86                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
87                        match self.0.poll_readable(cx) {
88                            Poll::Pending => return Poll::Pending,
89                            Poll::Ready(res) => res?,
90                        }
91                    }
92                    Err(e) => return Poll::Ready(Err(e.into())),
93                }
94            }
95        })
96        .await
97    }
98}
99
100impl AsFd for ReadHalf {
101    fn as_fd(&self) -> BorrowedFd<'_> {
102        self.0.as_ref().as_fd()
103    }
104}
105
106impl socket::UnixSocket for ReadHalf {}
107
108/// The [`WriteHalf`] implementation using Unix Domain Sockets.
109#[derive(Debug)]
110pub struct WriteHalf(Arc<Async<StdUnixStream>>);
111
112impl socket::WriteHalf for WriteHalf {
113    async fn write(
114        &mut self,
115        buf: &[u8],
116        fds: &[impl AsFd],
117        #[cfg(target_os = "linux")] creds: Option<&crate::connection::PassedCredentials>,
118    ) -> Result<()> {
119        use std::{future::poll_fn, task::Poll};
120
121        // Convert to BorrowedFd for rustix.
122        let borrowed_fds: Vec<BorrowedFd<'_>> = fds.iter().map(|f| f.as_fd()).collect();
123
124        let mut pos = 0;
125        while pos < buf.len() {
126            // Use FDs on first write, empty slice on subsequent writes.
127            let fds_to_send = if pos == 0 { &borrowed_fds[..] } else { &[] };
128
129            let n: usize = poll_fn(|cx| {
130                loop {
131                    match crate::unix_utils::sendmsg(
132                        self.0.as_ref(),
133                        &buf[pos..],
134                        fds_to_send,
135                        #[cfg(target_os = "linux")]
136                        creds,
137                    ) {
138                        Ok(bytes_sent) => return Poll::Ready(Ok::<_, crate::Error>(bytes_sent)),
139                        Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
140                            match self.0.poll_writable(cx) {
141                                Poll::Pending => return Poll::Pending,
142                                Poll::Ready(res) => res?,
143                            }
144                        }
145                        Err(e) => return Poll::Ready(Err(e.into())),
146                    }
147                }
148            })
149            .await?;
150
151            pos += n;
152        }
153
154        Ok(())
155    }
156}
157
158impl AsFd for WriteHalf {
159    fn as_fd(&self) -> BorrowedFd<'_> {
160        self.0.as_ref().as_fd()
161    }
162}
163
164impl socket::UnixSocket for WriteHalf {}