Skip to main content

minecraft_pinger/common/
protocol.rs

1use crate::error::PingError;
2use bytes::{Buf, BufMut, Bytes, BytesMut};
3
4pub const NEXT_BYTE_MASK: u32 = 0xFFFFFF80;
5pub const CONTINUATION_BIT: u8 = 0x80;
6pub const DATA_MASK: u8 = 0x7F;
7
8pub(crate) fn write_var_int(buffer: &mut BytesMut, x: i32) {
9    let mut ux = x as u32; // cast to u32 for logical shift
10    while (ux & NEXT_BYTE_MASK) != 0 {
11        buffer.put_u8((ux as u8 & DATA_MASK) | CONTINUATION_BIT); // (ux as u8 & DATA_MASK) get the 7 bits of data, | CONTINUATION_BIT set the continuation to 1
12        ux >>= 7; // remove the 7 bits
13    }
14    buffer.put_u8(ux as u8); // put the last, often 0x000..
15}
16
17pub(crate) fn write_string(buffer: &mut BytesMut, string: &str) {
18    write_var_int(buffer, string.len() as i32);
19    buffer.put_slice(string.as_bytes());
20}
21
22macro_rules! decode_var_int {
23    ($read_byte:expr, $err_msg:literal) => {{
24        let mut result = 0i32;
25        let mut shift = 0;
26        loop {
27            let byte =
28                $read_byte.map_err(|e| PingError::ReadPacket(format!("{} {e}", $err_msg)))?;
29
30            result |= ((byte & DATA_MASK) as i32) << shift;
31            if byte & CONTINUATION_BIT == 0 {
32                break Ok(result);
33            }
34            shift += 7;
35
36            // Security
37            if shift >= 35 {
38                break Err(PingError::ReadPacket("Shift security".to_string()));
39            }
40        }
41    }};
42}
43
44pub(crate) fn read_var_int(buf: &mut Bytes) -> Result<i32, PingError> {
45    decode_var_int!(buf.try_get_u8(), "Read varint length (v1)")
46}
47
48pub(crate) async fn read_var_int_stream<R: tokio::io::AsyncReadExt + Unpin>(
49    stream: &mut R,
50) -> Result<i32, PingError> {
51    decode_var_int!(stream.read_u8().await, "Read varint length (v2)")
52}
53
54pub(crate) fn read_string(buf: &mut Bytes) -> Result<String, PingError> {
55    let len = read_var_int(buf)? as usize;
56    if buf.remaining() < len {
57        return Err(PingError::ReadPacket("Remaining < len".to_string()));
58    }
59    let bytes = buf.split_to(len);
60    String::from_utf8(bytes.into())
61        .map_err(|_| PingError::ReadPacket("Parse to string error".to_string()))
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67    use bytes::BytesMut;
68
69    #[test]
70    fn test_var_int_serialization() {
71        let mut buf = BytesMut::new();
72        write_var_int(&mut buf, 255);
73        let mut data = buf.freeze();
74        assert_eq!(read_var_int(&mut data).unwrap(), 255);
75    }
76
77    #[test]
78    fn test_var_int_zero() {
79        let mut buf = BytesMut::new();
80        write_var_int(&mut buf, 0);
81        let mut data = buf.freeze();
82        assert_eq!(read_var_int(&mut data).unwrap(), 0);
83    }
84
85    #[test]
86    fn test_var_int_large() {
87        let mut buf = BytesMut::new();
88        write_var_int(&mut buf, 2147483647);
89        let mut data = buf.freeze();
90        assert_eq!(read_var_int(&mut data).unwrap(), 2147483647);
91    }
92
93    #[test]
94    fn test_var_int_negative() {
95        let mut buf = BytesMut::new();
96        write_var_int(&mut buf, -1);
97        let mut data = buf.freeze();
98        assert_eq!(read_var_int(&mut data).unwrap(), -1);
99    }
100
101    #[test]
102    fn test_var_int_multiple() {
103        let mut buf = BytesMut::new();
104        write_var_int(&mut buf, 1);
105        write_var_int(&mut buf, 2);
106        write_var_int(&mut buf, 3);
107        let mut data = buf.freeze();
108        assert_eq!(read_var_int(&mut data).unwrap(), 1);
109        assert_eq!(read_var_int(&mut data).unwrap(), 2);
110        assert_eq!(read_var_int(&mut data).unwrap(), 3);
111    }
112
113    #[test]
114    fn test_write_string() {
115        let mut buf = BytesMut::new();
116        write_string(&mut buf, "hello");
117        let mut data = buf.freeze();
118        assert_eq!(read_string(&mut data).unwrap(), "hello");
119    }
120
121    #[test]
122    fn test_write_empty_string() {
123        let mut buf = BytesMut::new();
124        write_string(&mut buf, "");
125        let mut data = buf.freeze();
126        assert_eq!(read_string(&mut data).unwrap(), "");
127    }
128
129    #[test]
130    fn test_write_unicode_string() {
131        let mut buf = BytesMut::new();
132        write_string(&mut buf, "héllo 🌍");
133        let mut data = buf.freeze();
134        assert_eq!(read_string(&mut data).unwrap(), "héllo 🌍");
135    }
136
137    #[test]
138    fn test_read_var_int_shift_security() {
139        let mut buf = BytesMut::new();
140        for _ in 0..6 {
141            buf.put_u8(0xFF);
142        }
143        buf.put_u8(0x01);
144        let mut data = buf.freeze();
145        let result = read_var_int(&mut data);
146        assert!(result.is_err());
147    }
148
149    #[test]
150    fn test_read_string_insufficient_data() {
151        let mut buf = BytesMut::new();
152        write_var_int(&mut buf, 10);
153        buf.put_slice(b"hi");
154        let mut data = buf.freeze();
155        let result = read_string(&mut data);
156        assert!(result.is_err());
157    }
158}