Skip to main content

codex_uds/
lib.rs

1//! Cross-platform async Unix domain socket helpers.
2
3use std::io::Result as IoResult;
4use std::path::Path;
5use std::pin::Pin;
6use std::task::Context;
7use std::task::Poll;
8
9use tokio::io::AsyncRead;
10use tokio::io::AsyncWrite;
11use tokio::io::ReadBuf;
12
13/// Creates `socket_dir` if needed and restricts it to the current user where
14/// the platform exposes Unix permissions.
15pub async fn prepare_private_socket_directory(socket_dir: impl AsRef<Path>) -> IoResult<()> {
16    platform::prepare_private_socket_directory(socket_dir.as_ref()).await
17}
18
19/// Returns whether `socket_path` points at a stale Unix socket rendezvous path.
20///
21/// On Unix this checks the file type. On Windows, `uds_windows` represents the
22/// rendezvous as a regular path, so existence is the only useful stale-path
23/// signal available.
24pub async fn is_stale_socket_path(socket_path: impl AsRef<Path>) -> IoResult<bool> {
25    platform::is_stale_socket_path(socket_path.as_ref()).await
26}
27
28/// Async Unix domain socket listener.
29pub struct UnixListener {
30    inner: platform::Listener,
31}
32
33impl UnixListener {
34    /// Binds a new listener at `socket_path`.
35    pub async fn bind(socket_path: impl AsRef<Path>) -> IoResult<Self> {
36        platform::bind_listener(socket_path.as_ref())
37            .await
38            .map(|inner| Self { inner })
39    }
40
41    /// Accepts the next incoming stream.
42    pub async fn accept(&mut self) -> IoResult<UnixStream> {
43        self.inner.accept().await.map(|inner| UnixStream { inner })
44    }
45}
46
47/// Async Unix domain socket stream.
48pub struct UnixStream {
49    inner: platform::Stream,
50}
51
52impl UnixStream {
53    /// Connects to `socket_path`.
54    pub async fn connect(socket_path: impl AsRef<Path>) -> IoResult<Self> {
55        platform::connect_stream(socket_path.as_ref())
56            .await
57            .map(|inner| Self { inner })
58    }
59}
60
61impl AsyncRead for UnixStream {
62    fn poll_read(
63        self: Pin<&mut Self>,
64        cx: &mut Context<'_>,
65        buf: &mut ReadBuf<'_>,
66    ) -> Poll<IoResult<()>> {
67        Pin::new(&mut self.get_mut().inner).poll_read(cx, buf)
68    }
69}
70
71impl AsyncWrite for UnixStream {
72    fn poll_write(self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &[u8]) -> Poll<IoResult<usize>> {
73        Pin::new(&mut self.get_mut().inner).poll_write(cx, buf)
74    }
75
76    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
77        Pin::new(&mut self.get_mut().inner).poll_flush(cx)
78    }
79
80    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
81        Pin::new(&mut self.get_mut().inner).poll_shutdown(cx)
82    }
83}
84
85#[cfg(unix)]
86mod platform {
87    use std::io;
88    use std::io::ErrorKind;
89    use std::io::Result as IoResult;
90    use std::os::unix::fs::FileTypeExt;
91    use std::os::unix::fs::PermissionsExt;
92    use std::path::Path;
93
94    use tokio::fs;
95    use tokio::net::UnixListener;
96    use tokio::net::UnixStream;
97
98    /// Owner-only access keeps the control socket directory private while
99    /// preserving owner traversal and socket path creation.
100    const SOCKET_DIR_MODE: u32 = 0o700;
101    const SOCKET_DIR_PERMISSION_BITS: u32 = 0o777;
102
103    pub(super) type Stream = UnixStream;
104
105    pub(super) struct Listener(UnixListener);
106
107    pub(super) async fn prepare_private_socket_directory(socket_dir: &Path) -> IoResult<()> {
108        let mut dir_builder = fs::DirBuilder::new();
109        dir_builder.mode(SOCKET_DIR_MODE);
110        match dir_builder.create(socket_dir).await {
111            Ok(()) => return Ok(()),
112            Err(err) if err.kind() == ErrorKind::AlreadyExists => {}
113            Err(err) => return Err(err),
114        }
115
116        let metadata = fs::symlink_metadata(socket_dir).await?;
117        if !metadata.is_dir() {
118            return Err(io::Error::new(
119                ErrorKind::AlreadyExists,
120                format!(
121                    "socket directory path exists and is not a directory: {}",
122                    socket_dir.display()
123                ),
124            ));
125        }
126
127        let permissions = metadata.permissions();
128        // The SSH-over-UDS control socket is reachable by path, so the
129        // rendezvous directory must be owner-traversable while denying
130        // group/other access; exact 0700 fixes insecure modes and unusable
131        // owner-only modes like 0600.
132        if permissions.mode() & SOCKET_DIR_PERMISSION_BITS != SOCKET_DIR_MODE {
133            fs::set_permissions(socket_dir, std::fs::Permissions::from_mode(SOCKET_DIR_MODE))
134                .await?;
135        }
136
137        Ok(())
138    }
139
140    pub(super) async fn bind_listener(socket_path: &Path) -> IoResult<Listener> {
141        UnixListener::bind(socket_path).map(Listener)
142    }
143
144    impl Listener {
145        pub(super) async fn accept(&mut self) -> IoResult<Stream> {
146            self.0.accept().await.map(|(stream, _addr)| stream)
147        }
148    }
149
150    pub(super) async fn connect_stream(socket_path: &Path) -> IoResult<Stream> {
151        UnixStream::connect(socket_path).await
152    }
153
154    pub(super) async fn is_stale_socket_path(socket_path: &Path) -> IoResult<bool> {
155        Ok(fs::symlink_metadata(socket_path)
156            .await?
157            .file_type()
158            .is_socket())
159    }
160}
161
162#[cfg(windows)]
163mod platform {
164    use std::io;
165    use std::io::Result as IoResult;
166    use std::net::Shutdown;
167    use std::ops::Deref;
168    use std::os::windows::io::AsRawSocket;
169    use std::os::windows::io::AsSocket;
170    use std::os::windows::io::BorrowedSocket;
171    use std::path::Path;
172    use std::pin::Pin;
173    use std::task::Context;
174    use std::task::Poll;
175    use std::task::ready;
176
177    use async_io::Async;
178    use tokio::io::AsyncRead;
179    use tokio::io::AsyncWrite;
180    use tokio::io::ReadBuf;
181    use tokio::task;
182    use tokio_util::compat::Compat;
183    use tokio_util::compat::FuturesAsyncReadCompatExt;
184
185    pub(super) struct Stream(Compat<Async<WindowsUnixStream>>);
186
187    pub(super) async fn prepare_private_socket_directory(socket_dir: &Path) -> IoResult<()> {
188        tokio::fs::create_dir_all(socket_dir).await
189    }
190
191    pub(super) struct Listener(Async<WindowsUnixListener>);
192
193    pub(super) async fn bind_listener(socket_path: &Path) -> IoResult<Listener> {
194        let socket_path = socket_path.to_path_buf();
195        let listener =
196            spawn_blocking_io(move || uds_windows::UnixListener::bind(socket_path)).await?;
197        Async::new(WindowsUnixListener::from(listener)).map(Listener)
198    }
199
200    impl Listener {
201        pub(super) async fn accept(&mut self) -> IoResult<Stream> {
202            let (stream, _addr) = self.0.read_with(|listener| listener.accept()).await?;
203            Async::new(WindowsUnixStream::from(stream))
204                .map(FuturesAsyncReadCompatExt::compat)
205                .map(Stream)
206        }
207    }
208
209    pub(super) async fn connect_stream(socket_path: &Path) -> IoResult<Stream> {
210        let socket_path = socket_path.to_path_buf();
211        let stream =
212            spawn_blocking_io(move || uds_windows::UnixStream::connect(socket_path)).await?;
213        Async::new(WindowsUnixStream::from(stream))
214            .map(FuturesAsyncReadCompatExt::compat)
215            .map(Stream)
216    }
217
218    pub(super) async fn is_stale_socket_path(socket_path: &Path) -> IoResult<bool> {
219        tokio::fs::try_exists(socket_path).await
220    }
221
222    async fn spawn_blocking_io<T>(
223        operation: impl FnOnce() -> IoResult<T> + Send + 'static,
224    ) -> IoResult<T>
225    where
226        T: Send + 'static,
227    {
228        task::spawn_blocking(operation)
229            .await
230            .map_err(|err| io::Error::other(format!("blocking socket task failed: {err}")))?
231    }
232
233    pub(super) struct WindowsUnixListener(uds_windows::UnixListener);
234
235    impl From<uds_windows::UnixListener> for WindowsUnixListener {
236        fn from(listener: uds_windows::UnixListener) -> Self {
237            Self(listener)
238        }
239    }
240
241    impl Deref for WindowsUnixListener {
242        type Target = uds_windows::UnixListener;
243
244        fn deref(&self) -> &Self::Target {
245            &self.0
246        }
247    }
248
249    impl AsSocket for WindowsUnixListener {
250        fn as_socket(&self) -> BorrowedSocket<'_> {
251            unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) }
252        }
253    }
254
255    pub(super) struct WindowsUnixStream(uds_windows::UnixStream);
256
257    impl From<uds_windows::UnixStream> for WindowsUnixStream {
258        fn from(stream: uds_windows::UnixStream) -> Self {
259            Self(stream)
260        }
261    }
262
263    impl Deref for WindowsUnixStream {
264        type Target = uds_windows::UnixStream;
265
266        fn deref(&self) -> &Self::Target {
267            &self.0
268        }
269    }
270
271    impl AsSocket for WindowsUnixStream {
272        fn as_socket(&self) -> BorrowedSocket<'_> {
273            unsafe { BorrowedSocket::borrow_raw(self.as_raw_socket()) }
274        }
275    }
276
277    impl io::Read for WindowsUnixStream {
278        fn read(&mut self, buf: &mut [u8]) -> IoResult<usize> {
279            io::Read::read(&mut self.0, buf)
280        }
281    }
282
283    impl io::Write for WindowsUnixStream {
284        fn write(&mut self, buf: &[u8]) -> IoResult<usize> {
285            io::Write::write(&mut self.0, buf)
286        }
287
288        fn flush(&mut self) -> IoResult<()> {
289            io::Write::flush(&mut self.0)
290        }
291    }
292
293    impl AsyncRead for Stream {
294        fn poll_read(
295            self: Pin<&mut Self>,
296            cx: &mut Context<'_>,
297            buf: &mut ReadBuf<'_>,
298        ) -> Poll<IoResult<()>> {
299            Pin::new(&mut self.get_mut().0).poll_read(cx, buf)
300        }
301    }
302
303    impl AsyncWrite for Stream {
304        fn poll_write(
305            self: Pin<&mut Self>,
306            cx: &mut Context<'_>,
307            buf: &[u8],
308        ) -> Poll<IoResult<usize>> {
309            Pin::new(&mut self.get_mut().0).poll_write(cx, buf)
310        }
311
312        fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
313            Pin::new(&mut self.get_mut().0).poll_flush(cx)
314        }
315
316        fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<IoResult<()>> {
317            let stream = &mut self.get_mut().0;
318            ready!(Pin::new(&mut *stream).poll_flush(cx))?;
319            // `Compat<Async<_>>` maps shutdown to `poll_close()`, which only
320            // flushes for `async_io::Async`; call the socket shutdown directly.
321            stream.get_ref().get_ref().shutdown(Shutdown::Write)?;
322            Poll::Ready(Ok(()))
323        }
324    }
325
326    unsafe impl async_io::IoSafe for WindowsUnixListener {}
327    unsafe impl async_io::IoSafe for WindowsUnixStream {}
328}
329
330#[cfg(test)]
331mod lib_tests;