Skip to main content

steel_registry/data_components/components/
custom_data.rs

1//! Vanilla `CustomData` item component value.
2
3use 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/// Immutable component wrapper around a Vanilla-normalized NBT compound.
18#[derive(Debug, Clone)]
19pub struct CustomData {
20    tag: NbtCompound,
21}
22
23impl CustomData {
24    /// Creates a component after validating strings and applying Vanilla map
25    /// semantics to duplicate compound keys.
26    #[must_use]
27    pub fn try_from_compound(tag: NbtCompound) -> Option<Self> {
28        normalize_nbt_compound(tag).map(|tag| Self { tag })
29    }
30
31    /// Decodes the persistent codec's compound-or-flattened-SNBT alternatives.
32    #[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    /// Mirrors `CustomData.matchedBy` and `NbtUtils.compareNbt`.
58    #[must_use]
59    pub fn matched_by(&self, expected: &NbtCompound) -> bool {
60        compare_nbt_compounds(expected, &self.tag, true)
61    }
62
63    /// Returns a copy updated through Vanilla `CompoundTag.merge` semantics.
64    #[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}