1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175
use std::io::{self, Read, Write}; use std::marker::PhantomData; use std::net::{Shutdown, SocketAddr}; use std::sync::{Mutex, MutexGuard}; use bincode::Result; use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt}; use serde::{de::DeserializeOwned, Serialize}; use socket2::{Domain, Socket, Type}; pub struct SocketClient<Req, Res> where Req: Serialize, Res: DeserializeOwned, { buffer: Mutex<Buffer>, stream: Socket, _request: PhantomData<Req>, _response: PhantomData<Res>, } impl<Req, Res> SocketClient<Req, Res> where Req: Serialize, Res: DeserializeOwned, { pub fn try_new(addr: SocketAddr) -> io::Result<Self> { let domain = match addr { SocketAddr::V4(_) => Domain::ipv4(), SocketAddr::V6(_) => Domain::ipv6(), }; let socket = Socket::new(domain, Type::stream(), None)?; socket.connect(&addr.into())?; Ok(Self { buffer: Mutex::default(), stream: socket, _request: PhantomData::default(), _response: PhantomData::default(), }) } pub(crate) fn try_from_stream(stream: Socket) -> io::Result<Self> { stream.set_nonblocking(true)?; Ok(Self { buffer: Mutex::default(), stream, _request: PhantomData::default(), _response: PhantomData::default(), }) } } impl<Req, Res> SocketClient<Req, Res> where Req: Serialize, Res: DeserializeOwned, { pub fn request(&self, request: &Req) -> Result<Res> { let mut buffer = self.buffer.lock().unwrap(); let stream = &mut &self.stream; buffer.data.clear(); bincode::serialize_into(&mut buffer.data, request)?; let size = buffer.data.len() as u64; assert_ne!(size, 0, "Message must have one or more bytes."); stream.write_u64::<NetworkEndian>(size)?; stream.write_all(&buffer.data)?; bincode::deserialize_from(stream) } pub(crate) fn response<F>(&self, handler: F) -> Result<SocketStatus> where F: FnMut(Res) -> Req, { let mut buffer = self.buffer.lock().unwrap(); let stream = &mut &self.stream; if buffer.size.is_some() { fill_buffer_and_handle(buffer, stream, handler) } else { let mut buf = [0; 8]; match stream.peek(&mut buf) { Ok(8) => { stream.read_exact(&mut buf)?; let size = buf.as_ref().read_u64::<NetworkEndian>()? as usize; if size == 0 { return Ok(SocketStatus::Closed); } buffer.offset = 0; buffer.size = Some(size); buffer.data.resize(size, 0); fill_buffer_and_handle(buffer, stream, handler) } Ok(_) => Ok(SocketStatus::Alive), Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => Ok(SocketStatus::Alive), Err(error) => Err(error.into()), } } } fn stop(&self) -> io::Result<()> { let stream = &mut &self.stream; stream.write_u64::<NetworkEndian>(0)?; stream.shutdown(Shutdown::Read)?; Ok(()) } } fn fill_buffer_and_handle<Res, Req, F>( mut buffer: MutexGuard<Buffer>, stream: &mut &Socket, mut handler: F, ) -> Result<SocketStatus> where Req: Serialize, Res: DeserializeOwned, F: FnMut(Res) -> Req, { let offset = buffer.offset; buffer.offset += match stream.read(&mut buffer.data[offset..]) { Ok(size) => size, Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => return Ok(SocketStatus::Alive), Err(error) => return Err(error.into()), }; if buffer.offset >= buffer.size.unwrap() { let request = bincode::deserialize_from(&buffer.data[..])?; buffer.size = None; bincode::serialize_into(stream, &handler(request))?; Ok(SocketStatus::Alive) } else { Ok(SocketStatus::Alive) } } impl<Req, Res> Drop for SocketClient<Req, Res> where Req: Serialize, Res: DeserializeOwned, { fn drop(&mut self) { match self.stop() { Ok(()) => (), Err(ref e) if e.kind() == io::ErrorKind::NotConnected => (), Err(error) => panic!(error), } } } pub enum SocketStatus { Alive, Closed, } #[derive(Default)] struct Buffer { data: Vec<u8>, offset: usize, size: Option<usize>, }