support for raw frames in PacketEncoder

crates/valence_protocol/src/decode.rs

@@ -15,13 +15,13 @@ type Cipher = cfb8::Decryptor<aes::Aes128>;
 
 #[derive(Default)]
 pub struct PacketDecoder {
-    buf: BytesMut,
+    pub buf: BytesMut,
     #[cfg(feature = "compression")]
-    decompress_buf: BytesMut,
+    pub decompress_buf: BytesMut,
     #[cfg(feature = "compression")]
-    threshold: CompressionThreshold,
+    pub threshold: CompressionThreshold,
     #[cfg(feature = "encryption")]
-    cipher: Option<Cipher>,
+    pub cipher: Option<Cipher>,
 }
 
 impl PacketDecoder {

crates/valence_protocol/src/encode.rs

@@ -8,6 +8,7 @@ use anyhow::ensure;
 use bytes::{BufMut, BytesMut};
 use tracing::warn;
 
+use crate::decode::PacketFrame;
 use crate::var_int::VarInt;
 use crate::{CompressionThreshold, Encode, Packet, MAX_PACKET_SIZE};
 
@@ -18,13 +19,13 @@ type Cipher = cfb8::Encryptor<aes::Aes128>;
 
 #[derive(Default)]
 pub struct PacketEncoder {
-    buf: BytesMut,
+    pub buf: BytesMut,
     #[cfg(feature = "compression")]
-    compress_buf: Vec<u8>,
+    pub compress_buf: Vec<u8>,
     #[cfg(feature = "compression")]
-    threshold: CompressionThreshold,
+    pub threshold: CompressionThreshold,
     #[cfg(feature = "encryption")]
-    cipher: Option<Cipher>,
+    pub cipher: Option<Cipher>,
 }
 
 impl PacketEncoder {
@@ -37,37 +38,12 @@ impl PacketEncoder {
         self.buf.extend_from_slice(bytes)
     }
 
-    pub fn prepend_packet<P>(&mut self, pkt: &P) -> anyhow::Result<()>
-    where
-        P: Packet + Encode,
-    {
-        let start_len = self.buf.len();
-        self.append_packet(pkt)?;
-
-        let end_len = self.buf.len();
-        let total_packet_len = end_len - start_len;
-
-        // 1) Move everything back by the length of the packet.
-        // 2) Move the packet to the new space at the front.
-        // 3) Truncate the old packet away.
-        self.buf.put_bytes(0, total_packet_len);
-        self.buf.copy_within(..end_len, total_packet_len);
-        self.buf.copy_within(total_packet_len + start_len.., 0);
-        self.buf.truncate(end_len);
-
-        Ok(())
-    }
-
-    #[allow(clippy::needless_borrows_for_generic_args)]
-    pub fn append_packet<P>(&mut self, pkt: &P) -> anyhow::Result<()>
-    where
-        P: Packet + Encode,
-    {
-        let start_len = self.buf.len();
-
-        pkt.encode_with_id((&mut self.buf).writer())?;
-
-        let data_len = self.buf.len() - start_len;
+    /// frames the bytes in a range from `from` to the end of the buffer:
+    ///     adding a packet length varint to the start of the frame
+    ///     adding a data length varint after the packet length, if compression is enabled
+    ///     compressing the packet, if compression is enabled
+    pub fn enframe_from(&mut self, from: usize) -> anyhow::Result<()> {
+        let data_len = self.buf.len() - from;
 
         #[cfg(feature = "compression")]
         if self.threshold.0 >= 0 {
@@ -77,7 +53,7 @@ impl PacketEncoder {
             use flate2::Compression;
 
             if data_len > self.threshold.0 as usize {
-                let mut z = ZlibEncoder::new(&self.buf[start_len..], Compression::new(4));
+                let mut z = ZlibEncoder::new(&self.buf[from..], Compression::new(4));
 
                 self.compress_buf.clear();
 
@@ -92,7 +68,7 @@ impl PacketEncoder {
 
                 drop(z);
 
-                self.buf.truncate(start_len);
+                self.buf.truncate(from);
 
                 let mut writer = (&mut self.buf).writer();
 
@@ -114,9 +90,9 @@ impl PacketEncoder {
 
                 self.buf.put_bytes(0, data_prefix_len);
                 self.buf
-                    .copy_within(start_len..start_len + data_len, start_len + data_prefix_len);
+                    .copy_within(from..from + data_len, from + data_prefix_len);
 
-                let mut front = &mut self.buf[start_len..];
+                let mut front = &mut self.buf[from..];
 
                 VarInt(packet_len as i32).encode(&mut front)?;
                 // Zero for no compression on this packet.
@@ -137,14 +113,77 @@ impl PacketEncoder {
 
         self.buf.put_bytes(0, packet_len_size);
         self.buf
-            .copy_within(start_len..start_len + data_len, start_len + packet_len_size);
+            .copy_within(from..from + data_len, from + packet_len_size);
 
-        let front = &mut self.buf[start_len..];
+        let front = &mut self.buf[from..];
         VarInt(packet_len as i32).encode(front)?;
 
         Ok(())
     }
 
+    fn move_to_back(&mut self, from: usize) {
+        // 1) Move everything back by the length of the packet.
+        // 2) Move the packet to the new space at the front.
+        // 3) Truncate the old packet away.
+        let to = self.buf.len();
+        let len = to - from;
+
+        self.buf.put_bytes(0, len);
+        self.buf.copy_within(..to, len);
+        self.buf.copy_within(to.., 0);
+        self.buf.truncate(to);
+    }
+
+    pub fn append_packet_frame(&mut self, frame: &PacketFrame) -> anyhow::Result<()> {
+        let start_len = self.buf.len();
+        VarInt(frame.id).encode((&mut self.buf).writer())?;
+        self.append_bytes(&frame.body);
+        self.enframe_from(start_len)?;
+        Ok(())
+    }
+
+    pub fn prepend_packet_frame<P>(&mut self, frame: &PacketFrame) -> anyhow::Result<()> {
+        let start_len = self.buf.len();
+        self.append_packet_frame(frame)?;
+        self.move_to_back(start_len);
+        Ok(())
+    }
+
+    pub fn append_raw_frame(&mut self, frame: &[u8]) -> anyhow::Result<()> {
+        let start_len = self.buf.len();
+        self.append_bytes(frame);
+        self.enframe_from(start_len)?;
+        Ok(())
+    }
+
+    pub fn prepend_raw_frame<P>(&mut self, frame: &[u8]) -> anyhow::Result<()> {
+        let start_len = self.buf.len();
+        self.append_raw_frame(frame)?;
+        self.move_to_back(start_len);
+        Ok(())
+    }
+
+    #[allow(clippy::needless_borrows_for_generic_args)]
+    pub fn append_packet<P>(&mut self, pkt: &P) -> anyhow::Result<()>
+    where
+        P: Packet + Encode,
+    {
+        let start_len = self.buf.len();
+        pkt.encode_with_id((&mut self.buf).writer())?;
+        self.enframe_from(start_len)?;
+        Ok(())
+    }
+
+    pub fn prepend_packet<P>(&mut self, pkt: &P) -> anyhow::Result<()>
+    where
+        P: Packet + Encode,
+    {
+        let start_len = self.buf.len();
+        self.append_packet(pkt)?;
+        self.move_to_back(start_len);
+        Ok(())
+    }
+
     /// Takes all the packets written so far and encrypts them if encryption is
     /// enabled.
     pub fn take(&mut self) -> BytesMut {

crates/valence_protocol/src/lib.rs

@@ -296,6 +296,12 @@ pub enum PacketState {
     Play,
 }
 
+impl Default for PacketState {
+    fn default() -> Self {
+        Self::Handshaking
+    }
+}
+
 #[allow(dead_code)]
 #[cfg(test)]
 mod tests {

crates/valence_protocol/src/var_long.rs

@@ -43,61 +43,6 @@ impl VarLong {
 }
 
 impl Encode for VarLong {
-    // Adapted from VarInt-Simd encode
-    // https://github.com/as-com/varint-simd/blob/0f468783da8e181929b01b9c6e9f741c1fe09825/src/encode/mod.rs#L71
-    #[cfg(all(
-        any(target_arch = "x86", target_arch = "x86_64"),
-        not(target_os = "macos")
-    ))]
-    fn encode(&self, mut w: impl Write) -> anyhow::Result<()> {
-        #[cfg(target_arch = "x86")]
-        use std::arch::x86::*;
-        #[cfg(target_arch = "x86_64")]
-        use std::arch::x86_64::*;
-
-        // Break the number into 7-bit parts and spread them out into a vector
-        let mut res = [0_u64; 2];
-        {
-            let x = self.0 as u64;
-
-            res[0] = unsafe { _pdep_u64(x, 0x7f7f7f7f7f7f7f7f) };
-            res[1] = unsafe { _pdep_u64(x >> 56, 0x000000000000017f) }
-        };
-        let stage1: __m128i = unsafe { std::mem::transmute(res) };
-
-        // Create a mask for where there exist values
-        // This signed comparison works because all MSBs should be cleared at this point
-        // Also handle the special case when num == 0
-        let minimum =
-            unsafe { _mm_set_epi8(0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0xff_u8 as i8) };
-        let exists = unsafe { _mm_or_si128(_mm_cmpgt_epi8(stage1, _mm_setzero_si128()), minimum) };
-        let bits = unsafe { _mm_movemask_epi8(exists) };
-
-        // Count the number of bytes used
-        let bytes_needed = 32 - bits.leading_zeros() as u8; // lzcnt on supported CPUs
-
-        // Fill that many bytes into a vector
-        let ascend = unsafe { _mm_setr_epi8(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15) };
-        let mask = unsafe { _mm_cmplt_epi8(ascend, _mm_set1_epi8(bytes_needed as i8)) };
-
-        // Shift it down 1 byte so the last MSB is the only one set, and make sure only
-        // the MSB is set
-        let shift = unsafe { _mm_bsrli_si128(mask, 1) };
-        let msbmask = unsafe { _mm_and_si128(shift, _mm_set1_epi8(128_u8 as i8)) };
-
-        // Merge the MSB bits into the vector
-        let merged = unsafe { _mm_or_si128(stage1, msbmask) };
-        let bytes = unsafe { std::mem::transmute::<__m128i, [u8; 16]>(merged) };
-
-        w.write_all(unsafe { bytes.get_unchecked(..bytes_needed as usize) })?;
-
-        Ok(())
-    }
-
-    #[cfg(any(
-        not(any(target_arch = "x86", target_arch = "x86_64")),
-        target_os = "macos"
-    ))]
     fn encode(&self, mut w: impl Write) -> anyhow::Result<()> {
         use byteorder::WriteBytesExt;