diff --git a/src/models/packet.rs b/src/models/packet.rs index 9df6b2a..6e9db31 100644 --- a/src/models/packet.rs +++ b/src/models/packet.rs @@ -295,16 +295,16 @@ impl<'a> OED<'a> { stream.send(&self.plain_data()?).await?; self.chunk_count = 0; - let mut remaining_data = &self.encrypted_data[..]; - while !remaining_data.is_empty() { - let chunk_size = remaining_data.len().min(1024); - - let chunk_length = chunk_size as u32; - - stream.send(&chunk_length.to_be_bytes()).await?; - stream.send(&remaining_data[..chunk_size]).await?; - - remaining_data = &remaining_data[chunk_size..]; + let mut chunks = self.encrypted_data.chunks(1024); + loop { + if let Some(chunk) = chunks.next() { + let chunk_size = chunk.len() as u32; + stream.send(&chunk_size.to_be_bytes()).await?; + stream.send(&chunk).await?; + self.chunk_count += 1; + } else { + break; + } } stream.send(&STOP_FLAG).await?; @@ -320,7 +320,11 @@ impl<'a> OED<'a> { Ok(plain_bytes) } - pub fn get_data(&self) -> Vec { - self.data.clone().unwrap() + pub fn take(&mut self) -> Vec { + self.data.take().unwrap() + } + + pub fn get_data(&self) -> &[u8] { + self.data.as_ref().unwrap() } } diff --git a/src/models/session.rs b/src/models/session.rs index 3e7937a..2ee28c3 100644 --- a/src/models/session.rs +++ b/src/models/session.rs @@ -193,7 +193,7 @@ impl Session { let content = OED::new(&self.aes_key) .from_stream(socket) .await? - .get_data(); + .take(); let status_code = OSC::from_stream(socket).await?.status_code; let response = Response::new(None, content, None, status_code, flag); diff --git a/src/utils/gear.rs b/src/utils/gear.rs index e488bc3..e8b6949 100644 --- a/src/utils/gear.rs +++ b/src/utils/gear.rs @@ -50,16 +50,17 @@ impl Socket { } } + #[inline] pub async fn peer_addr(&self) -> Result { Ok(self.writer.lock().await.peer_addr()?) } + #[inline] pub async fn recv_usize(&self) -> Result { let mut len_bytes = [0; 4]; #[cfg(not(feature = "perf"))] self.reader.lock().await.read_exact(&mut len_bytes).await?; #[cfg(feature = "perf")] - #[cfg(feature = "perf")] { use colored::Colorize; let now = tokio::time::Instant::now(); @@ -73,24 +74,28 @@ impl Socket { Ok(u32::from_be_bytes(len_bytes) as usize) } + #[inline] pub async fn recv_u32(&self) -> Result { let mut len_bytes = [0; 4]; self.reader.lock().await.read_exact(&mut len_bytes).await?; Ok(u32::from_be_bytes(len_bytes)) } + #[inline] pub async fn recv(&self, len: usize) -> Result> { let mut recv_bytes: Vec = vec![0; len]; self.reader.lock().await.read_exact(&mut recv_bytes).await?; Ok(recv_bytes) } + #[inline] pub async fn recv_str(&self, len: usize) -> Result { let mut recv_bytes: Vec = vec![0; len]; self.reader.lock().await.read_exact(&mut recv_bytes).await?; Ok(String::from_utf8(recv_bytes)?) } + #[inline] pub async fn send(&self, data: &[u8]) -> Result<()> { let mut writer = self.writer.lock().await; writer.write_all(&data).await?;