ghast_protocol/
codec.rs

1//! Low-level reading and writing of protocol primitives.
2//!
3//! [`Reader`] is a bounds-checked cursor over a [`Bytes`] buffer: every read returns
4//! `Result`, and variable-length data ([`Reader::read_bytes`], strings) is returned as
5//! zero-copy slices of the underlying buffer. [`McWrite`] adds the protocol's primitive
6//! encodings to [`BytesMut`].
7//!
8//! The [`Decode`] / [`Encode`] traits are implemented by every packet and shared type.
9
10use bytes::{BufMut, Bytes, BytesMut};
11use bytestring::ByteString;
12use uuid::Uuid;
13
14use crate::error::{DecodeError, EncodeError};
15
16/// Result alias for decoding.
17pub type DecodeResult<T> = Result<T, DecodeError>;
18/// Result alias for encoding.
19pub type EncodeResult = Result<(), EncodeError>;
20
21/// A value that can be decoded from protocol bytes.
22pub trait Decode: Sized {
23    /// Decodes one value, advancing the reader past it.
24    fn decode(r: &mut Reader<'_>) -> DecodeResult<Self>;
25}
26
27/// A value that can be encoded to protocol bytes.
28pub trait Encode {
29    /// Appends the encoding of `self` to `w`.
30    fn encode(&self, w: &mut BytesMut) -> EncodeResult;
31}
32
33/// Decodes a complete value from `buf`, failing if any bytes are left over.
34pub fn decode_exact<T: Decode>(buf: &Bytes) -> DecodeResult<T> {
35    let mut r = Reader::new(buf);
36    let v = T::decode(&mut r)?;
37    r.finish()?;
38    Ok(v)
39}
40
41/// Encodes a value into a fresh buffer.
42pub fn encode_to_bytes<T: Encode + ?Sized>(v: &T) -> Result<BytesMut, EncodeError> {
43    let mut w = BytesMut::new();
44    v.encode(&mut w)?;
45    Ok(w)
46}
47
48/// Bounds-checked read cursor over a [`Bytes`] buffer.
49#[derive(Debug, Clone)]
50pub struct Reader<'a> {
51    buf: &'a Bytes,
52    pos: usize,
53}
54
55impl<'a> Reader<'a> {
56    /// Creates a reader at the start of `buf`.
57    #[inline]
58    pub fn new(buf: &'a Bytes) -> Self {
59        Reader { buf, pos: 0 }
60    }
61
62    /// Number of unread bytes.
63    #[inline]
64    pub fn remaining(&self) -> usize {
65        self.buf.len() - self.pos
66    }
67
68    /// Current offset from the start of the buffer.
69    #[inline]
70    pub fn position(&self) -> usize {
71        self.pos
72    }
73
74    /// The unread bytes, without consuming them.
75    #[inline]
76    pub fn rest_slice(&self) -> &'a [u8] {
77        &self.buf[self.pos..]
78    }
79
80    /// Fails with [`DecodeError::TrailingBytes`] unless the input is fully consumed.
81    #[inline]
82    pub fn finish(&self) -> DecodeResult<()> {
83        match self.remaining() {
84            0 => Ok(()),
85            n => Err(DecodeError::TrailingBytes(n)),
86        }
87    }
88
89    #[inline]
90    fn eof(&self, n: usize) -> DecodeError {
91        DecodeError::UnexpectedEof {
92            needed: n - self.remaining(),
93        }
94    }
95
96    /// Consumes `n` bytes and returns them as a borrowed slice.
97    #[inline]
98    pub fn take(&mut self, n: usize) -> DecodeResult<&'a [u8]> {
99        if n > self.remaining() {
100            return Err(self.eof(n));
101        }
102        let buf: &'a Bytes = self.buf;
103        let s = &buf[self.pos..self.pos + n];
104        self.pos += n;
105        Ok(s)
106    }
107
108    /// Advances past `n` bytes.
109    #[inline]
110    pub fn skip(&mut self, n: usize) -> DecodeResult<()> {
111        if n > self.remaining() {
112            return Err(self.eof(n));
113        }
114        self.pos += n;
115        Ok(())
116    }
117
118    #[inline]
119    fn take_array<const N: usize>(&mut self) -> DecodeResult<[u8; N]> {
120        let s = self.take(N)?;
121        let mut a = [0u8; N];
122        a.copy_from_slice(s);
123        Ok(a)
124    }
125
126    /// Consumes `n` bytes and returns them as a zero-copy [`Bytes`] slice.
127    #[inline]
128    pub fn read_bytes(&mut self, n: usize) -> DecodeResult<Bytes> {
129        if n > self.remaining() {
130            return Err(self.eof(n));
131        }
132        let b = self.buf.slice(self.pos..self.pos + n);
133        self.pos += n;
134        Ok(b)
135    }
136
137    /// Zero-copy slice of everything read since offset `start` (from [`Reader::position`]).
138    #[inline]
139    pub fn bytes_since(&self, start: usize) -> Bytes {
140        self.buf.slice(start.min(self.pos)..self.pos)
141    }
142
143    /// Consumes and returns everything that is left (protodef `restBuffer`).
144    #[inline]
145    pub fn read_rest(&mut self) -> Bytes {
146        let b = self.buf.slice(self.pos..);
147        self.pos = self.buf.len();
148        b
149    }
150
151    /// Returns the next byte without consuming it.
152    #[inline]
153    pub fn peek_u8(&self) -> DecodeResult<u8> {
154        self.buf
155            .get(self.pos)
156            .copied()
157            .ok_or(DecodeError::UnexpectedEof { needed: 1 })
158    }
159
160    /// Reads an unsigned byte.
161    #[inline]
162    pub fn read_u8(&mut self) -> DecodeResult<u8> {
163        let b = self.peek_u8()?;
164        self.pos += 1;
165        Ok(b)
166    }
167
168    /// Reads a signed byte.
169    #[inline]
170    pub fn read_i8(&mut self) -> DecodeResult<i8> {
171        Ok(self.read_u8()? as i8)
172    }
173
174    /// Reads a boolean. Like vanilla, any non-zero byte is `true`.
175    #[inline]
176    pub fn read_bool(&mut self) -> DecodeResult<bool> {
177        Ok(self.read_u8()? != 0)
178    }
179
180    /// Reads a big-endian `u16`.
181    #[inline]
182    pub fn read_u16(&mut self) -> DecodeResult<u16> {
183        Ok(u16::from_be_bytes(self.take_array()?))
184    }
185
186    /// Reads a big-endian `i16`.
187    #[inline]
188    pub fn read_i16(&mut self) -> DecodeResult<i16> {
189        Ok(i16::from_be_bytes(self.take_array()?))
190    }
191
192    /// Reads a big-endian `i32`.
193    #[inline]
194    pub fn read_i32(&mut self) -> DecodeResult<i32> {
195        Ok(i32::from_be_bytes(self.take_array()?))
196    }
197
198    /// Reads a big-endian `i64`.
199    #[inline]
200    pub fn read_i64(&mut self) -> DecodeResult<i64> {
201        Ok(i64::from_be_bytes(self.take_array()?))
202    }
203
204    /// Reads a big-endian `u64`.
205    #[inline]
206    pub fn read_u64(&mut self) -> DecodeResult<u64> {
207        Ok(u64::from_be_bytes(self.take_array()?))
208    }
209
210    /// Reads a big-endian IEEE-754 `f32`.
211    #[inline]
212    pub fn read_f32(&mut self) -> DecodeResult<f32> {
213        Ok(f32::from_be_bytes(self.take_array()?))
214    }
215
216    /// Reads a big-endian IEEE-754 `f64`.
217    #[inline]
218    pub fn read_f64(&mut self) -> DecodeResult<f64> {
219        Ok(f64::from_be_bytes(self.take_array()?))
220    }
221
222    /// Reads a UUID (two big-endian 64-bit halves).
223    #[inline]
224    pub fn read_uuid(&mut self) -> DecodeResult<Uuid> {
225        Ok(Uuid::from_u128(u128::from_be_bytes(self.take_array()?)))
226    }
227
228    /// Reads a VarInt (at most 5 bytes).
229    ///
230    /// Like vanilla, the unused high bits of a 5th byte are ignored; a 6th byte is an
231    /// error ([`DecodeError::VarIntTooLong`]).
232    #[inline]
233    pub fn read_varint(&mut self) -> DecodeResult<i32> {
234        let rest = &self.buf[self.pos..];
235        if let Some(&b) = rest.first()
236            && b < 0x80
237        {
238            self.pos += 1;
239            return Ok(i32::from(b));
240        }
241        let (v, n) = varint_slow(rest)?;
242        self.pos += n;
243        Ok(v)
244    }
245
246    /// Reads a VarLong (at most 10 bytes).
247    #[inline]
248    pub fn read_varlong(&mut self) -> DecodeResult<i64> {
249        let rest = &self.buf[self.pos..];
250        let mut result: u64 = 0;
251        for i in 0..10 {
252            let Some(&b) = rest.get(i) else {
253                return Err(DecodeError::UnexpectedEof { needed: 1 });
254            };
255            result |= u64::from(b & 0x7f) << (7 * i);
256            if b & 0x80 == 0 {
257                self.pos += i + 1;
258                return Ok(result as i64);
259            }
260        }
261        Err(DecodeError::VarIntTooLong)
262    }
263
264    /// Reads a non-negative length of type `i32` and checks it against `limit`.
265    #[inline]
266    fn check_len(len: i64, limit: usize) -> DecodeResult<usize> {
267        if len < 0 {
268            return Err(DecodeError::NegativeLength(len));
269        }
270        let len = len as usize;
271        if len > limit {
272            return Err(DecodeError::LengthTooLarge { len, limit });
273        }
274        Ok(len)
275    }
276
277    /// Reads a VarInt length prefix and checks it is non-negative and at most `limit`.
278    #[inline]
279    pub fn read_len_varint(&mut self, limit: usize) -> DecodeResult<usize> {
280        let n = self.read_varint()?;
281        Self::check_len(i64::from(n), limit)
282    }
283
284    /// Reads an array count of the given prefix width and validates it.
285    ///
286    /// `min_elem_size` is the smallest number of bytes one element can occupy; when it is
287    /// non-zero the count is bounded by the remaining input, so a hostile count can never
288    /// cause a large allocation or a long loop.
289    #[inline]
290    pub fn check_count(&self, count: i64, min_elem_size: usize) -> DecodeResult<usize> {
291        let limit = self
292            .remaining()
293            .checked_div(min_elem_size)
294            .unwrap_or(ZERO_SIZE_ELEMENT_LIMIT);
295        Self::check_len(count, limit)
296    }
297
298    /// Reads a VarInt-prefixed byte array as a zero-copy slice.
299    #[inline]
300    pub fn read_byte_array(&mut self) -> DecodeResult<Bytes> {
301        let n = self.read_len_varint(self.remaining())?;
302        self.read_bytes(n)
303    }
304
305    /// Reads a VarInt-prefixed UTF-8 string as a zero-copy [`ByteString`].
306    #[inline]
307    pub fn read_string(&mut self) -> DecodeResult<ByteString> {
308        let n = self.read_len_varint(self.remaining())?;
309        let b = self.read_bytes(n)?;
310        ByteString::try_from(b).map_err(|_| DecodeError::InvalidUtf8)
311    }
312}
313
314/// Largest count accepted for arrays whose elements can be zero bytes long.
315pub const ZERO_SIZE_ELEMENT_LIMIT: usize = 1 << 16;
316
317#[inline(never)]
318fn varint_slow(rest: &[u8]) -> DecodeResult<(i32, usize)> {
319    let mut result: u32 = 0;
320    for i in 0..5 {
321        let Some(&b) = rest.get(i) else {
322            return Err(DecodeError::UnexpectedEof { needed: 1 });
323        };
324        result |= u32::from(b & 0x7f) << (7 * i);
325        if b & 0x80 == 0 {
326            return Ok((result as i32, i + 1));
327        }
328    }
329    Err(DecodeError::VarIntTooLong)
330}
331
332/// Decodes a VarInt from the start of a plain slice. Returns the value and its length.
333///
334/// Returns `Ok(None)` when the slice ends before the VarInt does.
335#[inline]
336pub fn read_varint_prefix(buf: &[u8]) -> DecodeResult<Option<(i32, usize)>> {
337    match varint_slow(buf) {
338        Ok(v) => Ok(Some(v)),
339        Err(DecodeError::UnexpectedEof { .. }) => Ok(None),
340        Err(e) => Err(e),
341    }
342}
343
344/// Number of bytes the VarInt encoding of `v` takes (1..=5).
345#[inline]
346pub const fn varint_len(v: i32) -> usize {
347    let v = v as u32;
348    match v {
349        0..=0x7f => 1,
350        0x80..=0x3fff => 2,
351        0x4000..=0x1f_ffff => 3,
352        0x20_0000..=0x0fff_ffff => 4,
353        _ => 5,
354    }
355}
356
357/// Encodes `v` as a VarInt into a fixed buffer; returns the buffer and the used length.
358#[inline]
359pub const fn varint_bytes(v: i32) -> ([u8; 5], usize) {
360    let mut v = v as u32;
361    let mut out = [0u8; 5];
362    let mut i = 0;
363    loop {
364        if v < 0x80 {
365            out[i] = v as u8;
366            return (out, i + 1);
367        }
368        out[i] = (v as u8 & 0x7f) | 0x80;
369        v >>= 7;
370        i += 1;
371    }
372}
373
374/// Protocol primitive encodings on top of [`BytesMut`].
375pub trait McWrite {
376    /// Writes a VarInt.
377    fn put_varint(&mut self, v: i32);
378    /// Writes a VarLong.
379    fn put_varlong(&mut self, v: i64);
380    /// Writes a boolean as one byte (0 or 1).
381    fn put_bool(&mut self, v: bool);
382    /// Writes a UUID as 16 big-endian bytes.
383    fn put_uuid(&mut self, v: &Uuid);
384    /// Writes a VarInt length prefix.
385    fn put_len_varint(&mut self, len: usize) -> EncodeResult;
386    /// Writes a VarInt-prefixed byte array.
387    fn put_byte_array(&mut self, b: &[u8]) -> EncodeResult;
388    /// Writes a VarInt-prefixed UTF-8 string.
389    fn put_string(&mut self, s: &str) -> EncodeResult;
390}
391
392impl McWrite for BytesMut {
393    #[inline]
394    fn put_varint(&mut self, v: i32) {
395        if (0..0x80).contains(&v) {
396            self.put_u8(v as u8);
397        } else {
398            let (b, n) = varint_bytes(v);
399            self.put_slice(&b[..n]);
400        }
401    }
402
403    #[inline]
404    fn put_varlong(&mut self, v: i64) {
405        let mut v = v as u64;
406        loop {
407            if v < 0x80 {
408                self.put_u8(v as u8);
409                return;
410            }
411            self.put_u8((v as u8 & 0x7f) | 0x80);
412            v >>= 7;
413        }
414    }
415
416    #[inline]
417    fn put_bool(&mut self, v: bool) {
418        self.put_u8(u8::from(v));
419    }
420
421    #[inline]
422    fn put_uuid(&mut self, v: &Uuid) {
423        self.put_u128(v.as_u128());
424    }
425
426    #[inline]
427    fn put_len_varint(&mut self, len: usize) -> EncodeResult {
428        let v = i32::try_from(len).map_err(|_| EncodeError::LengthOverflow {
429            len,
430            max: i32::MAX as usize,
431        })?;
432        self.put_varint(v);
433        Ok(())
434    }
435
436    #[inline]
437    fn put_byte_array(&mut self, b: &[u8]) -> EncodeResult {
438        self.put_len_varint(b.len())?;
439        self.put_slice(b);
440        Ok(())
441    }
442
443    #[inline]
444    fn put_string(&mut self, s: &str) -> EncodeResult {
445        self.put_byte_array(s.as_bytes())
446    }
447}
448
449/// Upper bound on the capacity reserved up front for a decoded array; larger arrays
450/// grow as elements are actually decoded, so a hostile count cannot force a big
451/// allocation.
452pub const PREALLOC_LIMIT: usize = 4096;
453
454fn overflow(len: usize, max: usize) -> EncodeError {
455    EncodeError::LengthOverflow { len, max }
456}
457
458/// Converts a collection length into an `i16` count prefix.
459#[inline]
460pub fn count_i16(len: usize) -> Result<i16, EncodeError> {
461    i16::try_from(len).map_err(|_| overflow(len, i16::MAX as usize))
462}
463
464/// Converts a collection length into an `i32` count prefix.
465#[inline]
466pub fn count_i32(len: usize) -> Result<i32, EncodeError> {
467    i32::try_from(len).map_err(|_| overflow(len, i32::MAX as usize))
468}
469
470/// Converts a collection length into a `u8` count prefix.
471#[inline]
472pub fn count_u8(len: usize) -> Result<u8, EncodeError> {
473    u8::try_from(len).map_err(|_| overflow(len, u8::MAX as usize))
474}
475
476#[cfg(test)]
477mod tests {
478    use super::*;
479
480    fn rd(bytes: &[u8]) -> DecodeResult<i32> {
481        let b = Bytes::copy_from_slice(bytes);
482        let mut r = Reader::new(&b);
483        let v = r.read_varint()?;
484        r.finish()?;
485        Ok(v)
486    }
487
488    #[test]
489    fn varint_known_answers() {
490        let cases: &[(i32, &[u8])] = &[
491            (0, &[0x00]),
492            (1, &[0x01]),
493            (2, &[0x02]),
494            (127, &[0x7f]),
495            (128, &[0x80, 0x01]),
496            (255, &[0xff, 0x01]),
497            (25565, &[0xdd, 0xc7, 0x01]),
498            (2097151, &[0xff, 0xff, 0x7f]),
499            (2147483647, &[0xff, 0xff, 0xff, 0xff, 0x07]),
500            (-1, &[0xff, 0xff, 0xff, 0xff, 0x0f]),
501            (-2147483648, &[0x80, 0x80, 0x80, 0x80, 0x08]),
502        ];
503        for &(v, bytes) in cases {
504            let mut w = BytesMut::new();
505            w.put_varint(v);
506            assert_eq!(&w[..], bytes, "encode {v}");
507            assert_eq!(varint_len(v), bytes.len());
508            assert_eq!(rd(bytes), Ok(v), "decode {bytes:02x?}");
509        }
510    }
511
512    #[test]
513    fn varint_rejects_too_long_and_eof() {
514        assert_eq!(
515            rd(&[0x80, 0x80, 0x80, 0x80, 0x80, 0x01]),
516            Err(DecodeError::VarIntTooLong)
517        );
518        assert!(matches!(
519            rd(&[0x80]),
520            Err(DecodeError::UnexpectedEof { .. })
521        ));
522        assert!(matches!(rd(&[]), Err(DecodeError::UnexpectedEof { .. })));
523        // Non-canonical (padded) encodings are accepted, like vanilla.
524        assert_eq!(rd(&[0x80, 0x00]), Ok(0));
525        assert_eq!(read_varint_prefix(&[0xff, 0xff]), Ok(None));
526    }
527
528    #[test]
529    fn varlong_known_answers() {
530        let cases: &[(i64, &[u8])] = &[
531            (0, &[0x00]),
532            (127, &[0x7f]),
533            (128, &[0x80, 0x01]),
534            (2147483647, &[0xff, 0xff, 0xff, 0xff, 0x07]),
535            (
536                9223372036854775807,
537                &[0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x7f],
538            ),
539            (
540                -1,
541                &[0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01],
542            ),
543            (
544                -9223372036854775808,
545                &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x01],
546            ),
547        ];
548        for &(v, bytes) in cases {
549            let mut w = BytesMut::new();
550            w.put_varlong(v);
551            assert_eq!(&w[..], bytes, "encode {v}");
552            let b = Bytes::copy_from_slice(bytes);
553            let mut r = Reader::new(&b);
554            assert_eq!(r.read_varlong(), Ok(v));
555            assert_eq!(r.remaining(), 0);
556        }
557        let b = Bytes::from_static(&[0x80; 11]);
558        assert_eq!(
559            Reader::new(&b).read_varlong(),
560            Err(DecodeError::VarIntTooLong)
561        );
562    }
563
564    #[test]
565    fn strings_are_bounded_and_validated() {
566        let b = Bytes::from_static(&[0x05, b'h', b'i']);
567        assert!(matches!(
568            Reader::new(&b).read_string(),
569            Err(DecodeError::LengthTooLarge { .. })
570        ));
571        let b = Bytes::from_static(&[0x02, 0xc3, 0x28]);
572        assert_eq!(Reader::new(&b).read_string(), Err(DecodeError::InvalidUtf8));
573        let b = Bytes::from_static(&[0xff, 0xff, 0xff, 0xff, 0x0f]);
574        assert_eq!(
575            Reader::new(&b).read_string(),
576            Err(DecodeError::NegativeLength(-1))
577        );
578        let b = Bytes::from_static(&[0x03, b'a', b'b', b'c', 0x00]);
579        let mut r = Reader::new(&b);
580        assert_eq!(&*r.read_string().unwrap(), "abc");
581        assert_eq!(r.finish(), Err(DecodeError::TrailingBytes(1)));
582    }
583
584    #[test]
585    fn fixed_width_known_answers() {
586        let b = Bytes::from_static(&[
587            0x12, 0x34, 0xff, 0xfe, 0x3f, 0x80, 0x00, 0x00, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55,
588            0x66, 0x77, 0x88, 0x99, 0xaa, 0xbb, 0xcc, 0xdd, 0xee, 0xff,
589        ]);
590        let mut r = Reader::new(&b);
591        assert_eq!(r.read_u16(), Ok(0x1234));
592        assert_eq!(r.read_i16(), Ok(-2));
593        assert_eq!(r.read_f32(), Ok(1.0));
594        assert_eq!(
595            r.read_uuid().unwrap().to_string(),
596            "00112233-4455-6677-8899-aabbccddeeff"
597        );
598        assert_eq!(r.remaining(), 0);
599    }
600}
ghast_protocol/codec.rs

ghast_protocol/
codec.rs
