Skip to main content

steel_protocol/
packet_traits.rs

1//! # Steel Protocol Packet Traits
2//!
3//! This module contains the traits for the packets.
4use std::{
5    io::{Cursor, Write},
6    num::NonZeroU32,
7    sync::Arc,
8};
9
10use flate2::{Compression, write::ZlibEncoder};
11use serde::Deserialize;
12use steel_utils::{
13    FrontVec,
14    codec::VarInt,
15    serial::{ReadFrom, WriteTo},
16};
17
18use crate::utils::{ConnectionProtocol, MAX_PACKET_DATA_SIZE, MAX_PACKET_SIZE, PacketError};
19
20// These are the network read/write traits
21/// A trait for packets sent from the server to the client.
22pub trait ServerPacket: ReadFrom {
23    /// Reads a packet from the given data.
24    fn read_packet(data: &mut Cursor<&[u8]>) -> Result<Self, PacketError> {
25        Self::read(data).map_err(PacketError::from)
26    }
27}
28
29/// A trait for packets sent from the client to the server.
30pub trait ClientPacket: WriteTo {
31    /// Writes the packet to the given writer.
32    ///
33    /// # Errors
34    /// - If the packet fails to write.
35    /// - If the protocol is invalid.
36    fn write_packet(
37        &self,
38        writer: &mut impl Write,
39        protocol: ConnectionProtocol,
40    ) -> Result<(), PacketError> {
41        let packet_id = self
42            .get_id(protocol)
43            .ok_or(PacketError::InvalidProtocol(format!(
44                "Invalid protocol {protocol:?}"
45            )))?;
46        VarInt(packet_id).write(writer)?;
47        self.write(writer).map_err(PacketError::from)
48    }
49
50    /// Gets the ID of the packet for the given protocol.
51    fn get_id(&self, protocol: ConnectionProtocol) -> Option<i32>;
52}
53
54/// Information about compression.
55#[derive(Copy, Clone, Debug, Deserialize)]
56pub struct CompressionInfo {
57    /// The compression threshold used when compression is enabled.
58    /// Its an `NonZeroU32` to allow for nullptr optimization in `Option<Self>` cases
59    pub threshold: NonZeroU32,
60    /// A value between `0..9`.
61    /// `1` = Optimize for the best speed of encoding.
62    /// `9` = Optimize for the size of data being encoded.
63    pub level: i32,
64}
65
66impl Default for CompressionInfo {
67    #[expect(
68        clippy::unwrap_used,
69        reason = "256 is a known nonzero compression threshold"
70    )]
71    fn default() -> Self {
72        Self {
73            threshold: NonZeroU32::new(256).unwrap(),
74            level: 4,
75        }
76    }
77}
78
79/// Represents an encoded clientbound packet, optionally applying compression based on threshold and level.
80///
81/// # Packet Size Limits
82/// - Maximum packet size: 2097151 bytes (2^21 - 1, max 3-byte `VarInt`)
83/// - Maximum uncompressed size for compressed packets: 8388608 bytes (2^23)
84/// - Length field must not exceed 3 bytes
85///
86/// # Packet Encoding Format
87///
88/// **Without Compression:**
89/// ```text
90/// [Length: VarInt]     Length of (Packet ID + Data)
91/// [Packet ID: VarInt]  Protocol ID from packet report
92/// [Data: Byte Array]   Packet payload
93/// ```
94///
95/// **With Compression (size >= threshold):**
96/// ```text
97/// [Length: VarInt]     Length of (Data Length + compressed data)
98/// [Data Length: VarInt] Length of uncompressed (Packet ID + Data)
99/// [Compressed Data]    zlib compressed (Packet ID + Data)
100/// ```
101///
102/// **With Compression (size < threshold):**
103/// ```text
104/// [Length: VarInt]     Length of (Data Length + uncompressed data)
105/// [Data Length: VarInt] 0 to indicate uncompressed
106/// [Packet ID: VarInt]  Protocol ID from packet report
107/// [Data: Byte Array]   Uncompressed packet payload
108/// ```
109///
110/// Compression is only applied when:
111/// 1. Compression is enabled via Set Compression packet
112/// 2. The uncompressed data length meets/exceeds the threshold
113/// 3. The threshold is non-negative
114#[derive(Clone)]
115pub struct EncodedPacket {
116    // This is optimized for reduces allocation
117    /// The encoded data.
118    pub encoded_data: Arc<FrontVec>,
119}
120
121impl EncodedPacket {
122    fn from_data_uncompressed(mut packet_data: FrontVec) -> Result<Self, PacketError> {
123        let data_len = packet_data.len();
124        let varint_size = VarInt::written_size(data_len as i32);
125
126        let complete_len = varint_size + data_len;
127        if complete_len > MAX_PACKET_SIZE {
128            return Err(PacketError::TooLong(complete_len));
129        }
130
131        VarInt(data_len as i32).set_in_front(&mut packet_data, varint_size);
132
133        Ok(Self {
134            encoded_data: Arc::new(packet_data),
135        })
136    }
137
138    fn from_packet_data(
139        mut packet_data: FrontVec,
140        compression: CompressionInfo,
141    ) -> Result<Self, PacketError> {
142        let data_len = packet_data.len();
143        // We dont need any more size check to convert to i32 as MAX_PACKET_DATA_SIZE < i32::MAX
144        if data_len + VarInt::MAX_SIZE * 2 > MAX_PACKET_DATA_SIZE {
145            Err(PacketError::TooLong(data_len))?;
146        }
147
148        if data_len >= compression.threshold.get() as _ {
149            let mut buf = FrontVec::new(10);
150            let mut compressor =
151                ZlibEncoder::new(&mut buf, Compression::new(compression.level as u32));
152
153            compressor
154                .write_all(&packet_data)
155                .map_err(|e| PacketError::CompressionFailed(e.to_string()))?;
156            compressor
157                .finish()
158                .map_err(|e| PacketError::CompressionFailed(e.to_string()))?;
159
160            // compressed data cant be larger so we dont need to check the size again
161            let varint_size = VarInt::written_size(data_len as i32);
162            let full_len = varint_size + buf.len();
163            let full_varint_size = VarInt::written_size(full_len as i32);
164
165            VarInt(data_len as i32).set_in_front(&mut buf, varint_size);
166            VarInt(full_len as i32).set_in_front(&mut buf, full_varint_size);
167            log::trace!(
168                "data length: {data_len}, full length: {full_len}, varint size: {varint_size}, full varint size: {full_varint_size}"
169            );
170
171            Ok(Self {
172                encoded_data: Arc::new(buf),
173            })
174        } else {
175            // Pushed before data:
176            // Length of (Data Length) + length of compressed (Packet ID + Data)
177            // 0 to indicate uncompressed
178
179            let data_len_with_header = data_len + 1;
180            let varint_size = VarInt::written_size(data_len_with_header as i32);
181
182            VarInt(0).set_in_front(&mut packet_data, 1);
183            VarInt(data_len_with_header as i32).set_in_front(&mut packet_data, varint_size);
184
185            Ok(Self {
186                encoded_data: Arc::new(packet_data),
187            })
188        }
189    }
190
191    /// Creates a new `EncodedPacket` from a bare packet.
192    ///
193    /// # Errors
194    /// - If the packet fails to write.
195    /// - If the packet fails to compress.
196    pub fn from_bare<P: ClientPacket>(
197        packet: P,
198        compression: Option<CompressionInfo>,
199        protocol: ConnectionProtocol,
200    ) -> Result<Self, PacketError> {
201        let buf = Self::write_vec(packet, protocol)?;
202        Self::from_data(buf, compression)
203    }
204
205    fn write_vec<P: ClientPacket>(
206        packet: P,
207        protocol: ConnectionProtocol,
208    ) -> Result<FrontVec, PacketError> {
209        let mut buf = FrontVec::new(6);
210        packet.write_packet(&mut buf, protocol)?;
211        Ok(buf)
212    }
213
214    fn from_data(buf: FrontVec, compression: Option<CompressionInfo>) -> Result<Self, PacketError> {
215        if let Some(compression) = compression {
216            Self::from_packet_data(buf, compression)
217        } else {
218            Self::from_data_uncompressed(buf)
219        }
220    }
221}