1 files changed, 41 insertions, 22 deletions

diff --git a/azalea-protocol/src/write.rs b/azalea-protocol/src/write.rs
index ae9a6829..30710f8b 100755
--- a/azalea-protocol/src/write.rs
+++ b/azalea-protocol/src/write.rs
@@ -3,49 +3,67 @@ use async_compression::tokio::bufread::ZlibEncoder;
 use azalea_buf::Writable;
 use azalea_crypto::Aes128CfbEnc;
 use std::fmt::Debug;
+use thiserror::Error;
 use tokio::io::{AsyncReadExt, AsyncWrite, AsyncWriteExt};
 
-fn frame_prepender(data: &mut Vec<u8>) -> Result<Vec<u8>, String> {
+fn frame_prepender(data: &mut Vec<u8>) -> Result<Vec<u8>, std::io::Error> {
     let mut buf = Vec::new();
-    buf.write_varint(data.len() as i32)
-        .map_err(|e| e.to_string())?;
+    buf.write_varint(data.len() as i32)?;
     buf.append(data);
     Ok(buf)
 }
 
-fn packet_encoder<P: ProtocolPacket + std::fmt::Debug>(packet: &P) -> Result<Vec<u8>, String> {
+#[derive(Error, Debug)]
+pub enum PacketEncodeError {
+    #[error("{0}")]
+    Io(#[from] std::io::Error),
+    #[error("Packet too big (is {actual} bytes, should be less than {maximum}): {packet_string}")]
+    TooBig {
+        actual: usize,
+        maximum: usize,
+        packet_string: String,
+    },
+}
+
+fn packet_encoder<P: ProtocolPacket + std::fmt::Debug>(
+    packet: &P,
+) -> Result<Vec<u8>, PacketEncodeError> {
     let mut buf = Vec::new();
-    buf.write_varint(packet.id() as i32)
-        .map_err(|e| e.to_string())?;
-    packet.write(&mut buf).map_err(|e| e.to_string())?;
+    buf.write_varint(packet.id() as i32)?;
+    packet.write(&mut buf)?;
     if buf.len() > MAXIMUM_UNCOMPRESSED_LENGTH as usize {
-        return Err(format!(
-            "Packet too big (is {} bytes, should be less than {}): {:?}",
-            buf.len(),
-            MAXIMUM_UNCOMPRESSED_LENGTH,
-            packet
-        ));
+        return Err(PacketEncodeError::TooBig {
+            actual: buf.len(),
+            maximum: MAXIMUM_UNCOMPRESSED_LENGTH as usize,
+            packet_string: format!("{:?}", packet),
+        });
     }
     Ok(buf)
 }
 
-async fn compression_encoder(data: &[u8], compression_threshold: u32) -> Result<Vec<u8>, String> {
+#[derive(Error, Debug)]
+pub enum PacketCompressError {
+    #[error("{0}")]
+    Io(#[from] std::io::Error),
+}
+
+async fn compression_encoder(
+    data: &[u8],
+    compression_threshold: u32,
+) -> Result<Vec<u8>, PacketCompressError> {
     let n = data.len();
     // if it's less than the compression threshold, don't compress
     if n < compression_threshold as usize {
         let mut buf = Vec::new();
-        buf.write_varint(0).map_err(|e| e.to_string())?;
-        buf.write_all(data).await.map_err(|e| e.to_string())?;
+        buf.write_varint(0)?;
+        buf.write_all(data).await?;
         Ok(buf)
     } else {
         // otherwise, compress
         let mut deflater = ZlibEncoder::new(data);
         // write deflated data to buf
         let mut buf = Vec::new();
-        deflater
-            .read_to_end(&mut buf)
-            .await
-            .map_err(|e| e.to_string())?;
+        deflater.read_to_end(&mut buf).await?;
         Ok(buf)
     }
 }
@@ -55,7 +73,8 @@ pub async fn write_packet<P, W>(
     stream: &mut W,
     compression_threshold: Option<u32>,
     cipher: &mut Option<Aes128CfbEnc>,
-) where
+) -> std::io::Result<()>
+where
     P: ProtocolPacket + Debug,
     W: AsyncWrite + Unpin + Send,
 {
@@ -68,5 +87,5 @@ pub async fn write_packet<P, W>(
     if let Some(cipher) = cipher {
         azalea_crypto::encrypt_packet(cipher, &mut buf);
     }
-    stream.write_all(&buf).await.unwrap();
+    stream.write_all(&buf).await
 }