From 7d8e352269c7c66822c82ed303b9b83371f0f83e Mon Sep 17 00:00:00 2001
From: Joey Sacchini <joey@sacchini.net>
Date: Tue, 15 Dec 2020 13:22:59 -0500
Subject: read untyped packets

---
 src/connection.rs | 10 +++++++-
 src/reader.rs     | 72 +++++++++++++++++++++++++++++++++++++++++++++++--------
 2 files changed, 71 insertions(+), 11 deletions(-)

diff --git a/src/connection.rs b/src/connection.rs
index f53f07d..486eed8 100644
--- a/src/connection.rs
+++ b/src/connection.rs
@@ -3,7 +3,7 @@ use crate::cfb8::CipherError;
 use crate::reader::{CraftReader, CraftSyncReader, ReadResult};
 use crate::wrapper::{CraftIo, CraftWrapper};
 use crate::writer::{CraftSyncWriter, CraftWriter, WriteResult};
-use mcproto_rs::protocol::{Packet, RawPacket, State};
+use mcproto_rs::protocol::{Packet, RawPacket, State, Id};
 #[cfg(feature = "gat")]
 use mcproto_rs::protocol::PacketKind;
 #[cfg(any(feature = "futures-io", feature = "tokio-io"))]
@@ -79,6 +79,10 @@ where
     {
         self.reader.read_raw_packet::<P>()
     }
+
+    fn read_raw_untyped_packet(&mut self) -> ReadResult<(Id, &[u8])> {
+        self.reader.read_raw_untyped_packet()
+    }
 }
 
 impl<R, W> CraftSyncWriter for CraftConnection<R, W>
@@ -141,6 +145,10 @@ where
     {
         self.reader.read_raw_packet_async::<P>().await
     }
+
+    async fn read_raw_untyped_packet_async(&mut self) -> ReadResult<(Id, &[u8])> {
+        self.reader.read_raw_untyped_packet_async().await
+    }
 }
 
 #[cfg(any(feature = "futures-io", feature = "tokio-io"))]
diff --git a/src/reader.rs b/src/reader.rs
index e6a00e1..4703a27 100644
--- a/src/reader.rs
+++ b/src/reader.rs
@@ -88,6 +88,8 @@ pub trait CraftAsyncReader {
     async fn read_raw_packet_async<P>(&mut self) -> ReadResult<P::RawPacket<'_>>
     where
         P: PacketKind;
+
+    async fn read_raw_untyped_packet_async(&mut self) -> ReadResult<(Id, &[u8])>;
 }
 
 pub trait CraftSyncReader {
@@ -116,6 +118,8 @@ pub trait CraftSyncReader {
     fn read_raw_packet<P>(&mut self) -> ReadResult<P::RawPacket<'_>>
     where
         P: PacketKind;
+
+    fn read_raw_untyped_packet(&mut self) -> ReadResult<(Id, &[u8])>;
 }
 
 pub struct CraftReader<R> {
@@ -196,6 +200,10 @@ where
     {
         self.read_raw_packet_inner::<P::RawPacket<'_>>()
     }
+
+    fn read_raw_untyped_packet(&mut self) -> ReadResult<(Id, &[u8])> {
+        self.read_untyped_packet_inner()
+    }
 }
 
 #[cfg(any(feature = "futures-io", feature = "tokio-io"))]
@@ -219,20 +227,40 @@ where
     {
         self.read_raw_packet_inner_async::<P::RawPacket<'_>>().await
     }
+
+    async fn read_raw_untyped_packet_async(&mut self) -> ReadResult<(Id, &[u8])> {
+        self.read_raw_untyped_packet_inner_async().await
+    }
 }
 
 impl<R> CraftReader<R>
 where
     R: io::Read,
 {
+    fn read_untyped_packet_inner(&mut self) -> ReadResult<(Id, &[u8])> {
+        if let Some(primary_packet_len) = self.read_raw_inner()? {
+            self.read_untyped_packet_in_buf(primary_packet_len)
+        } else {
+            Ok(None)
+        }
+    }
+
     fn read_raw_packet_inner<'a, P>(&'a mut self) -> ReadResult<P>
     where
         P: RawPacket<'a>
     {
+        if let Some(primary_packet_len) = self.read_raw_inner()? {
+            self.read_packet_in_buf(primary_packet_len)
+        } else {
+            Ok(None)
+        }
+    }
+
+    fn read_raw_inner(&mut self) -> ReadResult<usize> {
         self.move_ready_data_to_front();
         let primary_packet_len = rr_unwrap!(self.read_packet_len_sync()).0 as usize;
         self.ensure_n_ready_sync(primary_packet_len)?;
-        self.read_packet_in_buf(primary_packet_len)
+        Ok(Some(primary_packet_len))
     }
 
     fn read_packet_len_sync(&mut self) -> ReadResult<VarInt> {
@@ -266,10 +294,26 @@ where
     where
         P: RawPacket<'a>
     {
+        if let Some(primary_packet_len) = self.read_raw_inner_async().await? {
+            self.read_packet_in_buf(primary_packet_len)
+        } else {
+            Ok(None)
+        }
+    }
+
+    async fn read_raw_untyped_packet_inner_async(&mut self) -> ReadResult<(Id, &[u8])> {
+        if let Some(primary_packet_len) = self.read_raw_inner_async().await? {
+            self.read_untyped_packet_in_buf(primary_packet_len)
+        } else {
+            Ok(None)
+        }
+    }
+
+    async fn read_raw_inner_async(&mut self) -> ReadResult<usize> {
         self.move_ready_data_to_front();
         let primary_packet_len = rr_unwrap!(self.read_packet_len_async().await).0 as usize;
         self.ensure_n_ready_async(primary_packet_len).await?;
-        self.read_packet_in_buf(primary_packet_len)
+        Ok(Some(primary_packet_len))
     }
 
     async fn read_packet_len_async(&mut self) -> ReadResult<VarInt> {
@@ -391,9 +435,7 @@ impl<R> CraftReader<R> {
         }
     }
 
-    fn read_packet_in_buf<'a, P>(&'a mut self, size: usize) -> ReadResult<P>
-    where
-        P: RawPacket<'a>,
+    fn read_untyped_packet_in_buf(&mut self, size: usize) -> ReadResult<(Id, &[u8])>
     {
         // find data in buf
         let offset = self.raw_offset;
@@ -434,16 +476,26 @@ impl<R> CraftReader<R> {
         let packet_buf = buf;
 
         let (raw_id, body_buf) = dsz_unwrap!(packet_buf, VarInt);
-
         let id = Id {
             id: raw_id.0,
             state: self.state.clone(),
-            direction: self.direction.clone(),
+            direction: self.direction.clone()
         };
 
-        match P::create(id, body_buf) {
-            Ok(raw) => Ok(Some(raw)),
-            Err(err) => Err(err.into()),
+        Ok(Some((id, body_buf)))
+    }
+
+    fn read_packet_in_buf<'a, P>(&'a mut self, size: usize) -> ReadResult<P>
+    where
+        P: RawPacket<'a>,
+    {
+        if let Some((id, body_buf)) = self.read_untyped_packet_in_buf(size)? {
+            match P::create(id, body_buf) {
+                Ok(raw) => Ok(Some(raw)),
+                Err(err) => Err(err.into()),
+            }
+        } else {
+            None
         }
     }
 
-- 
cgit