Skip to main content

nbt_rust/
protocol_adapter.rs

1use std::io::Cursor;
2
3use serde::de::DeserializeOwned;
4use serde::Serialize;
5
6use crate::config::{NbtReadConfig, ParseMode};
7use crate::encoding::{BigEndian, LittleEndian, NetworkLittleEndian};
8use crate::error::{Error, Result};
9use crate::headless::{
10    read_headless_prefixed_with_config, read_headless_with_config, write_headless,
11    write_headless_prefixed,
12};
13use crate::limits::NbtLimits;
14use crate::root::{read_tag_with_config, write_tag, RootTag};
15use crate::serde_api::{from_root_tag, from_tag, to_root_tag, to_tag};
16use crate::tag::{Tag, TagType};
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum ProtocolNbtEncoding {
20    Network,
21    LittleEndian,
22    BigEndian,
23}
24
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct ProtocolNbtAdapter {
27    pub encoding: ProtocolNbtEncoding,
28    pub read_config: NbtReadConfig,
29}
30
31impl Default for ProtocolNbtAdapter {
32    fn default() -> Self {
33        Self::network()
34    }
35}
36
37impl ProtocolNbtAdapter {
38    pub fn network() -> Self {
39        Self {
40            encoding: ProtocolNbtEncoding::Network,
41            read_config: NbtReadConfig::default(),
42        }
43    }
44
45    pub fn little_endian() -> Self {
46        Self {
47            encoding: ProtocolNbtEncoding::LittleEndian,
48            read_config: NbtReadConfig::default(),
49        }
50    }
51
52    pub fn big_endian() -> Self {
53        Self {
54            encoding: ProtocolNbtEncoding::BigEndian,
55            read_config: NbtReadConfig::default(),
56        }
57    }
58
59    pub fn with_config(mut self, read_config: NbtReadConfig) -> Self {
60        self.read_config = read_config;
61        self
62    }
63
64    pub fn with_limits(mut self, limits: NbtLimits) -> Self {
65        self.read_config = self.read_config.with_limits(limits);
66        self
67    }
68
69    pub fn with_parse_mode(mut self, parse_mode: ParseMode) -> Self {
70        self.read_config = self.read_config.with_parse_mode(parse_mode);
71        self
72    }
73
74    pub fn decode_headless_tag(&self, tag_type: TagType, bytes: &[u8]) -> Result<Tag> {
75        let mut cursor = Cursor::new(bytes);
76        let tag = match self.encoding {
77            ProtocolNbtEncoding::Network => read_headless_with_config::<NetworkLittleEndian, _>(
78                &mut cursor,
79                tag_type,
80                &self.read_config,
81            ),
82            ProtocolNbtEncoding::LittleEndian => read_headless_with_config::<LittleEndian, _>(
83                &mut cursor,
84                tag_type,
85                &self.read_config,
86            ),
87            ProtocolNbtEncoding::BigEndian => {
88                read_headless_with_config::<BigEndian, _>(&mut cursor, tag_type, &self.read_config)
89            }
90        }?;
91        ensure_fully_consumed(bytes.len(), cursor.position() as usize)?;
92        Ok(tag)
93    }
94
95    pub fn decode_headless<T: DeserializeOwned>(
96        &self,
97        tag_type: TagType,
98        bytes: &[u8],
99    ) -> Result<T> {
100        let tag = self.decode_headless_tag(tag_type, bytes)?;
101        from_tag(&tag)
102    }
103
104    pub fn encode_headless_tag(&self, tag: &Tag) -> Result<Vec<u8>> {
105        let mut out = Vec::new();
106        match self.encoding {
107            ProtocolNbtEncoding::Network => {
108                write_headless::<NetworkLittleEndian, _>(&mut out, tag)?
109            }
110            ProtocolNbtEncoding::LittleEndian => write_headless::<LittleEndian, _>(&mut out, tag)?,
111            ProtocolNbtEncoding::BigEndian => write_headless::<BigEndian, _>(&mut out, tag)?,
112        }
113        Ok(out)
114    }
115
116    pub fn encode_headless<T: Serialize>(&self, value: &T) -> Result<(TagType, Vec<u8>)> {
117        let tag = to_tag(value)?;
118        let tag_type = tag.tag_type();
119        let bytes = self.encode_headless_tag(&tag)?;
120        Ok((tag_type, bytes))
121    }
122
123    pub fn decode_prefixed_tag(&self, bytes: &[u8]) -> Result<Tag> {
124        let mut cursor = Cursor::new(bytes);
125        let tag = match self.encoding {
126            ProtocolNbtEncoding::Network => read_headless_prefixed_with_config::<
127                NetworkLittleEndian,
128                _,
129            >(&mut cursor, &self.read_config),
130            ProtocolNbtEncoding::LittleEndian => read_headless_prefixed_with_config::<
131                LittleEndian,
132                _,
133            >(&mut cursor, &self.read_config),
134            ProtocolNbtEncoding::BigEndian => {
135                read_headless_prefixed_with_config::<BigEndian, _>(&mut cursor, &self.read_config)
136            }
137        }?;
138        ensure_fully_consumed(bytes.len(), cursor.position() as usize)?;
139        Ok(tag)
140    }
141
142    pub fn decode_prefixed<T: DeserializeOwned>(&self, bytes: &[u8]) -> Result<T> {
143        let tag = self.decode_prefixed_tag(bytes)?;
144        from_tag(&tag)
145    }
146
147    pub fn encode_prefixed_tag(&self, tag: &Tag) -> Result<Vec<u8>> {
148        let mut out = Vec::new();
149        match self.encoding {
150            ProtocolNbtEncoding::Network => {
151                write_headless_prefixed::<NetworkLittleEndian, _>(&mut out, tag)?
152            }
153            ProtocolNbtEncoding::LittleEndian => {
154                write_headless_prefixed::<LittleEndian, _>(&mut out, tag)?
155            }
156            ProtocolNbtEncoding::BigEndian => {
157                write_headless_prefixed::<BigEndian, _>(&mut out, tag)?
158            }
159        }
160        Ok(out)
161    }
162
163    pub fn encode_prefixed<T: Serialize>(&self, value: &T) -> Result<Vec<u8>> {
164        let tag = to_tag(value)?;
165        self.encode_prefixed_tag(&tag)
166    }
167
168    pub fn decode_root_tag(&self, bytes: &[u8]) -> Result<RootTag> {
169        let mut cursor = Cursor::new(bytes);
170        let root = match self.encoding {
171            ProtocolNbtEncoding::Network => {
172                read_tag_with_config::<NetworkLittleEndian, _>(&mut cursor, &self.read_config)
173            }
174            ProtocolNbtEncoding::LittleEndian => {
175                read_tag_with_config::<LittleEndian, _>(&mut cursor, &self.read_config)
176            }
177            ProtocolNbtEncoding::BigEndian => {
178                read_tag_with_config::<BigEndian, _>(&mut cursor, &self.read_config)
179            }
180        }?;
181        ensure_fully_consumed(bytes.len(), cursor.position() as usize)?;
182        Ok(root)
183    }
184
185    pub fn decode_root<T: DeserializeOwned>(&self, bytes: &[u8]) -> Result<T> {
186        let root = self.decode_root_tag(bytes)?;
187        from_root_tag(&root)
188    }
189
190    pub fn decode_root_named<T: DeserializeOwned>(&self, bytes: &[u8]) -> Result<(String, T)> {
191        let root = self.decode_root_tag(bytes)?;
192        let value = from_root_tag(&root)?;
193        Ok((root.name, value))
194    }
195
196    pub fn encode_root_tag(&self, root: &RootTag) -> Result<Vec<u8>> {
197        let mut out = Vec::new();
198        match self.encoding {
199            ProtocolNbtEncoding::Network => write_tag::<NetworkLittleEndian, _>(&mut out, root)?,
200            ProtocolNbtEncoding::LittleEndian => write_tag::<LittleEndian, _>(&mut out, root)?,
201            ProtocolNbtEncoding::BigEndian => write_tag::<BigEndian, _>(&mut out, root)?,
202        }
203        Ok(out)
204    }
205
206    pub fn encode_root<T: Serialize>(
207        &self,
208        root_name: impl Into<String>,
209        value: &T,
210    ) -> Result<Vec<u8>> {
211        let root = to_root_tag(root_name, value)?;
212        self.encode_root_tag(&root)
213    }
214}
215
216fn ensure_fully_consumed(total: usize, consumed: usize) -> Result<()> {
217    if consumed == total {
218        return Ok(());
219    }
220    Err(Error::TrailingPayloadBytes {
221        unread: total - consumed,
222    })
223}
224
225#[cfg(test)]
226mod tests {
227    use serde::{Deserialize, Serialize};
228
229    use crate::config::NbtReadConfig;
230    use crate::limits::NbtLimits;
231    use crate::tag::ListTag;
232
233    use super::*;
234
235    #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
236    struct PlayerState {
237        username: String,
238        hp: i32,
239        scores: Vec<i32>,
240        flags: Vec<u8>,
241    }
242
243    fn sample() -> PlayerState {
244        PlayerState {
245            username: "Alex".to_string(),
246            hp: 20,
247            scores: vec![10, 20, 30],
248            flags: vec![1, 0, 1, 1],
249        }
250    }
251
252    #[test]
253    fn network_headless_typed_roundtrip() {
254        let adapter = ProtocolNbtAdapter::network();
255        let input = sample();
256        let (tag_type, bytes) = adapter.encode_headless(&input).unwrap();
257        let output: PlayerState = adapter.decode_headless(tag_type, &bytes).unwrap();
258        assert_eq!(output, input);
259    }
260
261    #[test]
262    fn network_prefixed_typed_roundtrip() {
263        let adapter = ProtocolNbtAdapter::network();
264        let input = sample();
265        let bytes = adapter.encode_prefixed(&input).unwrap();
266        let output: PlayerState = adapter.decode_prefixed(&bytes).unwrap();
267        assert_eq!(output, input);
268    }
269
270    #[test]
271    fn little_endian_root_named_roundtrip() {
272        let adapter = ProtocolNbtAdapter::little_endian();
273        let input = sample();
274        let bytes = adapter.encode_root("PlayerState", &input).unwrap();
275        let (root_name, output): (String, PlayerState) = adapter.decode_root_named(&bytes).unwrap();
276        assert_eq!(root_name, "PlayerState");
277        assert_eq!(output, input);
278    }
279
280    #[test]
281    fn strict_vs_compatible_list_header_behavior() {
282        let payload = vec![0x00, 0x00, 0x00, 0x00, 0x01];
283
284        let strict = ProtocolNbtAdapter::big_endian();
285        let strict_err = strict
286            .decode_headless_tag(TagType::List, &payload)
287            .unwrap_err();
288        assert!(matches!(
289            strict_err.innermost(),
290            Error::InvalidListHeader { .. }
291        ));
292
293        let compat_cfg = NbtReadConfig::compatible(NbtLimits::default());
294        let compat = ProtocolNbtAdapter::big_endian().with_config(compat_cfg);
295        let compat_tag = compat.decode_headless_tag(TagType::List, &payload).unwrap();
296        assert_eq!(compat_tag, Tag::List(ListTag::empty(TagType::End)));
297    }
298
299    #[test]
300    fn prefixed_decode_rejects_trailing_bytes() {
301        let adapter = ProtocolNbtAdapter::network();
302        let bytes = {
303            let mut out = adapter.encode_prefixed(&sample()).unwrap();
304            out.extend_from_slice(&[0xAA, 0xBB]);
305            out
306        };
307        let err = adapter.decode_prefixed_tag(&bytes).unwrap_err();
308        assert!(matches!(
309            err.innermost(),
310            Error::TrailingPayloadBytes { unread: 2 }
311        ));
312    }
313}