Skip to main content

steel_registry/data_components/registry/
patch_network.rs

1use super::{
2    ComponentData, ComponentPatchEntry, Cursor, DataComponentPatch, Identifier, ReadFrom, Result,
3    VarInt, Write, WriteTo,
4};
5
6impl WriteTo for DataComponentPatch {
7    fn write(&self, writer: &mut impl Write) -> Result<()> {
8        use crate::{REGISTRY, RegistryExt};
9
10        let mut added: Vec<(&Identifier, &ComponentData)> = Vec::new();
11        let mut removed: Vec<&Identifier> = Vec::new();
12
13        for (key, entry) in &self.entries {
14            match entry {
15                ComponentPatchEntry::Set(data) => added.push((key, data)),
16                ComponentPatchEntry::Removed => removed.push(key),
17            }
18        }
19
20        let added_count = i32::try_from(added.len())
21            .map_err(|_| std::io::Error::other("Too many added data components"))?;
22        let removed_count = i32::try_from(removed.len())
23            .map_err(|_| std::io::Error::other("Too many removed data components"))?;
24        VarInt(added_count).write(writer)?;
25        VarInt(removed_count).write(writer)?;
26
27        // Write added components
28        for (key, data) in added {
29            let id = REGISTRY
30                .data_components
31                .id_from_key(key)
32                .ok_or_else(|| std::io::Error::other(format!("Unknown component key: {key:?}")))?;
33
34            let entry = REGISTRY
35                .data_components
36                .by_id(id)
37                .ok_or_else(|| std::io::Error::other(format!("No entry for component id: {id}")))?;
38
39            VarInt(id as i32).write(writer)?;
40
41            let mut buf = Vec::new();
42            entry.write_network(data, &mut buf)?;
43            writer.write_all(&buf)?;
44        }
45
46        // Write removed component IDs
47        for key in removed {
48            let id = REGISTRY
49                .data_components
50                .id_from_key(key)
51                .ok_or_else(|| std::io::Error::other(format!("Unknown component key: {key:?}")))?;
52            VarInt(id as i32).write(writer)?;
53        }
54
55        Ok(())
56    }
57}
58
59impl ReadFrom for DataComponentPatch {
60    fn read(data: &mut Cursor<&[u8]>) -> Result<Self> {
61        use crate::{REGISTRY, RegistryExt};
62
63        let added_count = read_component_count(data, "added")?;
64        let removed_count = read_component_count(data, "removed")?;
65
66        log::info!("Reading DataComponentPatch: added={added_count}, removed={removed_count}");
67
68        let mut patch = Self::new();
69
70        // Read added components
71        for i in 0..added_count {
72            let pos_before = data.position();
73            let type_id = read_non_negative_varint(data, "component type id")?;
74
75            let key = REGISTRY
76                .data_components
77                .get_key_by_id(type_id)
78                .ok_or_else(|| {
79                    std::io::Error::other(format!("Unknown component type ID: {type_id}"))
80                })?
81                .clone();
82
83            log::info!("  [{i}] Reading component {key} (id={type_id}) at pos {pos_before}");
84
85            let entry = REGISTRY
86                .data_components
87                .by_id(type_id)
88                .ok_or_else(|| std::io::Error::other(format!("No entry for component: {key}")))?;
89
90            let component_data = entry.read_network(data).map_err(|e| {
91                log::error!("    Failed to read component {key}: {e}");
92                e
93            })?;
94
95            let pos_after = data.position();
96            log::info!("    Read {} bytes for {key}", pos_after - pos_before);
97
98            patch
99                .entries
100                .insert(key, ComponentPatchEntry::Set(component_data));
101        }
102
103        // Read removed component IDs
104        for _ in 0..removed_count {
105            let type_id = read_non_negative_varint(data, "component type id")?;
106
107            let key = REGISTRY
108                .data_components
109                .get_key_by_id(type_id)
110                .ok_or_else(|| {
111                    std::io::Error::other(format!("Unknown component type ID: {type_id}"))
112                })?
113                .clone();
114
115            patch.entries.insert(key, ComponentPatchEntry::Removed);
116        }
117
118        Ok(patch)
119    }
120}
121
122impl DataComponentPatch {
123    /// Reads a patch where each component value is prefixed with a `VarInt` byte length.
124    ///
125    /// Vanilla uses this for untrusted client packets (e.g., creative mode slot)
126    /// via `DataComponentPatch.DELIMITED_STREAM_CODEC`.
127    pub fn read_delimited(data: &mut Cursor<&[u8]>) -> Result<Self> {
128        use crate::{REGISTRY, RegistryExt};
129        use std::io::Read;
130
131        let added_count = read_component_count(data, "added")?;
132        let removed_count = read_component_count(data, "removed")?;
133
134        const MAX_COMPONENTS: usize = 65_536;
135        const MAX_COMPONENT_BYTES: usize = 2 * 1024 * 1024;
136
137        if added_count.saturating_add(removed_count) > MAX_COMPONENTS {
138            return Err(std::io::Error::other(format!(
139                "Component patch too large: {added_count} added + {removed_count} removed > {MAX_COMPONENTS}"
140            )));
141        }
142
143        let mut patch = Self::new();
144
145        for _ in 0..added_count {
146            let type_id = read_non_negative_varint(data, "component type id")?;
147            let byte_len = read_non_negative_varint(data, "component byte length")?;
148
149            if byte_len > MAX_COMPONENT_BYTES {
150                return Err(std::io::Error::other(format!(
151                    "Component data too large: {byte_len} bytes > {MAX_COMPONENT_BYTES}"
152                )));
153            }
154
155            let key = REGISTRY
156                .data_components
157                .get_key_by_id(type_id)
158                .ok_or_else(|| {
159                    std::io::Error::other(format!("Unknown component type ID: {type_id}"))
160                })?
161                .clone();
162
163            let entry = REGISTRY
164                .data_components
165                .by_id(type_id)
166                .ok_or_else(|| std::io::Error::other(format!("No entry for component: {key}")))?;
167
168            // Read the component bytes into a sub-buffer
169            let mut buf = vec![0u8; byte_len];
170            data.read_exact(&mut buf)?;
171
172            let mut sub_cursor = Cursor::new(buf.as_slice());
173            let component_data = entry.read_network(&mut sub_cursor)?;
174            patch
175                .entries
176                .insert(key, ComponentPatchEntry::Set(component_data));
177        }
178
179        for _ in 0..removed_count {
180            let type_id = read_non_negative_varint(data, "component type id")?;
181            let key = REGISTRY
182                .data_components
183                .get_key_by_id(type_id)
184                .ok_or_else(|| {
185                    std::io::Error::other(format!("Unknown component type ID: {type_id}"))
186                })?
187                .clone();
188            patch.entries.insert(key, ComponentPatchEntry::Removed);
189        }
190
191        Ok(patch)
192    }
193}
194
195fn read_component_count(data: &mut Cursor<&[u8]>, kind: &str) -> Result<usize> {
196    read_non_negative_varint(data, &format!("{kind} component count"))
197}
198
199fn read_non_negative_varint(data: &mut Cursor<&[u8]>, name: &str) -> Result<usize> {
200    let value = VarInt::read(data)?.0;
201    usize::try_from(value).map_err(|_| std::io::Error::other(format!("Negative {name}: {value}")))
202}