use crate::{raw, NbtCompound, NbtList, NbtTag};
use flate2::{
read::{GzDecoder, ZlibDecoder},
write::{GzEncoder, ZlibEncoder},
Compression,
};
use std::{
error::Error,
fmt::{self, Display, Formatter},
io::{self, Read, Write},
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Flavor {
Uncompressed,
ZlibCompressed,
ZlibCompressedWith(Compression),
GzCompressed,
GzCompressedWith(Compression),
}
pub fn read_nbt<R: Read>(
reader: &mut R,
flavor: Flavor,
) -> Result<(NbtCompound, String), NbtIoError> {
match flavor {
Flavor::Uncompressed => read_nbt_uncompressed(reader),
Flavor::ZlibCompressed | Flavor::ZlibCompressedWith(_) =>
read_nbt_uncompressed(&mut ZlibDecoder::new(reader)),
Flavor::GzCompressed | Flavor::GzCompressedWith(_) =>
read_nbt_uncompressed(&mut GzDecoder::new(reader)),
}
}
fn read_nbt_uncompressed<R: Read>(reader: &mut R) -> Result<(NbtCompound, String), NbtIoError> {
let root_id = raw::read_u8(reader)?;
if root_id != 0xA {
return Err(NbtIoError::TagTypeMismatch {
expected: 0xA,
found: root_id,
});
}
let root_name = raw::read_string(reader)?;
match read_tag_body_const::<_, 0xA>(reader) {
Ok(NbtTag::Compound(compound)) => Ok((compound, root_name)),
Err(e) => Err(e),
_ => unreachable!(),
}
}
fn read_tag_body_dyn<R: Read>(reader: &mut R, tag_id: u8) -> Result<NbtTag, NbtIoError> {
macro_rules! drive_reader {
($($id:literal)*) => {
match tag_id {
$( $id => read_tag_body_const::<_, $id>(reader), )*
_ => Err(NbtIoError::InvalidTagId(tag_id))
}
};
}
drive_reader!(0x1 0x2 0x3 0x4 0x5 0x6 0x7 0x8 0x9 0xA 0xB 0xC)
}
#[inline]
fn read_tag_body_const<R: Read, const TAG_ID: u8>(reader: &mut R) -> Result<NbtTag, NbtIoError> {
let tag = match TAG_ID {
0x1 => NbtTag::Byte(raw::read_i8(reader)?),
0x2 => NbtTag::Short(raw::read_i16(reader)?),
0x3 => NbtTag::Int(raw::read_i32(reader)?),
0x4 => NbtTag::Long(raw::read_i64(reader)?),
0x5 => NbtTag::Float(raw::read_f32(reader)?),
0x6 => NbtTag::Double(raw::read_f64(reader)?),
0x7 => {
let len = raw::read_i32(reader)? as usize;
let mut array = vec![0u8; len];
reader.read_exact(&mut array)?;
NbtTag::ByteArray(raw::cast_byte_buf_to_signed(array))
}
0x8 => NbtTag::String(raw::read_string(reader)?),
0x9 => {
let tag_id = raw::read_u8(reader)?;
let len = raw::read_i32(reader)? as usize;
if tag_id > 0xC || (tag_id == 0 && len > 0) {
return Err(NbtIoError::InvalidTagId(tag_id));
}
if len == 0 {
return Ok(NbtTag::List(NbtList::new()));
}
let mut list = NbtList::with_capacity(len);
macro_rules! drive_reader {
($($id:literal)*) => {
match tag_id {
$(
$id => {
for _ in 0 .. len {
list.push(read_tag_body_const::<_, $id>(reader)?);
}
},
)*
_ => return Err(NbtIoError::InvalidTagId(tag_id))
}
};
}
drive_reader!(0x1 0x2 0x3 0x4 0x5 0x6 0x7 0x8 0x9 0xA 0xB 0xC);
NbtTag::List(list)
}
0xA => {
let mut compound = NbtCompound::new();
let mut tag_id = raw::read_u8(reader)?;
while tag_id != 0x0 {
let name = raw::read_string(reader)?;
let tag = read_tag_body_dyn(reader, tag_id)?;
compound.insert(name, tag);
tag_id = raw::read_u8(reader)?;
}
NbtTag::Compound(compound)
}
0xB => {
let len = raw::read_i32(reader)? as usize;
NbtTag::IntArray(raw::read_i32_array(reader, len)?)
}
0xC => {
let len = raw::read_i32(reader)? as usize;
NbtTag::LongArray(raw::read_i64_array(reader, len)?)
}
_ => unreachable!("read_tag_body_const called with unchecked TAG_ID"),
};
Ok(tag)
}
pub fn write_nbt<W: Write>(
writer: &mut W,
root_name: Option<&str>,
root: &NbtCompound,
flavor: Flavor,
) -> Result<(), NbtIoError> {
let (mode, compression) = match flavor {
Flavor::Uncompressed => {
return write_nbt_uncompressed(writer, root_name, root);
}
Flavor::ZlibCompressed => (2, Compression::default()),
Flavor::ZlibCompressedWith(compression) => (2, compression),
Flavor::GzCompressed => (1, Compression::default()),
Flavor::GzCompressedWith(compression) => (1, compression),
};
if mode == 1 {
write_nbt_uncompressed(&mut GzEncoder::new(writer, compression), root_name, root)
} else {
write_nbt_uncompressed(&mut ZlibEncoder::new(writer, compression), root_name, root)
}
}
fn write_nbt_uncompressed<W>(
writer: &mut W,
root_name: Option<&str>,
root: &NbtCompound,
) -> Result<(), NbtIoError>
where
W: Write,
{
raw::write_u8(writer, 0xA)?;
raw::write_string(writer, root_name.unwrap_or(""))?;
for (name, tag) in root.inner() {
raw::write_u8(writer, raw::id_for_tag(Some(tag)))?;
raw::write_string(writer, name)?;
write_tag_body(writer, tag)?;
}
raw::write_u8(writer, raw::id_for_tag(None))?;
Ok(())
}
fn write_tag_body<W: Write>(writer: &mut W, tag: &NbtTag) -> Result<(), NbtIoError> {
match tag {
&NbtTag::Byte(value) => raw::write_i8(writer, value)?,
&NbtTag::Short(value) => raw::write_i16(writer, value)?,
&NbtTag::Int(value) => raw::write_i32(writer, value)?,
&NbtTag::Long(value) => raw::write_i64(writer, value)?,
&NbtTag::Float(value) => raw::write_f32(writer, value)?,
&NbtTag::Double(value) => raw::write_f64(writer, value)?,
NbtTag::ByteArray(value) => {
raw::write_i32(writer, value.len() as i32)?;
writer.write_all(raw::cast_bytes_to_unsigned(value.as_slice()))?;
}
NbtTag::String(value) => raw::write_string(writer, value)?,
NbtTag::List(value) =>
if value.is_empty() {
writer.write_all(&[raw::id_for_tag(None), 0, 0, 0, 0])?;
} else {
let list_type = raw::id_for_tag(Some(&value[0]));
raw::write_u8(writer, list_type)?;
raw::write_i32(writer, value.len() as i32)?;
for sub_tag in value.as_ref() {
let tag_id = raw::id_for_tag(Some(sub_tag));
if tag_id != list_type {
return Err(NbtIoError::NonHomogenousList {
list_type,
encountered_type: tag_id,
});
}
write_tag_body(writer, sub_tag)?;
}
},
NbtTag::Compound(value) => {
for (name, tag) in value.inner() {
raw::write_u8(writer, raw::id_for_tag(Some(tag)))?;
raw::write_string(writer, name)?;
write_tag_body(writer, tag)?;
}
raw::write_u8(writer, raw::id_for_tag(None))?;
}
NbtTag::IntArray(value) => {
raw::write_i32(writer, value.len() as i32)?;
for &int in value.iter() {
raw::write_i32(writer, int)?;
}
}
NbtTag::LongArray(value) => {
raw::write_i32(writer, value.len() as i32)?;
for &long in value.iter() {
raw::write_i64(writer, long)?;
}
}
}
Ok(())
}
#[derive(Debug)]
pub enum NbtIoError {
StdIo(io::Error),
MissingRootTag,
NonHomogenousList {
list_type: u8,
encountered_type: u8,
},
OptionInList,
MissingLength,
InvalidTagId(u8),
TagTypeMismatch {
expected: u8,
found: u8,
},
ExpectedSeq,
ExpectedEnum,
InvalidKey,
InvalidEnumVariant,
InvalidCesu8String,
UnsupportedType(&'static str),
Custom(Box<str>),
}
#[cfg(feature = "serde")]
impl serde::ser::Error for NbtIoError {
fn custom<T>(msg: T) -> Self
where T: Display {
NbtIoError::Custom(msg.to_string().into_boxed_str())
}
}
#[cfg(feature = "serde")]
impl serde::de::Error for NbtIoError {
fn custom<T>(msg: T) -> Self
where T: Display {
NbtIoError::Custom(msg.to_string().into_boxed_str())
}
}
impl From<io::Error> for NbtIoError {
fn from(error: io::Error) -> Self {
NbtIoError::StdIo(error)
}
}
impl Display for NbtIoError {
fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
match self {
NbtIoError::StdIo(error) => write!(f, "{}", error),
NbtIoError::MissingRootTag =>
write!(f, "NBT tree does not start with a valid root tag."),
&NbtIoError::NonHomogenousList {
list_type,
encountered_type,
} => write!(
f,
"Encountered non-homogenous list or sequential type: expected {:X} but found {:X}",
list_type, encountered_type
),
NbtIoError::OptionInList => write!(
f,
"Minecraft's NBT format cannot support options in sequential data structures"
),
NbtIoError::MissingLength => write!(
f,
"Sequential types must have an initial computable length to be serializable"
),
&NbtIoError::InvalidTagId(id) => write!(
f,
"Encountered invalid tag ID 0x{:X} during deserialization",
id
),
&NbtIoError::TagTypeMismatch { expected, found } => write!(
f,
"Tag type mismatch: expected 0x{:X} but found 0x{:X}",
expected, found
),
NbtIoError::ExpectedSeq => write!(f, "Expected sequential tag type (array)"),
NbtIoError::ExpectedEnum => write!(
f,
"Encountered invalid enum representation in the NBT tag tree"
),
NbtIoError::InvalidKey => write!(f, "Map keys must be a valid string"),
NbtIoError::InvalidEnumVariant =>
write!(f, "Encountered invalid enum variant while deserializing"),
NbtIoError::InvalidCesu8String => write!(f, "Encountered invalid CESU8 string"),
NbtIoError::UnsupportedType(ty) =>
write!(f, "Type {} is not supported by Minecraft's NBT format", ty),
NbtIoError::Custom(msg) => write!(f, "{}", msg),
}
}
}
impl Error for NbtIoError {}