From 2dec0ce5c81d96426102a6cffbc8253a62347ec9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8B=8F=E5=90=91=E5=A4=9C?= Date: Mon, 15 Apr 2024 15:44:44 +0800 Subject: [PATCH] refactor: improve exception pass in packet.rs --- Cargo.lock | 12 +++++----- src/models/packet.rs | 57 ++++++++++++++++++-------------------------- 2 files changed, 29 insertions(+), 40 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dcf015e..a2cf3fd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -109,9 +109,9 @@ checksum = "514de17de45fdb8dc022b1a7975556c53c86f9f0aa5f534b98977b171857c2c9" [[package]] name = "cc" -version = "1.0.92" +version = "1.0.94" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2678b2e3449475e95b0aa6f9b506a28e61b3dc8996592b983695e8ebb58a8b41" +checksum = "17f6e324229dc011159fcc089755d1e2e216a90d43a7dea6853ca740b84f35e7" [[package]] name = "cfg-if" @@ -715,9 +715,9 @@ dependencies = [ [[package]] name = "proc-macro2" -version = "1.0.79" +version = "1.0.80" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e835ff2298f5721608eb1a980ecaee1aef2c132bf95ecc026a11b7bf3c01c02e" +checksum = "a56dea16b0a29e94408b9aa5e2940a4eedbd128a1ba20e8f7ae60fd3d465af0e" dependencies = [ "unicode-ident", ] @@ -1050,9 +1050,9 @@ checksum = "81cdd64d312baedb58e21336b31bc043b77e01cc99033ce76ef539f78e965ebc" [[package]] name = "syn" -version = "2.0.58" +version = "2.0.59" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "44cfb93f38070beee36b3fef7d4f5a16f27751d94b187b666a5cc5e9b0d30687" +checksum = "4a6531ffc7b071655e4ce2e04bd464c4830bb585a61cabb96cf808f05172615a" dependencies = [ "proc-macro2", "quote", diff --git a/src/models/packet.rs b/src/models/packet.rs index 60554b7..480d096 100644 --- a/src/models/packet.rs +++ b/src/models/packet.rs @@ -92,22 +92,15 @@ impl<'a> OKE<'a> { }) } - pub fn from_public_key_bytes( - &mut self, - public_key_bytes: &[u8], - ) -> Result<&mut Self, OblivionException> { - self.public_key = Some(PublicKey::from_sec1_bytes(public_key_bytes).unwrap()); + pub fn from_public_key_bytes(&mut self, public_key_bytes: &[u8]) -> Result<&mut Self> { + self.public_key = Some(PublicKey::from_sec1_bytes(public_key_bytes)?); Ok(self) } - pub async fn from_stream( - &mut self, - stream: &mut Socket, - ) -> Result<&mut Self, OblivionException> { + pub async fn from_stream(&mut self, stream: &mut Socket) -> Result<&mut Self> { let remote_public_key_length = stream.recv_len().await?; let remote_public_key_bytes = stream.recv(remote_public_key_length).await?; - self.remote_public_key = - Some(PublicKey::from_sec1_bytes(&remote_public_key_bytes).unwrap()); + self.remote_public_key = Some(PublicKey::from_sec1_bytes(&remote_public_key_bytes)?); self.shared_aes_key = Some(generate_shared_key( self.private_key.as_ref().unwrap(), self.remote_public_key.as_ref().unwrap(), @@ -122,44 +115,40 @@ impl<'a> OKE<'a> { ) -> Result<&mut Self, OblivionException> { let remote_public_key_length = stream.recv_len().await?; let remote_public_key_bytes = stream.recv(remote_public_key_length).await?; - self.remote_public_key = - Some(PublicKey::from_sec1_bytes(&remote_public_key_bytes).unwrap()); + self.remote_public_key = Some(PublicKey::from_sec1_bytes(&remote_public_key_bytes)?); let salt_length = stream.recv_len().await?; self.salt = Some(stream.recv(salt_length).await?); self.shared_aes_key = Some(generate_shared_key( - self.private_key.as_ref().unwrap(), - self.remote_public_key.as_ref().unwrap(), - &self.salt.as_mut().unwrap(), + self.private_key.unwrap(), + &self.remote_public_key.unwrap(), + self.salt.as_mut().unwrap(), )?); Ok(self) } - pub async fn to_stream(&mut self, stream: &mut Socket) -> Result<(), OblivionException> { - stream.send(&self.plain_data()).await?; + pub async fn to_stream(&mut self, stream: &mut Socket) -> Result<()> { + stream.send(&self.plain_data()?).await?; Ok(()) } - pub async fn to_stream_with_salt( - &mut self, - stream: &mut Socket, - ) -> Result<(), OblivionException> { - stream.send(&self.plain_data()).await?; - stream.send(&self.plain_salt()).await?; + pub async fn to_stream_with_salt(&mut self, stream: &mut Socket) -> Result<()> { + stream.send(&self.plain_data()?).await?; + stream.send(&self.plain_salt()?).await?; Ok(()) } - pub fn plain_data(&mut self) -> Vec { + pub fn plain_data(&mut self) -> Result> { let public_key_bytes = self.public_key.unwrap().to_sec1_bytes().to_vec(); - let mut plain_data_bytes = length(&public_key_bytes).unwrap(); + let mut plain_data_bytes = length(&public_key_bytes)?; plain_data_bytes.extend(public_key_bytes); - plain_data_bytes + Ok(plain_data_bytes) } - pub fn plain_salt(&mut self) -> Vec { + pub fn plain_salt(&mut self) -> Result> { let salt_bytes = self.salt.as_ref().unwrap(); - let mut plain_salt_bytes = length(&salt_bytes).unwrap(); + let mut plain_salt_bytes = length(&salt_bytes)?; plain_salt_bytes.extend(salt_bytes); - plain_salt_bytes + Ok(plain_salt_bytes) } pub fn get_aes_key(&mut self) -> Vec { @@ -188,7 +177,7 @@ impl OED { } } - fn serialize_bytes(&self, data: &[u8], size: Option) -> Vec> { + fn serialize_bytes(&self, data: &[u8], size: Option) -> Result>> { let size = if size.is_none() { let size: usize = 1024; size @@ -202,7 +191,7 @@ impl OED { for i in (0..data_size).step_by(size) { let buffer = &data[i..std::cmp::min(i + size, data_size)]; - let buffer_length = length(&buffer.to_vec()).unwrap(); + let buffer_length = length(&buffer.to_vec())?; let mut serialized_chunk = Vec::with_capacity(buffer_length.len() + buffer.len()); if i + size > data_size { @@ -219,7 +208,7 @@ impl OED { } serialized_bytes.push(b"0000".to_vec()); - serialized_bytes + Ok(serialized_bytes) } pub fn from_json_or_string( @@ -332,7 +321,7 @@ impl OED { self.chunk_size = 0; for bytes in self - .serialize_bytes(&self.encrypted_data.as_ref().unwrap(), None) + .serialize_bytes(&self.encrypted_data.as_ref().unwrap(), None)? .iter() { stream.send(&bytes).await?;