1use std::{
2 future::Future,
3 io,
4 path::Path,
5 pin::Pin,
6 task::{Context, Poll},
7};
8
9use compio_buf::{BufResult, IoBuf, IoBufMut, IoVectoredBuf, IoVectoredBufMut};
10use compio_driver::{
11 BufferRef, SharedFd, impl_raw_fd,
12 op::{RecvFlags, RecvMsgMultiResult, SendMsgZc, SendVectoredZc, SendZc},
13};
14use compio_io::{
15 AsyncRead, AsyncReadManaged, AsyncReadMulti, AsyncWrite, AsyncWriteZerocopy,
16 ancillary::{
17 AsyncReadAncillary, AsyncReadAncillaryManaged, AsyncReadAncillaryMulti,
18 AsyncWriteAncillary, AsyncWriteAncillaryZerocopy, ReturnFlags,
19 },
20 util::Splittable,
21};
22#[cfg(unix)]
23use compio_runtime::Runtime;
24use compio_runtime::fd::PollFd;
25use futures_util::{Stream, StreamExt, stream::FusedStream};
26#[cfg(unix)]
27use socket2::Socket as Socket2;
28use socket2::{Domain, SockAddr, Type};
29
30use crate::{
31 Extract, Incoming, MSG_NOSIGNAL, ReadHalf, Socket, SocketWithState, WriteHalf, Zerocopy,
32};
33
34#[derive(Debug, Clone)]
63pub struct UnixListener {
64 inner: Socket,
65}
66
67impl UnixListener {
68 pub async fn bind(path: impl AsRef<Path>) -> io::Result<Self> {
71 Self::bind_addr(&SockAddr::unix(path)?).await
72 }
73
74 pub async fn bind_addr(addr: &SockAddr) -> io::Result<Self> {
80 if !addr.is_unix() {
81 return Err(io::Error::new(
82 io::ErrorKind::InvalidInput,
83 "addr is not unix socket address",
84 ));
85 }
86
87 let socket = Socket::new(addr.domain(), Type::STREAM, None).await?;
88 socket.bind(addr).await?;
89 socket.listen(1024).await?;
90 Ok(UnixListener { inner: socket })
91 }
92
93 #[cfg(unix)]
94 pub fn from_std(stream: std::os::unix::net::UnixListener) -> io::Result<Self> {
96 if Runtime::with_current(|r| r.driver_type().is_polling()) {
97 stream.set_nonblocking(true)?;
98 }
99
100 Ok(Self {
101 inner: Socket::from_socket2(Socket2::from(stream))?,
102 })
103 }
104
105 pub fn close(self) -> impl Future<Output = io::Result<()>> {
112 self.inner.close()
113 }
114
115 pub async fn accept(&self) -> io::Result<(UnixStream, SockAddr)> {
121 let (socket, addr) = self.inner.accept().await?;
122 let stream = UnixStream { inner: socket };
123 Ok((stream, addr))
124 }
125
126 #[cfg(windows)]
129 pub async fn accept_with(&self, sock: UnixSocket) -> io::Result<(UnixStream, SockAddr)> {
130 let (socket, addr) = self.inner.accept_with(sock.inner).await?;
131 let stream = UnixStream { inner: socket };
132 Ok((stream, addr))
133 }
134
135 pub fn incoming(&self) -> UnixIncoming<'_> {
137 UnixIncoming {
138 inner: self.inner.incoming(),
139 }
140 }
141
142 pub fn local_addr(&self) -> io::Result<SockAddr> {
144 self.inner.local_addr()
145 }
146
147 pub fn take_error(&self) -> io::Result<Option<io::Error>> {
149 self.inner.socket.take_error()
150 }
151
152 pub fn to_poll_fd(&self) -> io::Result<PollFd<SocketWithState>> {
154 self.inner.to_poll_fd()
155 }
156
157 pub fn into_poll_fd(self) -> io::Result<PollFd<SocketWithState>> {
159 self.inner.into_poll_fd()
160 }
161}
162
163impl_raw_fd!(UnixListener, SocketWithState, inner, socket);
164
165pub struct UnixIncoming<'a> {
167 inner: Incoming<'a>,
168}
169
170impl Stream for UnixIncoming<'_> {
171 type Item = io::Result<UnixStream>;
172
173 fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
174 let this = self.get_mut();
175 this.inner.poll_next_unpin(cx).map(|res| {
176 res.map(|res| {
177 let socket = res?;
178 Ok(UnixStream { inner: socket })
179 })
180 })
181 }
182}
183
184impl FusedStream for UnixIncoming<'_> {
185 fn is_terminated(&self) -> bool {
186 self.inner.is_terminated()
187 }
188}
189
190#[derive(Debug, Clone)]
210pub struct UnixStream {
211 inner: Socket,
212}
213
214impl UnixStream {
215 pub async fn connect(path: impl AsRef<Path>) -> io::Result<Self> {
218 Self::connect_addr(&SockAddr::unix(path)?).await
219 }
220
221 pub async fn connect_addr(addr: &SockAddr) -> io::Result<Self> {
228 if !addr.is_unix() {
229 return Err(io::Error::new(
230 io::ErrorKind::InvalidInput,
231 "addr is not unix socket address",
232 ));
233 }
234 let socket = Socket::new(Domain::UNIX, Type::STREAM, None).await?;
235 #[cfg(windows)]
236 {
237 let new_addr = empty_unix_socket();
238 socket.bind(&new_addr).await?
239 }
240 socket.connect_async(addr).await?;
241 let unix_stream = UnixStream { inner: socket };
242 Ok(unix_stream)
243 }
244
245 #[cfg(unix)]
247 pub fn from_std(stream: std::os::unix::net::UnixStream) -> io::Result<Self> {
248 if Runtime::with_current(|r| r.driver_type().is_polling()) {
249 stream.set_nonblocking(true)?;
250 }
251
252 Ok(Self {
253 inner: Socket::from_socket2(Socket2::from(stream))?,
254 })
255 }
256
257 pub fn close(self) -> impl Future<Output = io::Result<()>> {
264 self.inner.close()
265 }
266
267 pub fn peer_addr(&self) -> io::Result<SockAddr> {
269 #[allow(unused_mut)]
270 let mut addr = self.inner.peer_addr()?;
271 #[cfg(windows)]
272 {
273 fix_unix_socket_length(&mut addr);
274 }
275 Ok(addr)
276 }
277
278 pub fn local_addr(&self) -> io::Result<SockAddr> {
280 self.inner.local_addr()
281 }
282
283 pub fn take_error(&self) -> io::Result<Option<io::Error>> {
285 self.inner.socket.take_error()
286 }
287
288 pub fn split(&self) -> (ReadHalf<'_, Self>, WriteHalf<'_, Self>) {
295 crate::split(self)
296 }
297
298 pub fn into_split(self) -> (Self, Self) {
304 (self.clone(), self)
305 }
306
307 pub fn to_poll_fd(&self) -> io::Result<PollFd<SocketWithState>> {
309 self.inner.to_poll_fd()
310 }
311
312 pub fn into_poll_fd(self) -> io::Result<PollFd<SocketWithState>> {
314 self.inner.into_poll_fd()
315 }
316
317 #[cfg(windows)]
322 pub async fn disconnect(self) -> io::Result<UnixSocket> {
323 self.inner.disconnect().await?;
324 Ok(UnixSocket { inner: self.inner })
325 }
326}
327
328impl AsyncRead for UnixStream {
329 #[inline]
330 async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
331 (&*self).read(buf).await
332 }
333
334 #[inline]
335 async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
336 (&*self).read_vectored(buf).await
337 }
338}
339
340impl AsyncRead for &UnixStream {
341 #[inline]
342 async fn read<B: IoBufMut>(&mut self, buf: B) -> BufResult<usize, B> {
343 self.inner.recv(buf, RecvFlags::empty()).await
344 }
345
346 #[inline]
347 async fn read_vectored<V: IoVectoredBufMut>(&mut self, buf: V) -> BufResult<usize, V> {
348 self.inner.recv_vectored(buf, RecvFlags::empty()).await
349 }
350}
351
352impl AsyncReadManaged for UnixStream {
353 type Buffer = BufferRef;
354
355 async fn read_managed(&mut self, len: usize) -> io::Result<Option<Self::Buffer>> {
356 (&*self).read_managed(len).await
357 }
358}
359
360impl AsyncReadManaged for &UnixStream {
361 type Buffer = BufferRef;
362
363 async fn read_managed(&mut self, len: usize) -> io::Result<Option<Self::Buffer>> {
364 self.inner.recv_managed(len, RecvFlags::empty()).await
365 }
366}
367
368impl AsyncReadMulti for UnixStream {
369 fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
370 self.inner.recv_multi(len, RecvFlags::empty())
371 }
372}
373
374impl AsyncReadMulti for &UnixStream {
375 fn read_multi(&mut self, len: usize) -> impl Stream<Item = io::Result<Self::Buffer>> {
376 self.inner.recv_multi(len, RecvFlags::empty())
377 }
378}
379
380impl AsyncReadAncillary for UnixStream {
381 #[inline]
382 async fn read_with_ancillary<T: IoBufMut, C: IoBufMut>(
383 &mut self,
384 buffer: T,
385 control: C,
386 ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
387 (&*self).read_with_ancillary(buffer, control).await
388 }
389
390 #[inline]
391 async fn read_vectored_with_ancillary<T: IoVectoredBufMut, C: IoBufMut>(
392 &mut self,
393 buffer: T,
394 control: C,
395 ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
396 (&*self).read_vectored_with_ancillary(buffer, control).await
397 }
398}
399
400impl AsyncReadAncillary for &UnixStream {
401 #[inline]
402 async fn read_with_ancillary<T: IoBufMut, C: IoBufMut>(
403 &mut self,
404 buffer: T,
405 control: C,
406 ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
407 self.inner
408 .recv_msg(buffer, control, RecvFlags::empty())
409 .await
410 .map_res(|(res, len, _addr, flags)| (res, len, flags))
411 }
412
413 #[inline]
414 async fn read_vectored_with_ancillary<T: IoVectoredBufMut, C: IoBufMut>(
415 &mut self,
416 buffer: T,
417 control: C,
418 ) -> BufResult<(usize, usize, ReturnFlags), (T, C)> {
419 self.inner
420 .recv_msg_vectored(buffer, control, RecvFlags::empty())
421 .await
422 .map_res(|(res, len, _addr, flags)| (res, len, flags))
423 }
424}
425
426impl AsyncReadAncillaryManaged for UnixStream {
427 #[inline]
428 async fn read_managed_with_ancillary<C: IoBufMut>(
429 &mut self,
430 len: usize,
431 control: C,
432 ) -> io::Result<Option<(Self::Buffer, C, ReturnFlags)>> {
433 (&*self).read_managed_with_ancillary(len, control).await
434 }
435}
436
437impl AsyncReadAncillaryManaged for &UnixStream {
438 #[inline]
439 async fn read_managed_with_ancillary<C: IoBufMut>(
440 &mut self,
441 len: usize,
442 control: C,
443 ) -> io::Result<Option<(Self::Buffer, C, ReturnFlags)>> {
444 self.inner
445 .recv_msg_managed(len, control, RecvFlags::empty())
446 .await
447 .map(|res| res.map(|(res, len, _addr, flags)| (res, len, flags)))
448 }
449}
450
451impl AsyncReadAncillaryMulti for UnixStream {
452 type Return = RecvMsgMultiResult;
453
454 #[inline]
455 fn read_multi_with_ancillary(
456 &mut self,
457 control_len: usize,
458 ) -> impl Stream<Item = io::Result<Self::Return>> {
459 self.inner.recv_msg_multi(control_len, RecvFlags::empty())
460 }
461}
462
463impl AsyncReadAncillaryMulti for &UnixStream {
464 type Return = RecvMsgMultiResult;
465
466 #[inline]
467 fn read_multi_with_ancillary(
468 &mut self,
469 control_len: usize,
470 ) -> impl Stream<Item = io::Result<Self::Return>> {
471 self.inner.recv_msg_multi(control_len, RecvFlags::empty())
472 }
473}
474
475impl AsyncWrite for UnixStream {
476 #[inline]
477 async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
478 (&*self).write(buf).await
479 }
480
481 #[inline]
482 async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
483 (&*self).write_vectored(buf).await
484 }
485
486 #[inline]
487 async fn flush(&mut self) -> io::Result<()> {
488 (&*self).flush().await
489 }
490
491 #[inline]
492 async fn shutdown(&mut self) -> io::Result<()> {
493 (&*self).shutdown().await
494 }
495}
496
497impl AsyncWrite for &UnixStream {
498 #[inline]
499 async fn write<T: IoBuf>(&mut self, buf: T) -> BufResult<usize, T> {
500 self.inner.send(buf, MSG_NOSIGNAL).await
501 }
502
503 #[inline]
504 async fn write_vectored<T: IoVectoredBuf>(&mut self, buf: T) -> BufResult<usize, T> {
505 self.inner.send_vectored(buf, MSG_NOSIGNAL).await
506 }
507
508 #[inline]
509 async fn flush(&mut self) -> io::Result<()> {
510 Ok(())
511 }
512
513 #[inline]
514 async fn shutdown(&mut self) -> io::Result<()> {
515 self.inner.shutdown().await
516 }
517}
518
519impl AsyncWriteZerocopy for UnixStream {
520 type BufferReadyFuture<T: IoBuf> = Zerocopy<SendZc<T, SharedFd<SocketWithState>>>;
521 type VectoredBufferReadyFuture<T: IoVectoredBuf> =
522 Zerocopy<SendVectoredZc<T, SharedFd<SocketWithState>>>;
523
524 async fn write_zerocopy<T: IoBuf>(
525 &mut self,
526 buf: T,
527 ) -> BufResult<usize, Self::BufferReadyFuture<T>> {
528 self.inner.send_zerocopy(buf, MSG_NOSIGNAL).await
529 }
530
531 async fn write_zerocopy_vectored<T: IoVectoredBuf>(
532 &mut self,
533 buf: T,
534 ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T>> {
535 self.inner.send_zerocopy_vectored(buf, MSG_NOSIGNAL).await
536 }
537}
538
539impl AsyncWriteZerocopy for &UnixStream {
540 type BufferReadyFuture<T: IoBuf> = Zerocopy<SendZc<T, SharedFd<SocketWithState>>>;
541 type VectoredBufferReadyFuture<T: IoVectoredBuf> =
542 Zerocopy<SendVectoredZc<T, SharedFd<SocketWithState>>>;
543
544 async fn write_zerocopy<T: IoBuf>(
545 &mut self,
546 buf: T,
547 ) -> BufResult<usize, Self::BufferReadyFuture<T>> {
548 self.inner.send_zerocopy(buf, MSG_NOSIGNAL).await
549 }
550
551 async fn write_zerocopy_vectored<T: IoVectoredBuf>(
552 &mut self,
553 buf: T,
554 ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T>> {
555 self.inner.send_zerocopy_vectored(buf, MSG_NOSIGNAL).await
556 }
557}
558
559impl AsyncWriteAncillaryZerocopy for UnixStream {
560 type BufferReadyFuture<T: IoBuf, C: IoBuf> =
561 Extract<Zerocopy<SendMsgZc<[T; 1], C, SharedFd<SocketWithState>>>, T, C>;
562 type VectoredBufferReadyFuture<T: IoVectoredBuf, C: IoBuf> =
563 Zerocopy<SendMsgZc<T, C, SharedFd<SocketWithState>>>;
564
565 async fn write_zerocopy_with_ancillary<T: IoBuf, C: IoBuf>(
566 &mut self,
567 buf: T,
568 control: C,
569 ) -> BufResult<usize, Self::BufferReadyFuture<T, C>> {
570 self.inner
571 .send_msg_zerocopy(buf, control, None, MSG_NOSIGNAL)
572 .await
573 }
574
575 async fn write_zerocopy_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
576 &mut self,
577 buf: T,
578 control: C,
579 ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T, C>> {
580 self.inner
581 .send_msg_zerocopy_vectored(buf, control, None, MSG_NOSIGNAL)
582 .await
583 }
584}
585
586impl AsyncWriteAncillaryZerocopy for &UnixStream {
587 type BufferReadyFuture<T: IoBuf, C: IoBuf> =
588 Extract<Zerocopy<SendMsgZc<[T; 1], C, SharedFd<SocketWithState>>>, T, C>;
589 type VectoredBufferReadyFuture<T: IoVectoredBuf, C: IoBuf> =
590 Zerocopy<SendMsgZc<T, C, SharedFd<SocketWithState>>>;
591
592 async fn write_zerocopy_with_ancillary<T: IoBuf, C: IoBuf>(
593 &mut self,
594 buf: T,
595 control: C,
596 ) -> BufResult<usize, Self::BufferReadyFuture<T, C>> {
597 self.inner
598 .send_msg_zerocopy(buf, control, None, MSG_NOSIGNAL)
599 .await
600 }
601
602 async fn write_zerocopy_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
603 &mut self,
604 buf: T,
605 control: C,
606 ) -> BufResult<usize, Self::VectoredBufferReadyFuture<T, C>> {
607 self.inner
608 .send_msg_zerocopy_vectored(buf, control, None, MSG_NOSIGNAL)
609 .await
610 }
611}
612
613impl AsyncWriteAncillary for UnixStream {
663 #[inline]
664 async fn write_with_ancillary<T: IoBuf, C: IoBuf>(
665 &mut self,
666 buffer: T,
667 control: C,
668 ) -> BufResult<usize, (T, C)> {
669 (&*self).write_with_ancillary(buffer, control).await
670 }
671
672 #[inline]
673 async fn write_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
674 &mut self,
675 buffer: T,
676 control: C,
677 ) -> BufResult<usize, (T, C)> {
678 (&*self)
679 .write_vectored_with_ancillary(buffer, control)
680 .await
681 }
682}
683
684impl AsyncWriteAncillary for &UnixStream {
685 #[inline]
686 async fn write_with_ancillary<T: IoBuf, C: IoBuf>(
687 &mut self,
688 buffer: T,
689 control: C,
690 ) -> BufResult<usize, (T, C)> {
691 self.inner
692 .send_msg(buffer, control, None, MSG_NOSIGNAL)
693 .await
694 }
695
696 #[inline]
697 async fn write_vectored_with_ancillary<T: IoVectoredBuf, C: IoBuf>(
698 &mut self,
699 buffer: T,
700 control: C,
701 ) -> BufResult<usize, (T, C)> {
702 self.inner
703 .send_msg_vectored(buffer, control, None, MSG_NOSIGNAL)
704 .await
705 }
706}
707
708impl Splittable for UnixStream {
709 type ReadHalf = Self;
710 type WriteHalf = Self;
711
712 fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
713 self.into_split()
714 }
715}
716
717impl<'a> Splittable for &'a UnixStream {
718 type ReadHalf = ReadHalf<'a, UnixStream>;
719 type WriteHalf = WriteHalf<'a, UnixStream>;
720
721 fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
722 crate::split(self)
723 }
724}
725
726impl<'a> Splittable for &'a mut UnixStream {
727 type ReadHalf = ReadHalf<'a, UnixStream>;
728 type WriteHalf = WriteHalf<'a, UnixStream>;
729
730 fn split(self) -> (Self::ReadHalf, Self::WriteHalf) {
731 crate::split(self)
732 }
733}
734
735impl_raw_fd!(UnixStream, SocketWithState, inner, socket);
736
737#[derive(Debug)]
740pub struct UnixSocket {
741 inner: Socket,
742}
743
744impl UnixSocket {
745 pub async fn new_stream() -> io::Result<UnixSocket> {
747 UnixSocket::new(socket2::Type::STREAM).await
748 }
749
750 async fn new(ty: socket2::Type) -> io::Result<UnixSocket> {
751 let inner = Socket::new(socket2::Domain::UNIX, ty, None).await?;
752 Ok(UnixSocket { inner })
753 }
754
755 pub fn local_addr(&self) -> io::Result<SockAddr> {
757 self.inner.local_addr()
758 }
759
760 pub fn take_error(&self) -> io::Result<Option<io::Error>> {
762 self.inner.socket.take_error()
763 }
764
765 pub async fn bind(&self, path: impl AsRef<Path>) -> io::Result<()> {
767 self.bind_addr(&SockAddr::unix(path)?).await
768 }
769
770 pub async fn bind_addr(&self, addr: &SockAddr) -> io::Result<()> {
772 if !addr.is_unix() {
773 return Err(io::Error::new(
774 io::ErrorKind::InvalidInput,
775 "addr is not unix socket address",
776 ));
777 }
778 self.inner.bind(addr).await
779 }
780
781 pub async fn listen(self, backlog: i32) -> io::Result<UnixListener> {
788 self.inner.listen(backlog).await?;
789 Ok(UnixListener { inner: self.inner })
790 }
791
792 pub async fn connect(self, path: impl AsRef<Path>) -> io::Result<UnixStream> {
797 self.connect_addr(&SockAddr::unix(path)?).await
798 }
799
800 pub async fn connect_addr(self, addr: &SockAddr) -> io::Result<UnixStream> {
810 if !addr.is_unix() {
811 return Err(io::Error::new(
812 io::ErrorKind::InvalidInput,
813 "addr is not unix socket address",
814 ));
815 }
816 self.inner.connect_async(addr).await?;
817 Ok(UnixStream { inner: self.inner })
818 }
819}
820
821impl_raw_fd!(UnixSocket, SocketWithState, inner, socket);
822
823#[cfg(windows)]
824#[inline]
825fn empty_unix_socket() -> SockAddr {
826 use windows_sys::Win32::Networking::WinSock::{AF_UNIX, SOCKADDR_UN};
827
828 unsafe {
830 SockAddr::try_init(|addr, len| {
831 let addr: *mut SOCKADDR_UN = addr.cast();
832 std::ptr::write(
833 addr,
834 SOCKADDR_UN {
835 sun_family: AF_UNIX,
836 sun_path: [0; 108],
837 },
838 );
839 std::ptr::write(len, 3);
840 Ok(())
841 })
842 }
843 .unwrap()
845 .1
846}
847
848#[cfg(windows)]
852#[inline]
853fn fix_unix_socket_length(addr: &mut SockAddr) {
854 use windows_sys::Win32::Networking::WinSock::SOCKADDR_UN;
855
856 let unix_addr: &SOCKADDR_UN = unsafe { &*addr.as_ptr().cast() };
858 let sun_path = unsafe {
859 std::slice::from_raw_parts(
860 unix_addr.sun_path.as_ptr() as *const u8,
861 unix_addr.sun_path.len(),
862 )
863 };
864 let addr_len = match std::ffi::CStr::from_bytes_until_nul(sun_path) {
865 Ok(str) => str.to_bytes_with_nul().len() + 2,
866 Err(_) => std::mem::size_of::<SOCKADDR_UN>(),
867 };
868 unsafe {
869 addr.set_length(addr_len as _);
870 }
871}