zlink_tokio/unix/
stream.rs1use 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
12pub type Connection = crate::Connection<Stream>;
14
15pub 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#[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#[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#[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 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 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 #[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 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 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}