1use std::{
2 collections::HashMap,
3 error::Error,
4 fmt,
5 io::{self, Read, Write},
6};
7
8use fastnbt::{
9 SerOpts, Tag,
10 stream::{Parser, Value as StreamValue},
11};
12use mcproto_codec::{
13 error::{CodecError, CodecKind, CodecOperation, InvalidEncodingReason},
14 io::write_all_counted,
15};
16
17use crate::{TypeCodec, component::NbtComponent};
18
19pub use fastnbt::Value as NbtValue;
20
21#[derive(Debug, Clone, PartialEq)]
23pub struct TextComponent(pub NbtComponent);
24
25impl TextComponent {
26 pub const MAX_DECODE_BYTES: usize = 2 * 1024 * 1024;
27 pub const MAX_DECODE_NODES: usize = 262_144;
28 const MAX_COMPONENT_DEPTH: usize = 512;
29
30 pub fn text(value: impl Into<String>) -> Self {
31 Self(NbtComponent::text(value))
32 }
33}
34
35impl Default for TextComponent {
36 fn default() -> Self {
37 Self::text("")
38 }
39}
40
41impl From<NbtComponent> for TextComponent {
42 fn from(value: NbtComponent) -> Self {
43 Self(value)
44 }
45}
46
47impl From<String> for TextComponent {
48 fn from(value: String) -> Self {
49 Self::text(value)
50 }
51}
52
53impl From<&str> for TextComponent {
54 fn from(value: &str) -> Self {
55 Self::text(value)
56 }
57}
58
59impl TypeCodec for TextComponent {
60 fn encode(&self, writer: &mut impl Write) -> Result<(), CodecError> {
61 self.0
62 .validate_depth(Self::MAX_COMPONENT_DEPTH)
63 .map_err(|source| invalid_nbt(CodecOperation::Write, 0, source))?;
64 self.0
65 .validate_dynamic_depth(Self::MAX_COMPONENT_DEPTH)
66 .map_err(|source| invalid_nbt(CodecOperation::Write, 0, source))?;
67 let normalized = self.0.normalized_root_for_nbt();
68 let value = fastnbt::to_value(&normalized)
69 .map_err(|source| invalid_nbt(CodecOperation::Write, 0, source))?;
70 validate_nbt_strings(&value, CodecOperation::Write, 0)?;
71 encode_root(&value, writer)
72 }
73
74 fn decode(reader: &mut impl Read) -> Result<Self, CodecError> {
75 let (value, bytes_processed) = decode_root(reader)?;
76 let component: NbtComponent = fastnbt::from_value(&value)
77 .map_err(|source| invalid_nbt(CodecOperation::Read, bytes_processed, source))?;
78 component
79 .validate_depth(Self::MAX_COMPONENT_DEPTH)
80 .map_err(|source| invalid_nbt(CodecOperation::Read, bytes_processed, source))?;
81 component
82 .validate_dynamic_depth(Self::MAX_COMPONENT_DEPTH)
83 .map_err(|source| invalid_nbt(CodecOperation::Read, bytes_processed, source))?;
84 Ok(Self(component.normalized_root_for_nbt()))
85 }
86}
87
88fn encode_root(value: &NbtValue, writer: &mut impl Write) -> Result<(), CodecError> {
89 let wrapper = HashMap::from([("", value)]);
90 let encoded = fastnbt::to_bytes_with_opts(&wrapper, SerOpts::network_nbt())
91 .map_err(|source| invalid_nbt(CodecOperation::Write, 0, source))?;
92
93 if encoded.len() < 5
94 || encoded[0] != Tag::Compound as u8
95 || encoded[2..4] != [0, 0]
96 || encoded.last() != Some(&(Tag::End as u8))
97 || !matches!(encoded[1], 8 | 10)
98 {
99 return Err(CodecError::invalid_encoding_for_operation(
100 CodecKind::TextComponent,
101 CodecOperation::Write,
102 0,
103 InvalidEncodingReason::InvalidNbt,
104 ));
105 }
106
107 write_all_counted(writer, &encoded[1..2], CodecKind::TextComponent, 0)?;
108 write_all_counted(
109 writer,
110 &encoded[4..encoded.len() - 1],
111 CodecKind::TextComponent,
112 1,
113 )
114}
115
116fn decode_root(reader: &mut impl Read) -> Result<(NbtValue, usize), CodecError> {
117 let mut reader = CountedReader::new(reader, TextComponent::MAX_DECODE_BYTES);
118 let mut root_tag = [0_u8; 1];
119 reader.read_exact(&mut root_tag).map_err(|source| {
120 CodecError::from_read_error(CodecKind::TextComponent, reader.processed, source)
121 })?;
122 if !matches!(root_tag[0], 8 | 10) {
123 return Err(CodecError::invalid_encoding(
124 CodecKind::TextComponent,
125 reader.processed,
126 InvalidEncodingReason::InvalidTextComponentRootTag { tag: root_tag[0] },
127 ));
128 }
129
130 let mut prefixed = PrefixReader {
132 prefix: &[root_tag[0], 0, 0],
133 inner: &mut reader,
134 };
135 let result = (|| {
136 let mut parser = Parser::new(&mut prefixed);
137 let first = parser.next().map_err(NbtTreeError::Parser)?;
138 let mut remaining_nodes = TextComponent::MAX_DECODE_NODES;
139 parse_stream_value(&mut parser, first, 1, &mut remaining_nodes)
140 })();
141
142 match result {
143 Ok(value) => Ok((value, prefixed.inner.processed)),
144 Err(error) => {
145 let processed = prefixed.inner.processed;
146 if prefixed.inner.limit_exceeded {
147 return Err(CodecError::invalid_encoding(
148 CodecKind::TextComponent,
149 processed,
150 InvalidEncodingReason::TooLong {
151 max_bytes: TextComponent::MAX_DECODE_BYTES,
152 },
153 ));
154 }
155 match prefixed.inner.failure.take() {
156 Some(source) => Err(CodecError::from_read_error(
157 CodecKind::TextComponent,
158 processed,
159 source,
160 )),
161 None => Err(invalid_nbt(CodecOperation::Read, processed, error)),
162 }
163 }
164 }
165}
166
167fn parse_stream_value<R: Read>(
168 parser: &mut Parser<R>,
169 value: StreamValue,
170 depth: usize,
171 remaining_nodes: &mut usize,
172) -> Result<NbtValue, NbtTreeError> {
173 *remaining_nodes = remaining_nodes
174 .checked_sub(1)
175 .ok_or(NbtTreeError::TooManyNodes)?;
176 if depth > TextComponent::MAX_COMPONENT_DEPTH {
177 return Err(NbtTreeError::TooDeep);
178 }
179 match value {
180 StreamValue::Byte(_, value) => Ok(NbtValue::Byte(value)),
181 StreamValue::Short(_, value) => Ok(NbtValue::Short(value)),
182 StreamValue::Int(_, value) => Ok(NbtValue::Int(value)),
183 StreamValue::Long(_, value) => Ok(NbtValue::Long(value)),
184 StreamValue::Float(_, value) => Ok(NbtValue::Float(value)),
185 StreamValue::Double(_, value) => Ok(NbtValue::Double(value)),
186 StreamValue::ByteArray(_, value) => Ok(NbtValue::ByteArray(fastnbt::ByteArray::new(value))),
187 StreamValue::String(_, value) => Ok(NbtValue::String(value)),
188 StreamValue::IntArray(_, value) => Ok(NbtValue::IntArray(fastnbt::IntArray::new(value))),
189 StreamValue::LongArray(_, value) => Ok(NbtValue::LongArray(fastnbt::LongArray::new(value))),
190 StreamValue::List(_, _, length) => {
191 let length = usize::try_from(length).map_err(|_| NbtTreeError::InvalidStructure)?;
192 let mut values = Vec::with_capacity(length.min(1024));
194 for _ in 0..length {
195 let value = parser.next().map_err(NbtTreeError::Parser)?;
196 values.push(parse_stream_value(
197 parser,
198 value,
199 depth + 1,
200 remaining_nodes,
201 )?);
202 }
203 if !matches!(parser.next(), Ok(StreamValue::ListEnd)) {
204 return Err(NbtTreeError::InvalidStructure);
205 }
206 Ok(NbtValue::List(values))
207 }
208 StreamValue::Compound(_) => {
209 let mut values = HashMap::new();
210 loop {
211 let value = parser.next().map_err(NbtTreeError::Parser)?;
212 if matches!(value, StreamValue::CompoundEnd) {
213 break;
214 }
215 let name = stream_name(&value).ok_or(NbtTreeError::InvalidStructure)?;
216 values.insert(
217 name.to_owned(),
218 parse_stream_value(parser, value, depth + 1, remaining_nodes)?,
219 );
220 }
221 Ok(NbtValue::Compound(values))
222 }
223 StreamValue::ListEnd | StreamValue::CompoundEnd => Err(NbtTreeError::InvalidStructure),
224 }
225}
226
227fn stream_name(value: &StreamValue) -> Option<&str> {
228 match value {
229 StreamValue::Byte(name, _)
230 | StreamValue::Short(name, _)
231 | StreamValue::Int(name, _)
232 | StreamValue::Long(name, _)
233 | StreamValue::Float(name, _)
234 | StreamValue::Double(name, _)
235 | StreamValue::ByteArray(name, _)
236 | StreamValue::String(name, _)
237 | StreamValue::List(name, _, _)
238 | StreamValue::Compound(name)
239 | StreamValue::IntArray(name, _)
240 | StreamValue::LongArray(name, _) => name.as_deref(),
241 StreamValue::ListEnd | StreamValue::CompoundEnd => None,
242 }
243}
244
245#[derive(Debug)]
246enum NbtTreeError {
247 Parser(fastnbt::stream::Error),
248 InvalidStructure,
249 TooDeep,
250 TooManyNodes,
251}
252
253impl fmt::Display for NbtTreeError {
254 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
255 match self {
256 Self::Parser(source) => source.fmt(formatter),
257 Self::InvalidStructure => formatter.write_str("invalid NBT tree structure"),
258 Self::TooDeep => formatter.write_str("NBT tree is nested too deeply"),
259 Self::TooManyNodes => formatter.write_str("NBT tree contains too many nodes"),
260 }
261 }
262}
263
264impl Error for NbtTreeError {
265 fn source(&self) -> Option<&(dyn Error + 'static)> {
266 match self {
267 Self::Parser(source) => Some(source),
268 _ => None,
269 }
270 }
271}
272
273fn invalid_nbt(
274 operation: CodecOperation,
275 bytes_processed: usize,
276 source: impl Error + Send + Sync + 'static,
277) -> CodecError {
278 CodecError::invalid_encoding_for_operation_with_source(
279 CodecKind::TextComponent,
280 operation,
281 bytes_processed,
282 InvalidEncodingReason::InvalidNbt,
283 source,
284 )
285}
286
287fn validate_nbt_strings(
288 root: &NbtValue,
289 operation: CodecOperation,
290 bytes_processed: usize,
291) -> Result<(), CodecError> {
292 let mut pending = vec![root];
293 while let Some(value) = pending.pop() {
294 match value {
295 NbtValue::String(value) => validate_nbt_string(value, operation, bytes_processed)?,
296 NbtValue::List(values) => pending.extend(values),
297 NbtValue::Compound(values) => {
298 for (name, value) in values {
299 validate_nbt_string(name, operation, bytes_processed)?;
300 pending.push(value);
301 }
302 }
303 _ => {}
304 }
305 }
306 Ok(())
307}
308
309fn validate_nbt_string(
310 value: &str,
311 operation: CodecOperation,
312 bytes_processed: usize,
313) -> Result<(), CodecError> {
314 let encoded_len = value.chars().try_fold(0_usize, |length, character| {
315 let width = match character as u32 {
316 0 => 2,
317 1..=0x7f => 1,
318 0x80..=0x7ff => 2,
319 0x800..=0xffff => 3,
320 _ => 6,
321 };
322 length.checked_add(width)
323 });
324 if !matches!(encoded_len, Some(length) if length <= u16::MAX as usize) {
325 return Err(CodecError::invalid_encoding_for_operation(
326 CodecKind::TextComponent,
327 operation,
328 bytes_processed,
329 InvalidEncodingReason::StringTooLong {
330 max_bytes: u16::MAX as usize,
331 },
332 ));
333 }
334 Ok(())
335}
336
337struct CountedReader<'a, R: ?Sized> {
338 inner: &'a mut R,
339 processed: usize,
340 byte_limit: usize,
341 limit_exceeded: bool,
342 failure: Option<io::Error>,
343}
344
345impl<'a, R: ?Sized> CountedReader<'a, R> {
346 fn new(inner: &'a mut R, byte_limit: usize) -> Self {
347 Self {
348 inner,
349 processed: 0,
350 byte_limit,
351 limit_exceeded: false,
352 failure: None,
353 }
354 }
355}
356
357impl<R: Read + ?Sized> Read for CountedReader<'_, R> {
358 fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
359 if buffer.is_empty() {
360 return Ok(0);
361 }
362 let remaining = self.byte_limit.saturating_sub(self.processed);
363 if remaining == 0 {
364 self.limit_exceeded = true;
365 return Err(io::Error::other("NBT byte limit exceeded"));
366 }
367 let buffer_len = buffer.len().min(remaining);
368 match self.inner.read(&mut buffer[..buffer_len]) {
369 Ok(0) => {
370 let error = io::Error::new(io::ErrorKind::UnexpectedEof, "unexpected end of NBT");
371 self.failure = Some(error);
372 Err(io::Error::new(
373 io::ErrorKind::UnexpectedEof,
374 "unexpected end of NBT",
375 ))
376 }
377 Ok(read) => {
378 self.processed += read;
379 Ok(read)
380 }
381 Err(error) if error.kind() == io::ErrorKind::Interrupted => Err(error),
382 Err(error) => {
383 let returned = clone_io_error(&error);
384 self.failure = Some(error);
385 Err(returned)
386 }
387 }
388 }
389}
390
391struct PrefixReader<'a, 'b, R: ?Sized> {
392 prefix: &'a [u8],
393 inner: &'a mut CountedReader<'b, R>,
394}
395
396impl<R: Read + ?Sized> Read for PrefixReader<'_, '_, R> {
397 fn read(&mut self, buffer: &mut [u8]) -> io::Result<usize> {
398 if buffer.is_empty() {
399 return Ok(0);
400 }
401 if !self.prefix.is_empty() {
402 let read = buffer.len().min(self.prefix.len());
403 buffer[..read].copy_from_slice(&self.prefix[..read]);
404 self.prefix = &self.prefix[read..];
405 return Ok(read);
406 }
407 self.inner.read(buffer)
408 }
409}
410
411fn clone_io_error(error: &io::Error) -> io::Error {
412 match error.raw_os_error() {
413 Some(code) => io::Error::from_raw_os_error(code),
414 None => io::Error::new(error.kind(), error.to_string()),
415 }
416}