Skip to main content

mcproto_types/
text_component.rs

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/// A network NBT representation of the current Java text component schema.
22#[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    // The stream parser requires a root name; network text components omit it.
131    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            // A declared length is untrusted and can exceed the remaining packet.
193            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}