mcproto_rs/
nbt.rs

1use crate::utils::take;
2use crate::{DeserializeErr, DeserializeResult, Deserialized};
3use alloc::{string::{String, ToString}, borrow::ToOwned, fmt, vec::Vec, vec, format};
4
5#[cfg(all(test, feature = "std"))]
6use crate::protocol::TestRandom;
7use crate::byte_order::{ProtoByteOrder, ByteOrder};
8
9#[derive(Clone, Debug, PartialEq)]
10pub struct NamedTag {
11    pub name: String,
12    pub payload: Tag,
13}
14
15impl NamedTag {
16    pub fn root_compound_tag_from_bytes(data: &[u8]) -> DeserializeResult<NamedTag> {
17        read_nbt_data(data)
18    }
19
20    pub fn is_end(&self) -> bool {
21        match self.payload {
22            Tag::End => true,
23            _ => false,
24        }
25    }
26}
27
28#[cfg(all(test, feature = "std"))]
29impl TestRandom for NamedTag {
30    fn test_gen_random() -> Self {
31        Self {
32            name: "".to_owned(),
33            payload: Tag::Compound(vec![Self {
34                name: String::test_gen_random(),
35                payload: Tag::test_gen_random(),
36            }]),
37        }
38    }
39}
40
41impl fmt::Display for NamedTag {
42    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
43        f.write_fmt(format_args!(
44            "TAG_{}('{}'): ",
45            self.payload.tag_type_name(),
46            self.name
47        ))?;
48        self.payload.write_contents(f)
49    }
50}
51
52#[derive(Clone, Debug, PartialEq)]
53pub enum Tag {
54    Byte(i8),
55    Short(i16),
56    Int(i32),
57    Long(i64),
58    Float(f32),
59    Double(f64),
60    ByteArray(Vec<u8>),
61    String(String),
62    List(Vec<Tag>),
63    Compound(Vec<NamedTag>),
64    IntArray(Vec<i32>),
65    LongArray(Vec<i64>),
66    End,
67}
68
69impl fmt::Display for Tag {
70    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
71        f.write_fmt(format_args!("TAG_{}: ", self.tag_type_name()))?;
72        self.write_contents(f)
73    }
74}
75
76impl Tag {
77    pub fn with_name(self, name: &str) -> NamedTag {
78        NamedTag {
79            name: name.into(),
80            payload: self,
81        }
82    }
83
84    pub fn tag_type_name(&self) -> &str {
85        match self {
86            Tag::Byte(_) => "Byte",
87            Tag::Short(_) => "Short",
88            Tag::Int(_) => "Int",
89            Tag::Long(_) => "Long",
90            Tag::Float(_) => "Float",
91            Tag::Double(_) => "Double",
92            Tag::ByteArray(_) => "Byte_Array",
93            Tag::String(_) => "String",
94            Tag::List(_) => "List",
95            Tag::Compound(_) => "Compound",
96            Tag::IntArray(_) => "Int_Array",
97            Tag::LongArray(_) => "Long_Array",
98            Tag::End => "END",
99        }
100    }
101
102    fn write_contents(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
103        match self {
104            Tag::Byte(v) => f.write_fmt(format_args!("{}", *v)),
105            Tag::Short(v) => f.write_fmt(format_args!("{}", *v)),
106            Tag::Int(v) => f.write_fmt(format_args!("{}", *v)),
107            Tag::Long(v) => f.write_fmt(format_args!("{}L", *v)),
108            Tag::Float(v) => f.write_fmt(format_args!("{}", *v)),
109            Tag::Double(v) => f.write_fmt(format_args!("{}", *v)),
110            Tag::ByteArray(v) => f.write_fmt(format_args!("[{} bytes]", v.len())),
111            Tag::String(v) => f.write_fmt(format_args!("\"{}\"", v)),
112            Tag::List(v) => {
113                let out = write_contents(v);
114                f.write_str(out.as_str())
115            }
116            Tag::Compound(v) => {
117                let out = write_contents(v);
118                f.write_str(out.as_str())
119            }
120            Tag::IntArray(v) => f.write_fmt(format_args!("[{} ints]", v.len())),
121            Tag::LongArray(v) => f.write_fmt(format_args!("[{} longs]", v.len())),
122            Tag::End => f.write_str("END"),
123        }
124    }
125}
126
127#[cfg(all(test, feature = "std"))]
128impl TestRandom for Tag {
129    fn test_gen_random() -> Self {
130        let random_idx = rand::random::<usize>() % 8;
131        match random_idx {
132            0 => Tag::Byte(i8::test_gen_random()),
133            1 => Tag::Short(i16::test_gen_random()),
134            2 => Tag::Int(i32::test_gen_random()),
135            3 => Tag::Long(i64::test_gen_random()),
136            4 => Tag::Float(f32::test_gen_random()),
137            5 => Tag::Double(f64::test_gen_random()),
138            6 => Tag::String(String::test_gen_random()),
139            7 => Tag::List({
140                let count = rand::random::<usize>() % 256;
141                let mut out = Vec::with_capacity(count);
142                let random_idx = rand::random::<usize>() % 6;
143                for _ in 0..count {
144                    out.push(match random_idx {
145                        0 => Tag::Byte(i8::test_gen_random()),
146                        1 => Tag::Short(i16::test_gen_random()),
147                        2 => Tag::Int(i32::test_gen_random()),
148                        3 => Tag::Long(i64::test_gen_random()),
149                        4 => Tag::Float(f32::test_gen_random()),
150                        5 => Tag::Double(f64::test_gen_random()),
151                        6 => Tag::String(String::test_gen_random()),
152                        other => panic!("impossible {}", other),
153                    });
154                }
155
156                out
157            }),
158            8 => Tag::Compound({
159                let count = rand::random::<usize>() % 256;
160                let mut out = Vec::with_capacity(count);
161                for _ in 0..count {
162                    out.push(NamedTag::test_gen_random());
163                }
164                out
165            }),
166            other => panic!("impossible {}", other),
167        }
168    }
169}
170
171fn write_contents<F>(contents: &Vec<F>) -> String
172    where
173        F: fmt::Display,
174{
175    format!(
176        "{} entries\n{{\n{}\n}}",
177        contents.len(),
178        contents
179            .iter()
180            .flat_map(move |elem| elem
181                .to_string()
182                .split("\n")
183                .map(String::from)
184                .collect::<Vec<String>>())
185            .map(move |line| "  ".to_owned() + line.as_str())
186            .collect::<Vec<String>>()
187            .join("\n")
188    )
189}
190
191// deserialization first
192
193// reads from the root level
194fn read_nbt_data(data: &[u8]) -> DeserializeResult<NamedTag> {
195    let Deserialized { value: tag_type_id, data: _ } = ProtoByteOrder::read_ubyte(data)?;
196    match tag_type_id {
197        0x0A => read_named_tag(data),
198        other => Err(DeserializeErr::NbtInvalidStartTag(other)),
199    }
200}
201
202// reads any named tag: read id -> read name -> read tag with id -> name tag with name
203pub fn read_named_tag(data: &[u8]) -> DeserializeResult<NamedTag> {
204    let Deserialized { value: tag_type_id, data } = ProtoByteOrder::read_ubyte(data)?;
205    if tag_type_id == 0x00 {
206        // tag end
207        Deserialized::ok(Tag::End.with_name(""), data)
208    } else {
209        let Deserialized { value: name, data } = read_string(data)?;
210        Ok(read_tag(tag_type_id, data)?.map(move |payload| NamedTag { name, payload }))
211    }
212}
213
214// reads any tag (given it's id)
215pub fn read_tag(tag_type_id: u8, data: &[u8]) -> DeserializeResult<Tag> {
216    match tag_type_id {
217        0x00 => Deserialized::ok(Tag::End, data),
218        0x01 => read_tag_byte(data),
219        0x02 => read_tag_short(data),
220        0x03 => read_tag_int(data),
221        0x04 => read_tag_long(data),
222        0x05 => read_tag_float(data),
223        0x06 => read_tag_double(data),
224        0x07 => read_tag_byte_array(data),
225        0x08 => read_tag_string(data),
226        0x09 => read_tag_list(data),
227        0x0A => read_tag_compound(data),
228        0x0B => read_tag_int_array(data),
229        0x0C => read_tag_long_array(data),
230        other => Err(DeserializeErr::NbtUnknownTagType(other)),
231    }
232}
233
234fn read_tag_byte(data: &[u8]) -> DeserializeResult<Tag> {
235    Ok(ProtoByteOrder::read_byte(data)?.map(Tag::Byte))
236}
237
238fn read_tag_short(data: &[u8]) -> DeserializeResult<Tag> {
239    Ok(ProtoByteOrder::read_short(data)?.map(Tag::Short))
240}
241
242fn read_tag_int(data: &[u8]) -> DeserializeResult<Tag> {
243    Ok(ProtoByteOrder::read_int(data)?.map(Tag::Int))
244}
245
246fn read_tag_long(data: &[u8]) -> DeserializeResult<Tag> {
247    Ok(ProtoByteOrder::read_long(data)?.map(Tag::Long))
248}
249
250fn read_tag_float(data: &[u8]) -> DeserializeResult<Tag> {
251    Ok(ProtoByteOrder::read_float(data)?.map(Tag::Float))
252}
253
254fn read_tag_double(data: &[u8]) -> DeserializeResult<Tag> {
255    Ok(ProtoByteOrder::read_double(data)?.map(Tag::Double))
256}
257
258fn read_tag_byte_array(data: &[u8]) -> DeserializeResult<Tag> {
259    Ok(ProtoByteOrder::read_int(data)?
260        .and_then(move |size, rest| take(size as usize, rest))?
261        .map(move |arr| Tag::ByteArray(Vec::from(arr))))
262}
263
264fn read_tag_string(data: &[u8]) -> DeserializeResult<Tag> {
265    Ok(read_string(data)?.map(move |str| Tag::String(str)))
266}
267
268fn read_tag_list(data: &[u8]) -> DeserializeResult<Tag> {
269    let Deserialized { value: contents_tag_type_id, data } = ProtoByteOrder::read_ubyte(data)?;
270    let Deserialized { value: list_length, data } = ProtoByteOrder::read_int(data)?;
271    if list_length == 0 {
272        Deserialized::ok(Tag::List(vec![]), data)
273    } else {
274        let mut out_vec = Vec::with_capacity(list_length as usize);
275        let mut remaining_data = data;
276        for _ in 0..list_length {
277            let Deserialized { value: element, data: rest } =
278                read_tag(contents_tag_type_id, &remaining_data)?;
279
280            out_vec.push(element);
281            remaining_data = rest;
282        }
283
284        Deserialized::ok(Tag::List(out_vec), remaining_data)
285    }
286}
287
288fn read_tag_compound(data: &[u8]) -> DeserializeResult<Tag> {
289    let mut out = Vec::new();
290    let mut remaining_data = data;
291    loop {
292        let Deserialized {
293            value: elem,
294            data: rest,
295        } = read_named_tag(remaining_data)?;
296        remaining_data = rest;
297        if elem.is_end() {
298            break;
299        }
300        out.push(elem);
301    }
302
303    Deserialized::ok(Tag::Compound(out), remaining_data)
304}
305
306fn read_tag_int_array(data: &[u8]) -> DeserializeResult<Tag> {
307    read_array_tag(data, ProtoByteOrder::read_int, Tag::IntArray)
308}
309
310fn read_tag_long_array(data: &[u8]) -> DeserializeResult<Tag> {
311    read_array_tag(data, ProtoByteOrder::read_long, Tag::LongArray)
312}
313
314fn read_array_tag<'a, R, F, M>(
315    data: &'a [u8],
316    parser: F,
317    finalizer: M,
318) -> DeserializeResult<'a, Tag>
319    where
320        F: Fn(&'a [u8]) -> DeserializeResult<'a, R>,
321        M: Fn(Vec<R>) -> Tag,
322{
323    let Deserialized { value: count, data } = ProtoByteOrder::read_int(data)?.map(move |v| v as i32);
324    if count < 0 {
325        Err(DeserializeErr::NbtBadLength(count as isize))
326    } else {
327        let mut out = Vec::with_capacity(count as usize);
328        let mut data_remaining = data;
329        for _ in 0..count {
330            let Deserialized {
331                value: elem,
332                data: rest,
333            } = parser(data_remaining)?;
334            data_remaining = rest;
335            out.push(elem);
336        }
337
338        Deserialized::ok(finalizer(out), data_remaining)
339    }
340}
341
342fn read_string(data: &[u8]) -> DeserializeResult<String> {
343    ProtoByteOrder::read_short(data)?
344        .and_then(move |length, data| take(length as usize, data))?
345        .try_map(move |bytes| {
346            String::from_utf8(Vec::from(bytes))
347                .map_err(move |err| DeserializeErr::BadStringEncoding(err))
348        })
349}
350
351// serialize
352impl NamedTag {
353    pub fn bytes(&self) -> Vec<u8> {
354        let type_id = self.payload.id();
355        if type_id == 0x00 {
356            vec![0x00]
357        } else {
358            let payload_bytes = self.payload.bytes();
359            let name_len = self.name.len();
360            let name_len_bytes = ProtoByteOrder::write_ushort(name_len as u16);
361            let mut out =
362                Vec::with_capacity(1 + name_len_bytes.len() + name_len + payload_bytes.len());
363            out.push(type_id);
364            out.extend_from_slice(&name_len_bytes);
365            out.extend(self.name.bytes());
366            out.extend(payload_bytes);
367            out
368        }
369    }
370}
371
372impl Tag {
373    pub fn id(&self) -> u8 {
374        match self {
375            Tag::Byte(_) => 0x01,
376            Tag::Short(_) => 0x02,
377            Tag::Int(_) => 0x03,
378            Tag::Long(_) => 0x04,
379            Tag::Float(_) => 0x05,
380            Tag::Double(_) => 0x06,
381            Tag::ByteArray(_) => 0x07,
382            Tag::String(_) => 0x08,
383            Tag::List(_) => 0x09,
384            Tag::Compound(_) => 0x0A,
385            Tag::IntArray(_) => 0x0B,
386            Tag::LongArray(_) => 0x0C,
387            Tag::End => 0x00,
388        }
389    }
390
391    pub fn bytes(&self) -> Vec<u8> {
392        match self {
393            Tag::Byte(b) => vec![*b as u8],
394            Tag::Short(v) => Vec::from(ProtoByteOrder::write_short(*v)),
395            Tag::Int(v) => Vec::from(ProtoByteOrder::write_int(*v)),
396            Tag::Long(v) => Vec::from(ProtoByteOrder::write_long(*v)),
397            Tag::Float(v) => Vec::from(ProtoByteOrder::write_float(*v)),
398            Tag::Double(v) => Vec::from(ProtoByteOrder::write_double(*v)),
399            Tag::ByteArray(v) => {
400                let n = v.len();
401                let mut out = Vec::with_capacity(n + 4);
402                let size_bytes = ProtoByteOrder::write_uint(n as u32);
403                out.extend_from_slice(&size_bytes);
404                out.extend(v);
405                out
406            }
407            Tag::String(v) => {
408                let n = v.len();
409                let mut out = Vec::with_capacity(n + 2);
410                let size_bytes = ProtoByteOrder::write_ushort(n as u16);
411                out.extend_from_slice(&size_bytes);
412                out.extend(v.bytes());
413                out
414            }
415            Tag::List(v) => {
416                let count = v.len();
417                let elem_id = {
418                    if count == 0 {
419                        0x00
420                    } else {
421                        let mut id = None;
422                        for elem in v {
423                            let elem_id = elem.id();
424                            if let Some(old_id) = id.replace(elem_id) {
425                                if old_id != elem_id {
426                                    panic!(
427                                        "list contains tags of different types, cannot serialize"
428                                    );
429                                }
430                            }
431                        }
432
433                        id.expect("there must be some elements in the list")
434                    }
435                };
436
437                let mut out = Vec::new();
438                out.push(elem_id);
439                let count_bytes = ProtoByteOrder::write_uint(count as u32);
440                out.extend_from_slice(&count_bytes);
441                out.extend(v.iter().flat_map(move |elem| elem.bytes().into_iter()));
442                out
443            }
444            Tag::Compound(v) => {
445                let mut out = Vec::new();
446                for elem in v {
447                    out.extend(elem.bytes());
448                }
449                out.extend(Tag::End.with_name("").bytes());
450                out
451            }
452            Tag::IntArray(v) => {
453                let n = v.len();
454                let mut out = Vec::with_capacity(4 + (4 * n));
455                let n_bytes = ProtoByteOrder::write_uint(n as u32);
456                out.extend_from_slice(&n_bytes);
457                for value in v {
458                    let bytes = ProtoByteOrder::write_int(*value);
459                    out.extend_from_slice(&bytes);
460                }
461                out
462            }
463            Tag::LongArray(v) => {
464                let n = v.len();
465                let mut out = Vec::with_capacity(4 + (8 * n));
466                let n_bytes = ProtoByteOrder::write_uint(n as u32);
467                out.extend_from_slice(&n_bytes);
468                for value in v {
469                    let bytes = ProtoByteOrder::write_long(*value);
470                    out.extend_from_slice(&bytes);
471                }
472                out
473            }
474            Tag::End => Vec::default(),
475        }
476    }
477}
478
479// test
480#[cfg(test)]
481mod tests {
482    use super::*;
483
484    #[cfg(feature = "std")]
485    #[test]
486    fn test_read_bignbt_example() {
487        let actual = read_bigtest();
488
489        let expected = Tag::Compound(vec!(
490            Tag::Long(9223372036854775807).with_name("longTest"),
491            Tag::Short(32767).with_name("shortTest"),
492            Tag::String("HELLO WORLD THIS IS A TEST STRING ÅÄÖ!".into()).with_name("stringTest"),
493            Tag::Float(0.49823147).with_name("floatTest"),
494            Tag::Int(2147483647).with_name("intTest"),
495            Tag::Compound(vec!(
496                Tag::Compound(vec!(
497                    Tag::String("Hampus".into()).with_name("name"),
498                    Tag::Float(0.75).with_name("value"),
499                )).with_name("ham"),
500                Tag::Compound(vec!(
501                    Tag::String("Eggbert".into()).with_name("name"),
502                    Tag::Float(0.5).with_name("value"),
503                )).with_name("egg")
504            )).with_name("nested compound test"),
505            Tag::List(vec!(
506                Tag::Long(11),
507                Tag::Long(12),
508                Tag::Long(13),
509                Tag::Long(14),
510                Tag::Long(15),
511            )).with_name("listTest (long)"),
512            Tag::List(vec!(
513                Tag::Compound(vec!(
514                    Tag::String("Compound tag #0".into()).with_name("name"),
515                    Tag::Long(1264099775885).with_name("created-on"),
516                )),
517                Tag::Compound(vec!(
518                    Tag::String("Compound tag #1".into()).with_name("name"),
519                    Tag::Long(1264099775885).with_name("created-on"),
520                ))
521            )).with_name("listTest (compound)"),
522            Tag::Byte(127).with_name("byteTest"),
523            Tag::ByteArray(bigtest_generate_byte_array()).with_name("byteArrayTest (the first 1000 values of (n*n*255+n*7)%100, starting with n=0 (0, 62, 34, 16, 8, ...))"),
524            Tag::Double(0.4931287132182315).with_name("doubleTest")
525        )).with_name("Level");
526
527        assert_eq!(actual, expected);
528    }
529
530    #[cfg(feature = "std")]
531    #[test]
532    fn test_serialize_bigtest() {
533        let (unzipped, result) = read_bigtest_with_bytes();
534        let serialized = result.bytes();
535        assert_eq!(unzipped, serialized);
536        let Deserialized {
537            value: unserialized,
538            data: _,
539        } = NamedTag::root_compound_tag_from_bytes(serialized.as_slice())
540            .expect("deserialize serialized nbt");
541        assert_eq!(unserialized, result);
542    }
543
544    #[test]
545    fn test_int_array() {
546        let original = Tag::Compound(vec![Tag::IntArray(vec![
547            1, 2, -5, 123127, -12373, 0, 0, 4, 2,
548        ])
549            .with_name("test ints")])
550            .with_name("test");
551
552        let bytes = original.bytes();
553        let Deserialized {
554            value: unserialized,
555            data: _,
556        } = NamedTag::root_compound_tag_from_bytes(bytes.as_slice())
557            .expect("deserialize int array");
558        assert_eq!(original, unserialized);
559    }
560
561    #[test]
562    fn test_long_array() {
563        let original = Tag::Compound(vec![Tag::LongArray(vec![
564            1,
565            2,
566            -5,
567            123127999999,
568            -1237399999,
569            0,
570            0,
571            4,
572            2,
573        ])
574            .with_name("test ints")])
575            .with_name("test");
576
577        let bytes = original.bytes();
578        let Deserialized {
579            value: unserialized,
580            data: _,
581        } = NamedTag::root_compound_tag_from_bytes(bytes.as_slice())
582            .expect("deserialize int array");
583        assert_eq!(original, unserialized);
584    }
585
586    #[cfg(feature = "std")]
587    #[test]
588    fn test_display() {
589        println!("{}", read_bigtest());
590    }
591
592    #[cfg(feature = "std")]
593    #[test]
594    fn test_debug() {
595        println!("{:?}", read_bigtest());
596    }
597
598    #[cfg(feature = "std")]
599    fn read_bigtest_with_bytes() -> (Vec<u8>, NamedTag) {
600        let unzipped = read_compressed_file("src/testdata/bigtest.nbt").expect("read nbt data");
601        let Deserialized {
602            value: result,
603            data: rest,
604        } = NamedTag::root_compound_tag_from_bytes(unzipped.as_slice()).expect("deserialize nbt");
605        assert_eq!(rest.len(), 0);
606
607        (unzipped, result)
608    }
609
610    #[cfg(feature = "std")]
611    fn read_bigtest() -> NamedTag {
612        let (_, result) = read_bigtest_with_bytes();
613        result
614    }
615
616    #[cfg(feature = "std")]
617    fn bigtest_generate_byte_array() -> Vec<u8> {
618        const COUNT: usize = 1000;
619        let mut out = Vec::with_capacity(COUNT);
620        for i in 0..COUNT {
621            out.push((((i * i * 255) + (i * 7)) % 100) as u8);
622        }
623        out
624    }
625
626    #[cfg(feature = "std")]
627    fn read_compressed_file(at: &str) -> std::io::Result<Vec<u8>> {
628        use flate2::read::GzDecoder;
629        use std::fs::File;
630        use std::io::Read;
631
632        let file = File::open(at)?;
633        let mut gz = GzDecoder::new(file);
634        let mut out = Vec::new();
635        gz.read_to_end(&mut out)?;
636        Ok(out)
637    }
638}