1 use crate::conn::Connection;
2 use crate::conn::ConnectionExt;
3 use std::io;
4 use std::os::unix::net::UnixStream;
5 
6 // TODO: Remove PeekExt once rust-lang/rust#73761 is stabilized
7 trait PeekExt {
peek(&self, buf: &mut [u8]) -> io::Result<usize>8     fn peek(&self, buf: &mut [u8]) -> io::Result<usize>;
9 }
10 
11 impl PeekExt for UnixStream {
12     #[cfg(feature = "paranoid_unsafe")]
13     #[allow(clippy::panic)]
peek(&self, _buf: &mut [u8]) -> io::Result<usize>14     fn peek(&self, _buf: &mut [u8]) -> io::Result<usize> {
15         panic!("cannot use `UnixStream::peek` with `paranoid_unsafe` until rust-lang/rust#73761 is stabilized");
16     }
17 
18     #[cfg(not(feature = "paranoid_unsafe"))]
19     #[allow(non_camel_case_types)]
peek(&self, buf: &mut [u8]) -> io::Result<usize>20     fn peek(&self, buf: &mut [u8]) -> io::Result<usize> {
21         use core::ffi::c_void;
22         use std::os::unix::io::AsRawFd;
23 
24         // Define some libc types inline (to avoid bringing in entire libc dep)
25 
26         // every platform supported by the libc crate uses c_int = i32
27         type c_int = i32;
28         type size_t = usize;
29         type ssize_t = isize;
30         const MSG_PEEK: c_int = 2;
31         extern "C" {
32             fn recv(socket: c_int, buf: *mut c_void, len: size_t, flags: c_int) -> ssize_t;
33         }
34 
35         // from std/sys/unix/mod.rs
36         pub fn cvt(t: isize) -> io::Result<isize> {
37             if t == -1 {
38                 Err(io::Error::last_os_error())
39             } else {
40                 Ok(t)
41             }
42         }
43 
44         // from std/sys/unix/net.rs
45         let ret = cvt(unsafe {
46             recv(
47                 self.as_raw_fd(),
48                 buf.as_mut_ptr() as *mut c_void,
49                 buf.len(),
50                 MSG_PEEK,
51             )
52         })?;
53         Ok(ret as usize)
54     }
55 }
56 
57 impl Connection for UnixStream {
58     type Error = std::io::Error;
59 
write(&mut self, byte: u8) -> Result<(), Self::Error>60     fn write(&mut self, byte: u8) -> Result<(), Self::Error> {
61         use std::io::Write;
62 
63         Write::write_all(self, &[byte])
64     }
65 
write_all(&mut self, buf: &[u8]) -> Result<(), Self::Error>66     fn write_all(&mut self, buf: &[u8]) -> Result<(), Self::Error> {
67         use std::io::Write;
68 
69         Write::write_all(self, buf)
70     }
71 
flush(&mut self) -> Result<(), Self::Error>72     fn flush(&mut self) -> Result<(), Self::Error> {
73         use std::io::Write;
74 
75         Write::flush(self)
76     }
77 }
78 
79 impl ConnectionExt for UnixStream {
read(&mut self) -> Result<u8, Self::Error>80     fn read(&mut self) -> Result<u8, Self::Error> {
81         use std::io::Read;
82 
83         self.set_nonblocking(false)?;
84 
85         let mut buf = [0u8];
86         match Read::read_exact(self, &mut buf) {
87             Ok(_) => Ok(buf[0]),
88             Err(e) => Err(e),
89         }
90     }
91 
peek(&mut self) -> Result<Option<u8>, Self::Error>92     fn peek(&mut self) -> Result<Option<u8>, Self::Error> {
93         self.set_nonblocking(true)?;
94 
95         let mut buf = [0u8];
96         match PeekExt::peek(self, &mut buf) {
97             Ok(_) => Ok(Some(buf[0])),
98             Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => Ok(None),
99             Err(e) => Err(e),
100         }
101     }
102 }
103