Skip to main content

zlink_tokio/unix/
stream.rs

1use crate::{
2    Result,
3    connection::socket::{self, Socket},
4};
5use std::os::{
6    fd::{AsFd, BorrowedFd},
7    unix::net::UnixStream as StdUnixStream,
8};
9use tokio::net::{UnixStream, unix};
10use zlink_core::connection::socket::ReadResult;
11
12/// The connection type that uses Unix Domain Sockets for transport.
13pub type Connection = crate::Connection<Stream>;
14
15/// Connect to Unix Domain Socket at the given path.
16pub async fn connect<P>(path: P) -> Result<Connection>
17where
18    P: AsRef<std::path::Path>,
19{
20    UnixStream::connect(path)
21        .await
22        .map_err(Into::into)
23        .and_then(TryInto::try_into)
24        .map(Connection::new)
25}
26
27/// The [`Socket`] implementation using Unix Domain Sockets.
28#[derive(Debug)]
29pub struct Stream(UnixStream);
30
31impl Socket for Stream {
32    type ReadHalf = ReadHalf;
33    type WriteHalf = WriteHalf;
34
35    const CAN_TRANSFER_FDS: bool = true;
36
37    fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
38        let (read, write) = self.0.into_split();
39
40        (ReadHalf(read), WriteHalf(write))
41    }
42}
43
44impl TryFrom<UnixStream> for Stream {
45    type Error = crate::Error;
46
47    fn try_from(stream: UnixStream) -> Result<Self> {
48        #[cfg(target_os = "linux")]
49        zlink_core::unix_utils::enable_passcred(&stream)?;
50        Ok(Self(stream))
51    }
52}
53
54impl TryFrom<StdUnixStream> for Stream {
55    type Error = crate::Error;
56
57    fn try_from(stream: StdUnixStream) -> Result<Self> {
58        stream.set_nonblocking(true)?;
59        UnixStream::from_std(stream)
60            .map_err(Into::into)
61            .and_then(TryInto::try_into)
62    }
63}
64
65impl socket::UnixSocket for Stream {}
66
67impl AsFd for Stream {
68    fn as_fd(&self) -> BorrowedFd<'_> {
69        self.0.as_fd()
70    }
71}
72
73/// The [`ReadHalf`] implementation using Unix Domain Sockets.
74#[derive(Debug)]
75pub struct ReadHalf(unix::OwnedReadHalf);
76
77impl socket::ReadHalf for ReadHalf {
78    async fn read(&mut self, buf: &mut [u8]) -> Result<ReadResult> {
79        use std::{future::poll_fn, task::Poll};
80
81        poll_fn(|cx| {
82            loop {
83                let stream: &UnixStream = self.0.as_ref();
84                match stream.try_io(tokio::io::Interest::READABLE, || {
85                    crate::unix_utils::recvmsg(stream, buf)
86                }) {
87                    Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
88                        match stream.poll_read_ready(cx) {
89                            Poll::Pending => return Poll::Pending,
90                            Poll::Ready(res) => res?,
91                        }
92                    }
93                    v => return Poll::Ready(v.map_err(Into::into)),
94                }
95            }
96        })
97        .await
98    }
99}
100
101impl AsFd for ReadHalf {
102    fn as_fd(&self) -> BorrowedFd<'_> {
103        let stream: &UnixStream = self.0.as_ref();
104        stream.as_fd()
105    }
106}
107
108impl socket::UnixSocket for ReadHalf {}
109
110/// The [`WriteHalf`] implementation using Unix Domain Sockets.
111#[derive(Debug)]
112pub struct WriteHalf(unix::OwnedWriteHalf);
113
114impl socket::WriteHalf for WriteHalf {
115    async fn write(
116        &mut self,
117        buf: &[u8],
118        fds: &[impl AsFd],
119        #[cfg(target_os = "linux")] creds: Option<&crate::connection::PassedCredentials>,
120    ) -> Result<()> {
121        use std::{future::poll_fn, task::Poll};
122
123        // Convert to BorrowedFd for rustix.
124        let borrowed_fds: Vec<BorrowedFd<'_>> = fds.iter().map(|f| f.as_fd()).collect();
125
126        let mut pos = 0;
127        while pos < buf.len() {
128            // Use FDs on first write, empty slice on subsequent writes.
129            let fds_to_send = if pos == 0 { &borrowed_fds[..] } else { &[] };
130
131            let n: usize = poll_fn(|cx| {
132                loop {
133                    let stream: &UnixStream = self.0.as_ref();
134                    match stream.try_io(tokio::io::Interest::WRITABLE, || {
135                        crate::unix_utils::sendmsg(
136                            stream,
137                            &buf[pos..],
138                            fds_to_send,
139                            #[cfg(target_os = "linux")]
140                            creds,
141                        )
142                    }) {
143                        Ok(bytes_sent) => return Poll::Ready(Ok::<_, crate::Error>(bytes_sent)),
144                        Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
145                            match stream.poll_write_ready(cx) {
146                                Poll::Pending => return Poll::Pending,
147                                Poll::Ready(res) => res?,
148                            }
149                        }
150                        Err(e) => return Poll::Ready(Err(e.into())),
151                    }
152                }
153            })
154            .await?;
155
156            pos += n;
157        }
158
159        Ok(())
160    }
161}
162
163impl AsFd for WriteHalf {
164    fn as_fd(&self) -> BorrowedFd<'_> {
165        let stream: &UnixStream = self.0.as_ref();
166        stream.as_fd()
167    }
168}
169
170impl socket::UnixSocket for WriteHalf {}
171
172#[cfg(test)]
173mod tests {
174    use super::*;
175    use std::{
176        io::Write,
177        os::{
178            fd::{FromRawFd, IntoRawFd},
179            unix::net::UnixStream as StdUnixStream,
180        },
181    };
182
183    /// Verify that FD passing works when sender and receiver use **separate** connections (as in
184    /// cross-process communication). On macOS, same-connection FD passing requires a workaround
185    /// (see `WriteConnection::held_fds`), but separate connections should work without it because
186    /// FD number reuse between `sendmsg` and `recvmsg` cannot happen across different FD tables.
187    ///
188    /// This test intentionally uses split halves (`WriteConnection::send_reply`) rather than
189    /// `Connection::send_reply`, so the `drain_held_fds` workaround is never invoked. This
190    /// proves the workaround is not needed for cross-connection FD passing.
191    #[tokio::test]
192    async fn fd_passing_across_separate_connections() {
193        let (std_a, std_b) = StdUnixStream::pair().unwrap();
194        std_a.set_nonblocking(true).unwrap();
195        std_b.set_nonblocking(true).unwrap();
196
197        let conn_a = Connection::new(UnixStream::from_std(std_a).unwrap().try_into().unwrap());
198        let conn_b = Connection::new(UnixStream::from_std(std_b).unwrap().try_into().unwrap());
199
200        let (_, mut write_a) = conn_a.split();
201        let (mut read_b, _) = conn_b.split();
202
203        // Send 3 FDs one at a time, closing the sender's copy after each sendmsg.
204        for name in ["alpha", "beta", "gamma"] {
205            let (r, mut w) = StdUnixStream::pair().unwrap();
206            w.write_all(name.as_bytes()).unwrap();
207            drop(w);
208
209            let reply = crate::Reply::new(Some(name.to_string())).set_continues(Some(false));
210            write_a.send_reply(&reply, vec![r.into()]).await.unwrap();
211        }
212
213        // Receive and verify each FD has the correct data.
214        for name in ["alpha", "beta", "gamma"] {
215            let (reply, fds) = read_b.receive_reply::<String, ()>().await.unwrap();
216            let params = reply.unwrap().into_parameters().unwrap();
217            assert_eq!(params, name);
218            assert_eq!(fds.len(), 1);
219
220            let recv_fd = fds.into_iter().next().unwrap();
221            let mut stream = unsafe { StdUnixStream::from_raw_fd(recv_fd.into_raw_fd()) };
222            let mut buf = String::new();
223            std::io::Read::read_to_string(&mut stream, &mut buf).unwrap();
224            assert_eq!(buf, name, "FD data mismatch for {name:?}");
225        }
226    }
227}