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