diff --git a/src/data.rs b/src/data.rs index 2bfe7bc..cee7423 100644 --- a/src/data.rs +++ b/src/data.rs @@ -1,5 +1,7 @@ use std::io::{Read, Write}; +use tokio::io::{AsyncRead, AsyncReadExt}; + use crate::inet::InetError; @@ -13,23 +15,67 @@ pub enum DataError { Inet(InetError), } +pub trait ReadLike { + type Output; + fn _read_bytes(self, size: usize) -> Self::Output; +} +impl ReadLike for &mut R { + type Output = Result, DataError>; + + fn _read_bytes(self, size: usize) -> Self::Output { + let mut buf = vec![0; size]; + let mut read = 0; + while read < size { + match self.read(&mut buf[read..]) { + Ok(0) => return Err(DataError::Inet(InetError::ConnectionClosed)), + Ok(n) => read+=n, + Err(_) => return Err(DataError::ReadError) + } + }; Ok(buf) + } +} + +#[cfg(feature = "async")] +#[async_trait] +impl ReadLike for &mut R { + type Output = Result, DataError>; + + async fn _read_bytes(self, size: usize) -> Self::Output { + let mut buf = vec![0; size]; + let mut read = 0; + while read < size { + match AsyncReadExt::read(&mut self, &mut buf[read..]).await { + Ok(0) => return Err(DataError::Inet(InetError::ConnectionClosed)), + Ok(n) => read += n, + Err(_) => return Err(DataError::ReadError), + } + } + Ok(buf) + } +} + +#[cfg_attr(feature = "async", async_trait)] pub trait DataReader { + #[cfg(not(feature = "async"))] fn read_bytes(&mut self, size: usize) -> Result, DataError>; + #[cfg(feature = "async")] + async fn read_bytes(&mut self, size: usize) -> Result, DataError>; + fn read_byte(&mut self) -> Result { Ok(self.read_bytes(1)?[0]) } - fn read_byte_signed(&mut self) -> Result { + fn read_signed_byte(&mut self) -> Result { Ok(self.read_byte()? as i8) } fn read_short(&mut self) -> Result { - Ok((self.read_byte()? as u16) + ((self.read_byte()? as u16) << 8)) + Ok(u16::from_be_bytes(self.read_bytes(2)?.try_into().unwrap())) } - fn read_short_signed(&mut self) -> Result { + fn read_signed_short(&mut self) -> Result { Ok(self.read_short()? as i16) } @@ -53,17 +99,17 @@ pub trait DataReader { } impl DataReader for R { + #[cfg(not(feature = "async"))] fn read_bytes(&mut self, size: usize) -> Result, DataError> { - let mut buf = vec![0; size]; - match self.read(&mut buf) { - Err(_) => return Err(DataError::ReadError), - Ok(n) => if n == 0 { - return Err(DataError::Inet(InetError::ConnectionClosed)); - } else if n < size { - buf.truncate(n); - buf.append(&mut self.read_bytes(size - n)?); - } - }; Ok(buf) + self._read_bytes(size) + } +} + +#[cfg(feature = "async")] +#[async_trait] +impl DataReader for R { + async fn _read_bytes(&mut self, size: usize) -> Result, DataError> { + self.read_bytes(size).await } } diff --git a/src/main.rs b/src/main.rs index ec98674..b0baffa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,8 @@ -use std::{sync::Arc, time::Duration}; +use std::{io::Cursor, sync::Arc, time::Duration}; use data::{DataError, DataReader}; use inet::Server; -use tokio::{net::TcpStream, time::sleep}; +use tokio::{io::AsyncReadExt, net::TcpStream, time::sleep}; @@ -19,17 +19,14 @@ async fn main() { loop { let stream = server.accept().await; - tokio::spawn(test(Arc::new(stream))); + tokio::spawn(test(stream)); } } -async fn test(stream: Arc) { +async fn test(stream: TcpStream) { let Ok(addr) = stream.peer_addr() else {return;}; println!("Подключение: {addr}"); - read_first_packet(stream.clone()); + // читаем первый пакет + let size = stream.read_varint(); println!("Отключение: {addr}"); -} - -fn read_first_packet(stream: Arc) -> Result<(), DataError>{ - let size = stream.read_varint()?; } \ No newline at end of file