diff --git a/Cargo.lock b/Cargo.lock index 52da5b346..2497dd281 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2951,6 +2951,7 @@ version = "0.1.0-dev+26.1" dependencies = [ "bytes", "cesu8", + "criterion", "flate2", "pumpkin-codecs", "serde", diff --git a/pumpkin-nbt/Cargo.toml b/pumpkin-nbt/Cargo.toml index 50182eec9..592357132 100644 --- a/pumpkin-nbt/Cargo.toml +++ b/pumpkin-nbt/Cargo.toml @@ -17,6 +17,12 @@ tracing.workspace = true [dev-dependencies] tempfile.workspace = true +criterion.workspace = true +serde = { version = "1.0", features = ["derive"] } + +[[bench]] +name = "nbt" +harness = false [lints] workspace = true diff --git a/pumpkin-nbt/benches/nbt.rs b/pumpkin-nbt/benches/nbt.rs new file mode 100644 index 000000000..ae0d636fc --- /dev/null +++ b/pumpkin-nbt/benches/nbt.rs @@ -0,0 +1,114 @@ +use criterion::{Criterion, criterion_group, criterion_main}; +use pumpkin_nbt::{Nbt, NbtCompound, deserializer, serializer, tag::NbtTag}; +use serde::{Deserialize, Serialize}; +use std::io::Cursor; + +#[derive(Serialize, Deserialize, PartialEq, Debug, Clone)] +struct LargeTest { + byte: i8, + short: i16, + int: i32, + long: i64, + float: f32, + double: f64, + string: String, + byte_array: Vec, + int_array: Vec, + long_array: Vec, + list_string: Vec, + nested: Option>, +} + +fn create_large_test_data(depth: usize) -> LargeTest { + LargeTest { + byte: 123, + short: 1342, + int: 4313, + long: 34, + float: 1.00, + double: 69.42, + string: "Hello test benchmark data".to_string(), + byte_array: vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + int_array: vec![10, 11, 12, 13, 14, 15, 16, 17, 18, 19], + long_array: vec![20, 21, 22, 23, 24, 25, 26, 27, 28, 29], + list_string: vec!["one".to_string(), "two".to_string(), "three".to_string()], + nested: (depth > 0).then(|| Box::new(create_large_test_data(depth - 1))), + } +} + +fn create_large_compound(depth: usize) -> NbtCompound { + let mut compound = NbtCompound::new(); + compound.put_byte("byte", 123); + compound.put_short("short", 1342); + compound.put_int("int", 4313); + compound.put_long("long", 34); + compound.put_float("float", 1.00); + compound.put_double("double", 69.42); + compound.put_string("string", "Hello test benchmark data".to_string()); + compound.put( + "byte_array", + NbtTag::ByteArray(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), + ); + compound.put( + "int_array", + NbtTag::IntArray(vec![10, 11, 12, 13, 14, 15, 16, 17, 18, 19]), + ); + compound.put( + "long_array", + NbtTag::LongArray(vec![20, 21, 22, 23, 24, 25, 26, 27, 28, 29]), + ); + + let list = vec![ + NbtTag::String("one".to_string()), + NbtTag::String("two".to_string()), + NbtTag::String("three".to_string()), + ]; + compound.put_list("list_string", list); + + if depth > 0 { + compound.put_compound("nested", create_large_compound(depth - 1)); + } + compound +} + +pub fn bench_nbt(c: &mut Criterion) { + let test_data = create_large_test_data(5); + let mut ser_bytes = Vec::new(); + serializer::to_bytes(&test_data, &mut ser_bytes).unwrap(); + + let compound_data = create_large_compound(5); + let nbt_wrapper = Nbt::new(String::new(), compound_data.clone()); + let wrapper_bytes = nbt_wrapper.write(); + + c.bench_function("nbt/serialize/serde", |b| { + b.iter(|| { + let mut out = Vec::new(); + serializer::to_bytes(&test_data, &mut out).unwrap(); + }); + }); + + c.bench_function("nbt/deserialize/serde", |b| { + b.iter(|| { + let cursor = Cursor::new(&ser_bytes); + let _: LargeTest = deserializer::from_bytes(cursor).unwrap(); + }); + }); + + c.bench_function("nbt/serialize/raw", |b| { + b.iter(|| { + let nbt = Nbt::new(String::new(), compound_data.clone()); + let _ = nbt.write(); + }); + }); + + c.bench_function("nbt/deserialize/raw", |b| { + b.iter(|| { + let mut cursor = Cursor::new(&wrapper_bytes); + let mut reader = deserializer::NbtReadHelper::new(&mut cursor); + Nbt::read(&mut reader).unwrap(); + }); + }); +} + +criterion_group!(benches, bench_nbt); +criterion_main!(benches); diff --git a/pumpkin-nbt/fuzz/.gitignore b/pumpkin-nbt/fuzz/.gitignore new file mode 100644 index 000000000..a0925114d --- /dev/null +++ b/pumpkin-nbt/fuzz/.gitignore @@ -0,0 +1,3 @@ +target +corpus +artifacts diff --git a/pumpkin-nbt/fuzz/Cargo.lock b/pumpkin-nbt/fuzz/Cargo.lock new file mode 100644 index 000000000..da69628bc --- /dev/null +++ b/pumpkin-nbt/fuzz/Cargo.lock @@ -0,0 +1,362 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d036a3c4ab069c7b410a2ce876bd74808d2d0888a82667669f8e783a898bf1" +dependencies = [ + "derive_arbitrary", +] + +[[package]] +name = "bytes" +version = "1.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" + +[[package]] +name = "cc" +version = "1.2.62" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1dce859f0832a7d088c4f1119888ab94ef4b5d6795d1ce05afb7fe159d79f98" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cesu8" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "derive_arbitrary" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e567bd82dcff979e4b03460c307b3cdc9e96fde3d73bed1496d2bc75d9dd62a" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasip2", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom", + "libc", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libfuzzer-sys" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f12a681b7dd8ce12bff52488013ba614b869148d54dd79836ab85aafdd53f08d" +dependencies = [ + "arbitrary", + "cc", +] + +[[package]] +name = "memchr" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pumpkin-codecs" +version = "0.1.0-dev+26.1" +dependencies = [ + "either", + "serde_json", + "tracing", +] + +[[package]] +name = "pumpkin-nbt" +version = "0.1.0-dev+26.1" +dependencies = [ + "bytes", + "cesu8", + "flate2", + "pumpkin-codecs", + "serde", + "thiserror", + "tracing", +] + +[[package]] +name = "pumpkin-nbt-fuzz" +version = "0.0.0" +dependencies = [ + "arbitrary", + "libfuzzer-sys", + "pumpkin-nbt", + "serde", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "wasip2" +version = "1.0.3+wasi-0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "20064672db26d7cdc89c7798c48a0fdfac8213434a1186e5ef29fd560ae223d6" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/pumpkin-nbt/fuzz/Cargo.toml b/pumpkin-nbt/fuzz/Cargo.toml new file mode 100644 index 000000000..5eed6150b --- /dev/null +++ b/pumpkin-nbt/fuzz/Cargo.toml @@ -0,0 +1,46 @@ +[package] +name = "pumpkin-nbt-fuzz" +version = "0.0.0" +publish = false +edition = "2024" + +[workspace] + +[package.metadata] +cargo-fuzz = true + +[dependencies] +libfuzzer-sys = "0.4" +serde = { version = "1.0", features = ["derive"] } +arbitrary = { version = "1.3", features = ["derive"] } + +[dependencies.pumpkin-nbt] +path = ".." + +[[bin]] +name = "deserialize" +path = "fuzz_targets/deserialize.rs" +test = false +doc = false +bench = false + +[[bin]] +name = "deserialize_serde" +path = "fuzz_targets/deserialize_serde.rs" +test = false +doc = false +bench = false + +[[bin]] +name = "decompress_gzip" +path = "fuzz_targets/decompress_gzip.rs" +test = false +doc = false +bench = false + +[[bin]] +name = "roundtrip" +path = "fuzz_targets/roundtrip.rs" +test = false +doc = false +bench = false diff --git a/pumpkin-nbt/fuzz/fuzz_targets/decompress_gzip.rs b/pumpkin-nbt/fuzz/fuzz_targets/decompress_gzip.rs new file mode 100644 index 000000000..ab1ebcfb1 --- /dev/null +++ b/pumpkin-nbt/fuzz/fuzz_targets/decompress_gzip.rs @@ -0,0 +1,9 @@ +#![no_main] +use libfuzzer_sys::fuzz_target; +use pumpkin_nbt::nbt_compress::read_gzip_compound_tag; +use std::io::Cursor; + +fuzz_target!(|data: &[u8]| { + let cursor = Cursor::new(data); + let _ = read_gzip_compound_tag(cursor); +}); diff --git a/pumpkin-nbt/fuzz/fuzz_targets/deserialize.rs b/pumpkin-nbt/fuzz/fuzz_targets/deserialize.rs new file mode 100644 index 000000000..fb36a16e0 --- /dev/null +++ b/pumpkin-nbt/fuzz/fuzz_targets/deserialize.rs @@ -0,0 +1,15 @@ +#![no_main] +use libfuzzer_sys::fuzz_target; +use pumpkin_nbt::Nbt; +use pumpkin_nbt::deserializer::NbtReadHelper; +use std::io::Cursor; + +fuzz_target!(|data: &[u8]| { + let mut cursor = Cursor::new(data); + let mut reader = NbtReadHelper::new(&mut cursor); + let _ = Nbt::read(&mut reader); + + let mut cursor_unnamed = Cursor::new(data); + let mut reader_unnamed = NbtReadHelper::new(&mut cursor_unnamed); + let _ = Nbt::read_unnamed(&mut reader_unnamed); +}); diff --git a/pumpkin-nbt/fuzz/fuzz_targets/deserialize_serde.rs b/pumpkin-nbt/fuzz/fuzz_targets/deserialize_serde.rs new file mode 100644 index 000000000..c5b99fe0f --- /dev/null +++ b/pumpkin-nbt/fuzz/fuzz_targets/deserialize_serde.rs @@ -0,0 +1,29 @@ +#![no_main] +use libfuzzer_sys::fuzz_target; +use pumpkin_nbt::deserializer::{from_bytes, from_bytes_unnamed}; +use serde::Deserialize; +use std::io::Cursor; + +#[derive(Deserialize, Debug, PartialEq)] +struct FuzzStruct { + byte: i8, + short: i16, + int: i32, + long: i64, + float: f32, + double: f64, + string: String, + byte_array: Vec, + int_array: Vec, + long_array: Vec, + list_string: Vec, + nested: Option>, +} + +fuzz_target!(|data: &[u8]| { + let cursor = Cursor::new(data); + let _: Result = from_bytes(cursor); + + let cursor_unnamed = Cursor::new(data); + let _: Result = from_bytes_unnamed(cursor_unnamed); +}); diff --git a/pumpkin-nbt/fuzz/fuzz_targets/roundtrip.rs b/pumpkin-nbt/fuzz/fuzz_targets/roundtrip.rs new file mode 100644 index 000000000..cea53751d --- /dev/null +++ b/pumpkin-nbt/fuzz/fuzz_targets/roundtrip.rs @@ -0,0 +1,33 @@ +#![no_main] +use arbitrary::Arbitrary; +use libfuzzer_sys::fuzz_target; +use pumpkin_nbt::deserializer::from_bytes; +use pumpkin_nbt::serializer::to_bytes; +use serde::{Deserialize, Serialize}; +use std::io::Cursor; + +#[derive(Arbitrary, Serialize, Deserialize, Debug, PartialEq)] +struct RoundtripStruct { + byte: i8, + short: i16, + int: i32, + long: i64, + float: f32, + double: f64, + string: String, + byte_array: Vec, + int_array: Vec, + long_array: Vec, + list_string: Vec, + // Nested struct to test recursion + nested: Option>, +} + +fuzz_target!(|data: RoundtripStruct| { + let mut bytes = Vec::new(); + if to_bytes(&data, &mut bytes).is_ok() { + let cursor = Cursor::new(bytes); + let reconstructed: RoundtripStruct = from_bytes(cursor).expect("Failed to deserialize valid NBT"); + assert_eq!(data, reconstructed); + } +}); diff --git a/pumpkin-nbt/src/deserializer.rs b/pumpkin-nbt/src/deserializer.rs index 057803070..df30a3a2e 100644 --- a/pumpkin-nbt/src/deserializer.rs +++ b/pumpkin-nbt/src/deserializer.rs @@ -3,7 +3,7 @@ use std::io::{Seek, SeekFrom}; use crate::{ BYTE_ARRAY_ID, BYTE_ID, COMPOUND_ID, END_ID, Error, INT_ARRAY_ID, INT_ID, LIST_ID, - LONG_ARRAY_ID, LONG_ID, NbtTag, get_nbt_string, io, + LONG_ARRAY_ID, LONG_ID, MAX_ARRAY_LENGTH, NbtTag, get_nbt_string, io, }; use io::Read; use serde::de::{self, DeserializeSeed, IntoDeserializer, MapAccess, SeqAccess, Visitor}; @@ -150,13 +150,18 @@ impl<'de, R: Read + Seek> de::Deserializer<'de> for &mut Deserializer { return Err(Error::NegativeLength(remaining_values)); } + let remaining_values = remaining_values as usize; + if remaining_values > MAX_ARRAY_LENGTH { + return Err(Error::LargeLength(remaining_values)); + } + //TODO this is a bit hacky but I couldn't think of a better way // This flag gets auto cleared in visit_seq set_curr_visitor_seq_list_id(Some(list_type)); let result = visitor.visit_seq(ListAccess { de: self, list_type, - remaining_values: remaining_values as usize, + remaining_values, })?; Ok(result) } diff --git a/pumpkin-nbt/src/lib.rs b/pumpkin-nbt/src/lib.rs index 6e771c3b1..63d43446d 100644 --- a/pumpkin-nbt/src/lib.rs +++ b/pumpkin-nbt/src/lib.rs @@ -5,7 +5,6 @@ use std::{ }; use bytes::Bytes; -use compound::NbtCompound; use deserializer::NbtReadHelper; use serde::{de, ser}; use serializer::WriteAdaptor; @@ -19,6 +18,7 @@ pub mod nbt_ops; pub mod serializer; pub mod tag; +pub use compound::NbtCompound; pub use deserializer::{from_bytes, from_bytes_unnamed}; pub use serializer::{to_bytes, to_bytes_named, to_bytes_unnamed}; @@ -38,6 +38,8 @@ pub const COMPOUND_ID: u8 = 0x0A; pub const INT_ARRAY_ID: u8 = 0x0B; pub const LONG_ARRAY_ID: u8 = 0x0C; +pub const MAX_ARRAY_LENGTH: usize = 2_000_000; + #[derive(Error, Debug)] pub enum Error { #[error("The root tag of the NBT file is not a compound tag. Received tag id: {0}")] diff --git a/pumpkin-nbt/src/nbt_compress.rs b/pumpkin-nbt/src/nbt_compress.rs index bdaee61c2..a2ab71735 100644 --- a/pumpkin-nbt/src/nbt_compress.rs +++ b/pumpkin-nbt/src/nbt_compress.rs @@ -14,7 +14,7 @@ use std::io::{Cursor, Read, Seek, Write}; /// A Result containing either the parsed `NbtCompound` or an Error pub fn read_gzip_compound_tag(input: impl Read + Seek) -> Result { // Create a GZip decoder and directly chain it to the NBT reader - let mut decoder = GzDecoder::new(input); + let mut decoder = GzDecoder::new(input).take(64 * 1024 * 1024); // 64 MB limit let mut buf = Vec::new(); decoder.read_to_end(&mut buf).map_err(Error::Incomplete)?; let mut reader = NbtReadHelper::new(Cursor::new(buf)); @@ -69,7 +69,7 @@ pub fn write_gzip_compound_tag_to_bytes(compound: NbtCompound) -> Result /// A Result containing either the deserialized type or an Error pub fn from_gzip_bytes<'a, T: serde::Deserialize<'a>, R: Read>(input: R) -> Result { // Create a GZip decoder and directly use it for deserialization - let mut decoder = GzDecoder::new(input); + let mut decoder = GzDecoder::new(input).take(64 * 1024 * 1024); // 64 MB limit let mut buf = Vec::new(); decoder.read_to_end(&mut buf).map_err(Error::Incomplete)?; deserializer::from_bytes(Cursor::new(buf)) diff --git a/pumpkin-nbt/src/tag.rs b/pumpkin-nbt/src/tag.rs index a211f46af..80348a12d 100644 --- a/pumpkin-nbt/src/tag.rs +++ b/pumpkin-nbt/src/tag.rs @@ -6,8 +6,8 @@ use serializer::WriteAdaptor; use crate::{ BYTE_ARRAY_ID, BYTE_ID, COMPOUND_ID, DOUBLE_ID, END_ID, Error, FLOAT_ID, INT_ARRAY_ID, INT_ID, - LIST_ID, LONG_ARRAY_ID, LONG_ID, SHORT_ID, STRING_ID, Seek, Write, compound, deserializer, - get_nbt_string, io, nbt_byte_array, nbt_int_array, nbt_long_array, serializer, + LIST_ID, LONG_ARRAY_ID, LONG_ID, MAX_ARRAY_LENGTH, SHORT_ID, STRING_ID, Seek, Write, compound, + deserializer, get_nbt_string, io, nbt_byte_array, nbt_int_array, nbt_long_array, serializer, }; #[derive(Clone, Debug, PartialEq)] @@ -286,6 +286,9 @@ impl NbtTag { } let len = len as usize; + if len > MAX_ARRAY_LENGTH { + return Err(Error::LargeLength(len)); + } let mut byte_array = Vec::with_capacity(len); for _ in 0..len { let byte = reader.get_i8_be()?; @@ -301,7 +304,12 @@ impl NbtTag { return Err(Error::NegativeLength(len)); } - let mut list = Vec::with_capacity(len as usize); + let len = len as usize; + if len > MAX_ARRAY_LENGTH { + return Err(Error::LargeLength(len)); + } + + let mut list = Vec::with_capacity(len); for _ in 0..len { let tag = Self::deserialize_data(reader, tag_type_id)?; assert_eq!(tag.get_type_id(), tag_type_id); @@ -318,6 +326,9 @@ impl NbtTag { } let len = len as usize; + if len > MAX_ARRAY_LENGTH { + return Err(Error::LargeLength(len)); + } let mut int_array = Vec::with_capacity(len); for _ in 0..len { let int = reader.get_i32_be()?; @@ -332,6 +343,9 @@ impl NbtTag { } let len = len as usize; + if len > MAX_ARRAY_LENGTH { + return Err(Error::LargeLength(len)); + } let mut long_array = Vec::with_capacity(len); for _ in 0..len { let long = reader.get_i64_be()?;