steel_registry/data_components/components/
custom_data.rs1use std::io::{Cursor, Error, Result, Write};
4
5use simdnbt::owned::{NbtCompound, NbtTag, read_tag};
6use simdnbt::{FromNbtTag, ToNbtTag};
7use steel_utils::{
8 hash::{ComponentHasher, HashComponent},
9 nbt::{
10 compare_nbt_compounds, merge_nbt_compounds, normalize_nbt_compound, vanilla_nbt_heap_size,
11 },
12 serial::{ReadFrom, WriteTo},
13};
14
15const DEFAULT_NBT_QUOTA: u64 = 2_097_152;
16
17#[derive(Debug, Clone)]
19pub struct CustomData {
20 tag: NbtCompound,
21}
22
23impl CustomData {
24 #[must_use]
27 pub fn try_from_compound(tag: NbtCompound) -> Option<Self> {
28 normalize_nbt_compound(tag).map(|tag| Self { tag })
29 }
30
31 #[must_use]
33 pub fn from_nbt_value(tag: &NbtTag) -> Option<Self> {
34 Self::from_codec_tag(tag.clone())
35 }
36
37 #[must_use]
38 pub fn is_empty(&self) -> bool {
39 self.tag.is_empty()
40 }
41
42 #[must_use]
43 pub fn copy_tag(&self) -> NbtCompound {
44 self.tag.clone()
45 }
46
47 #[must_use]
48 pub const fn as_compound(&self) -> &NbtCompound {
49 &self.tag
50 }
51
52 pub(crate) fn without_field(mut self, name: &str) -> Self {
53 self.tag.remove(name);
54 self
55 }
56
57 #[must_use]
59 pub fn matched_by(&self, expected: &NbtCompound) -> bool {
60 compare_nbt_compounds(expected, &self.tag, true)
61 }
62
63 #[must_use]
65 pub fn merged_with(&self, other: &Self) -> Self {
66 let mut tag = self.tag.clone();
67 merge_nbt_compounds(&mut tag, &other.tag);
68 Self { tag }
69 }
70
71 pub(crate) fn read_codec_network(data: &mut Cursor<&[u8]>) -> Result<Self> {
72 Self::from_codec_tag(read_network_tag(data)?).ok_or_else(|| {
73 Error::other("Custom data network value is not a compound or SNBT string")
74 })
75 }
76
77 fn from_codec_tag(tag: NbtTag) -> Option<Self> {
78 match tag {
79 NbtTag::Compound(compound) => Self::try_from_compound(compound),
80 NbtTag::String(value) => {
81 let value = value.try_into_string().ok()?;
82 let compound = steel_utils::nbt::parse_snbt_compound(&value).ok()?;
83 Self::try_from_compound(compound)
84 }
85 _ => None,
86 }
87 }
88}
89
90impl Default for CustomData {
91 fn default() -> Self {
92 Self {
93 tag: NbtCompound::new(),
94 }
95 }
96}
97
98impl PartialEq for CustomData {
99 fn eq(&self, other: &Self) -> bool {
100 steel_utils::nbt::nbt_compounds_equal(&self.tag, &other.tag)
101 }
102}
103
104impl WriteTo for CustomData {
105 fn write(&self, writer: &mut impl Write) -> Result<()> {
106 let mut encoded = Vec::new();
107 NbtTag::Compound(self.tag.clone()).write(&mut encoded);
108 writer.write_all(&encoded)
109 }
110}
111
112impl ReadFrom for CustomData {
113 fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
114 let NbtTag::Compound(compound) = read_network_tag(data)? else {
115 return Err(Error::other(
116 "Bucket entity data network value is not a compound",
117 ));
118 };
119 Self::try_from_compound(compound)
120 .ok_or_else(|| Error::other("Bucket entity data contains malformed modified UTF-8"))
121 }
122}
123
124fn read_network_tag(data: &mut Cursor<&[u8]>) -> Result<NbtTag> {
125 let tag = read_tag(data).map_err(|error| Error::other(format!("Invalid NBT: {error:?}")))?;
126 let Some(heap_size) = vanilla_nbt_heap_size(&tag) else {
127 return Err(Error::other("NBT contains malformed modified UTF-8"));
128 };
129 if heap_size > DEFAULT_NBT_QUOTA {
130 return Err(Error::other(format!(
131 "NBT exceeds Vanilla's {DEFAULT_NBT_QUOTA}-byte heap quota"
132 )));
133 }
134 Ok(tag)
135}
136
137impl ToNbtTag for CustomData {
138 fn to_nbt_tag(self) -> NbtTag {
139 NbtTag::Compound(self.tag)
140 }
141}
142
143impl FromNbtTag for CustomData {
144 fn from_nbt_tag(tag: simdnbt::borrow::NbtTag<'_, '_>) -> Option<Self> {
145 Self::from_codec_tag(tag.to_owned())
146 }
147}
148
149impl HashComponent for CustomData {
150 fn hash_component(&self, hasher: &mut ComponentHasher) {
151 NbtTag::Compound(self.tag.clone()).hash_component(hasher);
152 }
153}
154
155#[cfg(test)]
156mod tests {
157 use std::io::Cursor;
158
159 use simdnbt::owned::{NbtCompound, NbtList, NbtTag};
160 use steel_utils::{
161 hash::HashComponent as _,
162 serial::{ReadFrom as _, WriteTo as _},
163 };
164
165 use super::CustomData;
166
167 fn sample_compound() -> NbtCompound {
168 let mut nested = NbtCompound::new();
169 nested.insert("name", "steel");
170 let mut compound = NbtCompound::new();
171 compound.insert("value", 7);
172 compound.insert("nested", nested);
173 compound
174 }
175
176 #[test]
177 fn persistent_codec_accepts_compounds_and_flattened_snbt() {
178 let direct = CustomData::from_nbt_value(&NbtTag::Compound(sample_compound()))
179 .expect("compound should decode");
180 let flattened =
181 CustomData::from_nbt_value(&NbtTag::String("{value:7,nested:{name:'steel'}}".into()))
182 .expect("flattened SNBT should decode");
183
184 assert_eq!(direct, flattened);
185 assert!(CustomData::from_nbt_value(&NbtTag::Int(7)).is_none());
186 }
187
188 #[test]
189 fn network_codecs_differ_only_on_the_flattened_alternative() {
190 let value = CustomData::try_from_compound(sample_compound())
191 .expect("sample compound should be valid");
192 let mut encoded = Vec::new();
193 value
194 .write(&mut encoded)
195 .expect("custom data should encode");
196
197 assert_eq!(
198 CustomData::read(&mut Cursor::new(encoded.as_slice()))
199 .expect("raw compound stream should decode"),
200 value
201 );
202 assert_eq!(
203 CustomData::read_codec_network(&mut Cursor::new(encoded.as_slice()))
204 .expect("codec-derived stream should decode"),
205 value
206 );
207
208 let mut flattened = Vec::new();
209 NbtTag::String("{value:7}".into()).write(&mut flattened);
210 assert!(CustomData::read(&mut Cursor::new(flattened.as_slice())).is_err());
211 assert!(CustomData::read_codec_network(&mut Cursor::new(flattened.as_slice())).is_ok());
212 }
213
214 #[test]
215 fn network_decode_enforces_vanilla_nbt_heap_quota() {
216 let mut compound = NbtCompound::new();
217 compound.insert("values", NbtList::String(vec!["".into(); 60_000]));
218 let mut encoded = Vec::new();
219 NbtTag::Compound(compound).write(&mut encoded);
220
221 assert!(CustomData::read_codec_network(&mut Cursor::new(encoded.as_slice())).is_err());
222 }
223
224 #[test]
225 fn persistent_hash_uses_the_compound_codec_shape() {
226 let compound = sample_compound();
227 let value = CustomData::try_from_compound(compound.clone())
228 .expect("sample compound should be valid");
229
230 assert_eq!(
231 value.compute_hash(),
232 NbtTag::Compound(compound).compute_hash()
233 );
234 }
235}