Skip to main content

azalea_protocol/
read.rs

1//! Read packets from a stream.
2#![allow(clippy::redundant_field_names, reason = "Triggered by `Error` derive")]
3
4use std::{
5    backtrace::Backtrace,
6    env,
7    fmt::Debug,
8    io,
9    io::{Cursor, Read},
10    sync::LazyLock,
11};
12
13use azalea_buf::{AzBufVar, BufReadError};
14use azalea_crypto::Aes128CfbDec;
15use flate2::read::ZlibDecoder;
16use futures::StreamExt;
17use futures_lite::future;
18use thiserror::Error;
19use tokio::io::AsyncRead;
20use tokio_util::{
21    bytes::Buf,
22    codec::{BytesCodec, FramedRead},
23};
24use tracing::trace;
25
26use crate::packets::ProtocolPacket;
27
28#[derive(Debug, Error)]
29pub enum ReadPacketError {
30    #[error("Error reading packet {packet_name} (id {packet_id}): {source}")]
31    Parse {
32        packet_id: u32,
33        packet_name: String,
34        backtrace: Box<Backtrace>,
35        source: BufReadError,
36    },
37    #[error("Unknown packet id {id} in state {state_name}")]
38    UnknownPacketId { state_name: String, id: u32 },
39    #[error("Couldn't read packet id")]
40    ReadPacketId { source: BufReadError },
41    #[error(transparent)]
42    Decompress {
43        #[from]
44        #[backtrace]
45        source: DecompressionError,
46    },
47    #[error(transparent)]
48    FrameSplitter {
49        #[from]
50        #[backtrace]
51        source: FrameSplitterError,
52    },
53    #[error("Leftover data after reading packet {packet_name}: {data:?}")]
54    LeftoverData { data: Vec<u8>, packet_name: String },
55    #[error(transparent)]
56    IoError {
57        #[from]
58        #[backtrace]
59        source: io::Error,
60    },
61    #[error("Connection closed")]
62    ConnectionClosed,
63}
64
65#[derive(Debug, Error)]
66pub enum FrameSplitterError {
67    #[error("Couldn't read VarInt length for packet. The previous packet may have been corrupted")]
68    LengthRead {
69        #[from]
70        source: BufReadError,
71    },
72    #[error("Io error")]
73    Io {
74        #[from]
75        #[backtrace]
76        source: io::Error,
77    },
78    #[error("Packet is longer than {max} bytes (is {size})")]
79    BadLength { max: usize, size: usize },
80    #[error("Connection reset by peer")]
81    ConnectionReset,
82    #[error("Connection closed")]
83    ConnectionClosed,
84}
85
86/// Read a length, then read that amount of bytes from the `Cursor<Vec<u8>>`. If
87/// there's not enough data, `None` is returned.
88fn parse_frame(buffer: &mut Cursor<Vec<u8>>) -> Result<Box<[u8]>, FrameSplitterError> {
89    // copy the buffer first and read from the copy, then once we make sure
90    // the packet is all good we read it fully
91    let mut buffer_copy = Cursor::new(&buffer.get_ref()[buffer.position() as usize..]);
92    // Packet Length
93    let length = match u32::azalea_read_var(&mut buffer_copy) {
94        Ok(length) => length as usize,
95        Err(err) => match err {
96            BufReadError::Io { source } => return Err(FrameSplitterError::Io { source }),
97            _ => return Err(err.into()),
98        },
99    };
100
101    if length > buffer_copy.remaining() {
102        return Err(FrameSplitterError::BadLength {
103            max: buffer_copy.remaining(),
104            size: length,
105        });
106    }
107
108    // we read from the copy and we know it's legit, so we can take those bytes
109    // from the real buffer now
110
111    // the length of the varint that says the length of the whole packet
112    let varint_length = buffer.remaining() - buffer_copy.remaining();
113
114    buffer.advance(varint_length);
115    let data =
116        buffer.get_ref()[buffer.position() as usize..buffer.position() as usize + length].to_vec();
117    buffer.advance(length);
118
119    if buffer.position() == buffer.get_ref().len() as u64 {
120        // reset the inner vec once we've reached the end of the buffer so we don't keep
121        // leaking memory
122        buffer.get_mut().clear();
123
124        // we just cap the capacity to 64KB instead of resetting it to save some
125        // allocations.
126        // and the reason we bother capping it at all is to avoid wasting memory if we
127        // get a big packet once and then never again.
128        buffer.get_mut().shrink_to(1024 * 64);
129
130        buffer.set_position(0);
131    }
132
133    Ok(data.into_boxed_slice())
134}
135
136fn frame_splitter(buffer: &mut Cursor<Vec<u8>>) -> Result<Option<Box<[u8]>>, FrameSplitterError> {
137    // https://tokio.rs/tokio/tutorial/framing
138    let read_frame = parse_frame(buffer);
139    match read_frame {
140        Ok(frame) => return Ok(Some(frame)),
141        Err(err) => match err {
142            FrameSplitterError::BadLength { .. } | FrameSplitterError::Io { .. } => {
143                // we probably just haven't read enough yet
144            }
145            _ => return Err(err),
146        },
147    }
148
149    Ok(None)
150}
151
152pub fn deserialize_packet<P: ProtocolPacket + Debug>(
153    stream: &mut Cursor<&[u8]>,
154) -> Result<P, Box<ReadPacketError>> {
155    // Packet ID
156    let packet_id =
157        u32::azalea_read_var(stream).map_err(|e| ReadPacketError::ReadPacketId { source: e })?;
158    P::read(packet_id, stream)
159}
160
161// this is always true in multiplayer, false in singleplayer
162static VALIDATE_DECOMPRESSED: bool = true;
163
164pub static MAXIMUM_UNCOMPRESSED_LENGTH: u32 = 8_388_608;
165
166#[derive(Debug, Error)]
167pub enum DecompressionError {
168    #[error("Couldn't read VarInt length for data")]
169    LengthReadError {
170        #[from]
171        source: BufReadError,
172    },
173    #[error("Io error")]
174    Io {
175        #[from]
176        #[backtrace]
177        source: io::Error,
178    },
179    #[error("Badly compressed packet - size of {size} is below server threshold of {threshold}")]
180    BelowCompressionThreshold { size: u32, threshold: u32 },
181    #[error(
182        "Badly compressed packet - size of {size} is larger than protocol maximum of {maximum}"
183    )]
184    AboveCompressionThreshold { size: u32, maximum: u32 },
185}
186
187/// Get the decompressed bytes from a packet.
188///
189/// The stream must have already been decrypted before passing it to this
190/// function.
191pub fn compression_decoder(
192    stream: &mut Cursor<&[u8]>,
193    compression_threshold: u32,
194) -> Result<Box<[u8]>, DecompressionError> {
195    // Data Length
196    let n = u32::azalea_read_var(stream)?;
197    if n == 0 {
198        // no data size, no compression
199        let buf = stream.get_ref()[stream.position() as usize..]
200            .to_vec()
201            .into_boxed_slice();
202        stream.set_position(stream.get_ref().len() as u64);
203        return Ok(buf);
204    }
205
206    if VALIDATE_DECOMPRESSED {
207        if n < compression_threshold {
208            return Err(DecompressionError::BelowCompressionThreshold {
209                size: n,
210                threshold: compression_threshold,
211            });
212        }
213        if n > MAXIMUM_UNCOMPRESSED_LENGTH {
214            return Err(DecompressionError::AboveCompressionThreshold {
215                size: n,
216                maximum: MAXIMUM_UNCOMPRESSED_LENGTH,
217            });
218        }
219    }
220
221    // VALIDATE_DECOMPRESSED should always be true, so the max they can make us
222    // allocate here is 2mb
223    let mut decoded_buf = Vec::with_capacity(n as usize);
224
225    let mut decoder = ZlibDecoder::new(stream);
226    decoder.read_to_end(&mut decoded_buf)?;
227
228    Ok(decoded_buf.into_boxed_slice())
229}
230
231/// Read a single packet from a stream.
232///
233/// The buffer is required because servers may send multiple packets in the
234/// same frame, so we need to store the packet data that's left to read.
235///
236/// The current protocol state must be passed as a generic.
237///
238/// For the non-waiting version, see [`try_read_packet`].
239pub async fn read_packet<P: ProtocolPacket + Debug, R>(
240    stream: &mut R,
241    buffer: &mut Cursor<Vec<u8>>,
242    compression_threshold: Option<u32>,
243    cipher: &mut Option<Aes128CfbDec>,
244) -> Result<P, Box<ReadPacketError>>
245where
246    R: AsyncRead + Unpin + Send + Sync,
247{
248    let raw_packet = read_raw_packet(stream, buffer, compression_threshold, cipher).await?;
249    let packet = deserialize_packet(&mut Cursor::new(&raw_packet))?;
250    Ok(packet)
251}
252
253/// Try to read a single packet from a stream. Returns None if we haven't
254/// received a full packet yet.
255pub fn try_read_packet<P: ProtocolPacket + Debug, R>(
256    stream: &mut R,
257    buffer: &mut Cursor<Vec<u8>>,
258    compression_threshold: Option<u32>,
259    cipher: &mut Option<Aes128CfbDec>,
260) -> Result<Option<P>, Box<ReadPacketError>>
261where
262    R: AsyncRead + Unpin + Send + Sync,
263{
264    let Some(raw_packet) = try_read_raw_packet(stream, buffer, compression_threshold, cipher)?
265    else {
266        return Ok(None);
267    };
268    let packet = deserialize_packet(&mut Cursor::new(&raw_packet))?;
269    Ok(Some(packet))
270}
271
272pub async fn read_raw_packet<R>(
273    stream: &mut R,
274    buffer: &mut Cursor<Vec<u8>>,
275    compression_threshold: Option<u32>,
276    // this has to be a &mut Option<T> instead of an Option<&mut T> because
277    // otherwise the borrow checker complains about the cipher being moved
278    cipher: &mut Option<Aes128CfbDec>,
279) -> Result<Box<[u8]>, Box<ReadPacketError>>
280where
281    R: AsyncRead + Unpin + Send + Sync,
282{
283    loop {
284        if let Some(buf) = read_raw_packet_from_buffer::<R>(buffer, compression_threshold)? {
285            // we got a full packet!!
286            return Ok(buf);
287        };
288
289        let bytes = read_and_decrypt_frame(stream, cipher).await?;
290        buffer.get_mut().extend_from_slice(&bytes);
291    }
292}
293/// Read a packet from the stream, then if necessary decrypt it, decompress
294/// it, and split it.
295pub fn try_read_raw_packet<R>(
296    stream: &mut R,
297    buffer: &mut Cursor<Vec<u8>>,
298    compression_threshold: Option<u32>,
299    cipher: &mut Option<Aes128CfbDec>,
300) -> Result<Option<Box<[u8]>>, Box<ReadPacketError>>
301where
302    R: AsyncRead + Unpin + Send + Sync,
303{
304    loop {
305        if let Some(buf) = read_raw_packet_from_buffer::<R>(buffer, compression_threshold)? {
306            // we got a full packet!!
307            return Ok(Some(buf));
308        };
309        let Some(bytes) = try_read_and_decrypt_frame(stream, cipher)? else {
310            // no data received
311            return Ok(None);
312        };
313        // we got some data, so add it to the buffer and try again
314        buffer.get_mut().extend_from_slice(&bytes);
315    }
316}
317
318async fn read_and_decrypt_frame<R>(
319    stream: &mut R,
320    cipher: &mut Option<Aes128CfbDec>,
321) -> Result<Box<[u8]>, Box<ReadPacketError>>
322where
323    R: AsyncRead + Unpin + Send + Sync,
324{
325    let mut framed = FramedRead::new(stream, BytesCodec::new());
326
327    let Some(message) = framed.next().await else {
328        return Err(Box::new(ReadPacketError::ConnectionClosed));
329    };
330    let bytes = message.map_err(ReadPacketError::from)?;
331
332    let mut bytes = bytes.to_vec().into_boxed_slice();
333
334    // decrypt if necessary
335    if let Some(cipher) = cipher {
336        azalea_crypto::decrypt_packet(cipher, &mut bytes);
337    }
338
339    Ok(bytes)
340}
341fn try_read_and_decrypt_frame<R>(
342    stream: &mut R,
343    cipher: &mut Option<Aes128CfbDec>,
344) -> Result<Option<Box<[u8]>>, Box<ReadPacketError>>
345where
346    R: AsyncRead + Unpin + Send + Sync,
347{
348    let mut framed = FramedRead::new(stream, BytesCodec::new());
349
350    let Some(message) = future::block_on(future::poll_once(framed.next())) else {
351        // nothing yet
352        return Ok(None);
353    };
354    let Some(message) = message else {
355        return Err(Box::new(ReadPacketError::ConnectionClosed));
356    };
357    let bytes = message.map_err(ReadPacketError::from)?;
358    let mut bytes = bytes.to_vec().into_boxed_slice();
359
360    // decrypt if necessary
361    if let Some(cipher) = cipher {
362        azalea_crypto::decrypt_packet(cipher, &mut bytes);
363    }
364
365    Ok(Some(bytes))
366}
367
368pub fn read_raw_packet_from_buffer<R>(
369    buffer: &mut Cursor<Vec<u8>>,
370    compression_threshold: Option<u32>,
371) -> Result<Option<Box<[u8]>>, Box<ReadPacketError>>
372where
373    R: AsyncRead + Unpin + Send + Sync,
374{
375    let Some(mut buf) = frame_splitter(buffer).map_err(ReadPacketError::from)? else {
376        // no full packet yet :(
377        return Ok(None);
378    };
379
380    if let Some(compression_threshold) = compression_threshold {
381        buf = compression_decoder(&mut Cursor::new(&buf[..]), compression_threshold)
382            .map_err(ReadPacketError::from)?;
383    }
384
385    if tracing::enabled!(tracing::Level::TRACE) {
386        static DO_NOT_CUT_OFF_PACKET_LOGS: LazyLock<bool> = LazyLock::new(|| {
387            env::var("AZALEA_DO_NOT_CUT_OFF_PACKET_LOGS")
388                .map(|s| s == "1" || s == "true")
389                .unwrap_or(false)
390        });
391
392        let buf_string: String = {
393            if !*DO_NOT_CUT_OFF_PACKET_LOGS && buf.len() > 500 {
394                let cut_off_buf = &buf[..500];
395                format!("{cut_off_buf:?}...")
396            } else {
397                format!("{buf:?}")
398            }
399        };
400        trace!("Reading packet with bytes: {buf_string}");
401    };
402
403    Ok(Some(buf))
404}
405
406#[cfg(test)]
407mod tests {
408    use std::io::Cursor;
409
410    use azalea_buf::AzBuf as _;
411
412    use crate::{packets::game::ClientboundGamePacket, read::deserialize_packet};
413
414    #[test]
415    fn fuzzed_1() {
416        // oom: checks for unbounded TagMap
417        let _ = deserialize_packet::<ClientboundGamePacket>(&mut Cursor::new(
418            [132, 1, 255, 255, 255, 255, 255].as_slice(),
419        ));
420    }
421    #[test]
422    fn fuzzed_2() {
423        // oom: also checks for unbounded TagMap
424        let _ = deserialize_packet::<ClientboundGamePacket>(&mut Cursor::new(
425            [132, 1, 75, 0, 255, 255, 255, 255, 24, 0].as_slice(),
426        ));
427    }
428    #[test]
429    fn fuzzed_3() {
430        // panic: integer overflow in HolderSet::azalea_read
431        let _ = deserialize_packet::<ClientboundGamePacket>(&mut Cursor::new(
432            [
433                94, 44, 157, 38, 61, 37, 37, 37, 37, 37, 37, 65, 128, 128, 1, 1, 255, 252, 128,
434                128, 128, 128, 128, 128, 128, 40, 0,
435            ]
436            .as_slice(),
437        ));
438    }
439    #[test]
440    fn fuzzed_4() {
441        // memory leak in DataComponentPatch
442        let _ = deserialize_packet::<ClientboundGamePacket>(&mut Cursor::new(
443            [94, 94, 70, 52, 0, 6, 0].as_slice(),
444        ));
445    }
446    #[test]
447    fn fuzzed_5() {
448        // also a memory leak in DataComponentPatch
449        let _ = deserialize_packet::<ClientboundGamePacket>(&mut Cursor::new(
450            [94, 94, 70, 52, 0, 6, 0, 6, 0].as_slice(),
451        ));
452    }
453    #[test]
454    fn fuzzed_6() {
455        // memory leak in simdnbt
456        let _ = simdnbt::owned::Nbt::azalea_read(&mut Cursor::new([10, 10, 0, 0, 0].as_slice()));
457    }
458}