From b030b9de9345d7b1cfef205e5b9a1e2c7dc6025e Mon Sep 17 00:00:00 2001 From: mat Date: Sat, 18 Dec 2021 20:33:16 -0600 Subject: [PATCH] nbt --- Cargo.lock | 1 + azalea-nbt/Cargo.toml | 1 + azalea-nbt/src/decode.rs | 14 +++ azalea-nbt/src/encode.rs | 197 +++++++++++++++++------------------ azalea-nbt/src/error.rs | 1 + azalea-nbt/src/lib.rs | 1 + azalea-nbt/src/tag.rs | 20 ++++ azalea-nbt/tests/decode.rs | 47 ++++++++- azalea-protocol/src/read.rs | 4 +- azalea-protocol/src/write.rs | 4 +- 10 files changed, 181 insertions(+), 109 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 6c46763e..01a1ed53 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -85,6 +85,7 @@ name = "azalea-nbt" version = "0.1.0" dependencies = [ "byteorder", + "flate2", "num-derive", "num-traits", ] diff --git a/azalea-nbt/Cargo.toml b/azalea-nbt/Cargo.toml index 38fa4123..2e9fdb3b 100644 --- a/azalea-nbt/Cargo.toml +++ b/azalea-nbt/Cargo.toml @@ -7,5 +7,6 @@ version = "0.1.0" [dependencies] byteorder = "^1.4.3" +flate2 = "^1.0.22" num-derive = "^0.3.3" num-traits = "^0.2.14" diff --git a/azalea-nbt/src/decode.rs b/azalea-nbt/src/decode.rs index 704c8e2a..f3fbae54 100644 --- a/azalea-nbt/src/decode.rs +++ b/azalea-nbt/src/decode.rs @@ -1,10 +1,12 @@ use crate::Error; use crate::Tag; use byteorder::{ReadBytesExt, BE}; +use flate2::read::{GzDecoder, ZlibDecoder}; use std::{collections::HashMap, io::Read}; impl Tag { fn read_known(stream: &mut impl Read, id: u8) -> Result { + println!("read_known: id={}", id); let tag = match id { // Signifies the end of a TAG_Compound. It is only ever used inside // a TAG_Compound, and is not named despite being in a TAG_Compound @@ -65,7 +67,9 @@ impl Tag { 10 => { let mut map = HashMap::new(); loop { + println!("compound loop"); let tag_id = stream.read_u8().unwrap_or(0); + println!("compound loop tag_id={}", tag_id); if tag_id == 0 { break; } @@ -108,4 +112,14 @@ impl Tag { // default to compound tag Tag::read_known(stream, 10) } + + pub fn read_zlib(stream: &mut impl Read) -> Result { + let mut gz = ZlibDecoder::new(stream); + Tag::read(&mut gz) + } + + pub fn read_gzip(stream: &mut impl Read) -> Result { + let mut gz = GzDecoder::new(stream); + Tag::read(&mut gz) + } } diff --git a/azalea-nbt/src/encode.rs b/azalea-nbt/src/encode.rs index 86c0d868..d787e15f 100644 --- a/azalea-nbt/src/encode.rs +++ b/azalea-nbt/src/encode.rs @@ -1,117 +1,112 @@ -use byteorder::{ReadBytesExt, BE}; -use error::Error; -use std::{collections::HashMap, io::Read}; -use tag::Tag; +use crate::Error; +use crate::Tag; +use byteorder::{WriteBytesExt, BE}; +use flate2::write::{GzEncoder, ZlibEncoder}; +use std::io::Write; impl Tag { - fn write(&self, stream: &mut impl Read) -> Result<(), Error> { - println!("read_known: id={}", id); - let tag = match id { - // Signifies the end of a TAG_Compound. It is only ever used inside - // a TAG_Compound, and is not named despite being in a TAG_Compound - 0 => Tag::End, - // A single signed byte - 1 => Tag::Byte(stream.read_i8().map_err(|_| Error::InvalidTag)?), - // A single signed, big endian 16 bit integer - 2 => Tag::Short(stream.read_i16::().map_err(|_| Error::InvalidTag)?), - // A single signed, big endian 32 bit integer - 3 => Tag::Int(stream.read_i32::().map_err(|_| Error::InvalidTag)?), - // A single signed, big endian 64 bit integer - 4 => Tag::Long(stream.read_i64::().map_err(|_| Error::InvalidTag)?), - // A single, big endian IEEE-754 single-precision floating point - // number (NaN possible) - 5 => Tag::Float(stream.read_f32::().map_err(|_| Error::InvalidTag)?), - // A single, big endian IEEE-754 double-precision floating point - // number (NaN possible) - 6 => Tag::Double(stream.read_f64::().map_err(|_| Error::InvalidTag)?), - // A length-prefixed array of signed bytes. The prefix is a signed - // integer (thus 4 bytes) - 7 => { - let length = stream.read_i32::().map_err(|_| Error::InvalidTag)?; - let mut bytes = Vec::new(); - for _ in 0..length { - bytes.push(stream.read_i8().map_err(|_| Error::InvalidTag)?); + pub fn write_without_end(&self, writer: &mut dyn Write) -> Result<(), Error> { + match self { + Tag::End => {} + Tag::Byte(value) => writer.write_i8(*value).map_err(|_| Error::WriteError)?, + Tag::Short(value) => writer + .write_i16::(*value) + .map_err(|_| Error::WriteError)?, + Tag::Int(value) => writer + .write_i32::(*value) + .map_err(|_| Error::WriteError)?, + Tag::Long(value) => writer + .write_i64::(*value) + .map_err(|_| Error::WriteError)?, + Tag::Float(value) => writer + .write_f32::(*value) + .map_err(|_| Error::WriteError)?, + Tag::Double(value) => writer + .write_f64::(*value) + .map_err(|_| Error::WriteError)?, + Tag::ByteArray(value) => { + writer + .write_i32::(value.len() as i32) + .map_err(|_| Error::WriteError)?; + for byte in value { + writer.write_i8(*byte).map_err(|_| Error::WriteError)?; } - Tag::ByteArray(bytes) } - // A length-prefixed modified UTF-8 string. The prefix is an - // unsigned short (thus 2 bytes) signifying the length of the - // string in bytes - 8 => { - let length = stream.read_u16::().map_err(|_| Error::InvalidTag)?; - let mut bytes = Vec::new(); - for _ in 0..length { - bytes.push(stream.read_u8().map_err(|_| Error::InvalidTag)?); + Tag::String(value) => { + writer + .write_i16::(value.len() as i16) + .map_err(|_| Error::WriteError)?; + writer + .write_all(value.as_bytes()) + .map_err(|_| Error::WriteError)?; + } + Tag::List(value) => { + // we just get the type from the first item, or default the type to END + let type_id = value.first().and_then(|f| Some(f.id())).unwrap_or(0); + writer.write_u8(type_id).map_err(|_| Error::WriteError)?; + writer + .write_i32::(value.len() as i32) + .map_err(|_| Error::WriteError)?; + for tag in value { + tag.write_without_end(writer)?; } - Tag::String(String::from_utf8(bytes).map_err(|_| Error::InvalidTag)?) } - // A list of nameless tags, all of the same type. The list is - // prefixed with the Type ID of the items it contains (thus 1 - // byte), and the length of the list as a signed integer (a further - // 4 bytes). If the length of the list is 0 or negative, the type - // may be 0 (TAG_End) but otherwise it must be any other type. (The - // notchian implementation uses TAG_End in that situation, but - // another reference implementation by Mojang uses 1 instead; - // parsers should accept any type if the length is <= 0). - 9 => { - let type_id = stream.read_u8().map_err(|_| Error::InvalidTag)?; - let length = stream.read_i32::().map_err(|_| Error::InvalidTag)?; - let mut list = Vec::new(); - for _ in 0..length { - list.push(Tag::read_known(stream, type_id)?); + Tag::Compound(value) => { + for (key, tag) in value { + writer.write_u8(tag.id()).map_err(|_| Error::WriteError)?; + Tag::String(key.clone()).write_without_end(writer)?; + tag.write_without_end(writer)?; } - Tag::List(list) + writer + .write_u8(Tag::End.id()) + .map_err(|_| Error::WriteError)?; } - // Effectively a list of a named tags. Order is not guaranteed. - 10 => { - println!("reading compound {{"); - let mut map = HashMap::new(); - loop { - let tag_id = stream.read_u8().unwrap_or(0); - println!("compound tag id: {}", tag_id); - if tag_id == 0 { - break; - } - let name = match Tag::read_known(stream, 8)? { - Tag::String(name) => name, - _ => panic!("Expected a string tag"), - }; - println!("compound name: {}", name); - let tag = Tag::read_known(stream, tag_id).map_err(|_| Error::InvalidTag)?; - println!("aight read tag: {:?}", tag); - map.insert(name, tag); + Tag::IntArray(value) => { + writer + .write_i32::(value.len() as i32) + .map_err(|_| Error::WriteError)?; + for int in value { + writer + .write_i32::(*int) + .map_err(|_| Error::WriteError)?; } - println!("}} compound map: {:?}", map); - Tag::Compound(map) } - // A length-prefixed array of signed integers. The prefix is a - // signed integer (thus 4 bytes) and indicates the number of 4 byte - // integers. - 11 => { - let length = stream.read_i32::().map_err(|_| Error::InvalidTag)?; - let mut ints = Vec::new(); - for _ in 0..length { - ints.push(stream.read_i32::().map_err(|_| Error::InvalidTag)?); + Tag::LongArray(value) => { + writer + .write_i32::(value.len() as i32) + .map_err(|_| Error::WriteError)?; + for long in value { + writer + .write_i64::(*long) + .map_err(|_| Error::WriteError)?; } - Tag::IntArray(ints) } - // A length-prefixed array of signed longs. The prefix is a signed - // integer (thus 4 bytes) and indicates the number of 8 byte longs. - 12 => { - let length = stream.read_i32::().map_err(|_| Error::InvalidTag)?; - let mut longs = Vec::new(); - for _ in 0..length { - longs.push(stream.read_i64::().map_err(|_| Error::InvalidTag)?); - } - Tag::LongArray(longs) - } - _ => return Err(Error::InvalidTagType(id)), - }; - Ok(tag) + } + + Ok(()) } - pub fn read(stream: &mut impl Read) -> Result { - // default to compound tag - Tag::read_known(stream, 10) + pub fn write(&self, writer: &mut impl Write) -> Result<(), Error> { + match self { + Tag::Compound(value) => { + for (key, tag) in value { + writer.write_u8(tag.id()).map_err(|_| Error::WriteError)?; + Tag::String(key.clone()).write_without_end(writer)?; + tag.write_without_end(writer)?; + } + Ok(()) + } + _ => Err(Error::InvalidTag), + } + } + + pub fn write_zlib(&self, writer: &mut impl Write) -> Result<(), Error> { + let mut encoder = ZlibEncoder::new(writer, flate2::Compression::default()); + self.write(&mut encoder) + } + + pub fn write_gzip(&self, writer: &mut impl Write) -> Result<(), Error> { + let mut encoder = GzEncoder::new(writer, flate2::Compression::default()); + self.write(&mut encoder) } } diff --git a/azalea-nbt/src/error.rs b/azalea-nbt/src/error.rs index 149a00e9..3ada7cf7 100644 --- a/azalea-nbt/src/error.rs +++ b/azalea-nbt/src/error.rs @@ -2,4 +2,5 @@ pub enum Error { InvalidTagType(u8), InvalidTag, + WriteError, } diff --git a/azalea-nbt/src/lib.rs b/azalea-nbt/src/lib.rs index adeb6e56..d14fd929 100644 --- a/azalea-nbt/src/lib.rs +++ b/azalea-nbt/src/lib.rs @@ -1,4 +1,5 @@ mod decode; +mod encode; mod error; mod tag; diff --git a/azalea-nbt/src/tag.rs b/azalea-nbt/src/tag.rs index 1f496c9d..7c59ea87 100644 --- a/azalea-nbt/src/tag.rs +++ b/azalea-nbt/src/tag.rs @@ -16,3 +16,23 @@ pub enum Tag { IntArray(Vec), // 11 LongArray(Vec), // 12 } + +impl Tag { + pub fn id(&self) -> u8 { + match self { + Tag::End => 0, + Tag::Byte(value) => 1, + Tag::Short(value) => 2, + Tag::Int(value) => 3, + Tag::Long(value) => 4, + Tag::Float(value) => 5, + Tag::Double(value) => 6, + Tag::ByteArray(value) => 7, + Tag::String(value) => 8, + Tag::List(value) => 9, + Tag::Compound(value) => 10, + Tag::IntArray(value) => 11, + Tag::LongArray(value) => 12, + } + } +} diff --git a/azalea-nbt/tests/decode.rs b/azalea-nbt/tests/decode.rs index f4060512..2c69745b 100644 --- a/azalea-nbt/tests/decode.rs +++ b/azalea-nbt/tests/decode.rs @@ -1,8 +1,15 @@ use azalea_nbt::Tag; -use std::collections::HashMap; +use flate2::{ + read::{GzDecoder, ZlibDecoder}, + write::{GzEncoder, ZlibEncoder}, +}; +use std::{ + collections::HashMap, + io::{Cursor, Read}, +}; #[test] -fn test_hello_world() { +fn test_decode_hello_world() { // read hello_world.nbt let mut file = std::fs::File::open("tests/hello_world.nbt").unwrap(); let tag = Tag::read(&mut file).unwrap(); @@ -17,3 +24,39 @@ fn test_hello_world() { )])) ); } + +#[test] +fn test_roundtrip_hello_world() { + let mut file = std::fs::File::open("tests/hello_world.nbt").unwrap(); + let mut original = Vec::new(); + file.read_to_end(&mut original).unwrap(); + + let mut original_stream = Cursor::new(original.clone()); + let tag = Tag::read(&mut original_stream).unwrap(); + + println!("ok read {:?}", tag); + + // write hello_world.nbt + let mut result = Cursor::new(Vec::new()); + tag.write(&mut result).unwrap(); + + assert_eq!(result.into_inner(), original); +} + +#[test] +fn test_bigtest() { + // read bigtest.nbt + let mut file = std::fs::File::open("tests/bigtest.nbt").unwrap(); + let mut original = Vec::new(); + file.read_to_end(&mut original).unwrap(); + + let mut original_stream = Cursor::new(original.clone()); + let original_tag = Tag::read_gzip(&mut original_stream).unwrap(); + + let mut result = Vec::new(); + original_tag.write(&mut result).unwrap(); + + let decoded_tag = Tag::read(&mut Cursor::new(result)).unwrap(); + + assert_eq!(decoded_tag, original_tag); +} diff --git a/azalea-protocol/src/read.rs b/azalea-protocol/src/read.rs index 7b94135a..3d6a41d6 100644 --- a/azalea-protocol/src/read.rs +++ b/azalea-protocol/src/read.rs @@ -1,8 +1,6 @@ -use std::io::Cursor; - use crate::{connect::PacketFlow, mc_buf::Readable, packets::ProtocolPacket}; use async_compression::tokio::bufread::ZlibDecoder; -use tokio::io::{AsyncRead, AsyncReadExt, BufReader}; +use tokio::io::{AsyncRead, AsyncReadExt}; async fn frame_splitter(stream: &mut R) -> Result, String> where diff --git a/azalea-protocol/src/write.rs b/azalea-protocol/src/write.rs index 4ae9f1c1..3898e74a 100644 --- a/azalea-protocol/src/write.rs +++ b/azalea-protocol/src/write.rs @@ -1,9 +1,7 @@ -use std::io::Read; - use crate::{mc_buf::Writable, packets::ProtocolPacket, read::MAXIMUM_UNCOMPRESSED_LENGTH}; use async_compression::tokio::bufread::ZlibEncoder; use tokio::{ - io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}, + io::{AsyncReadExt, AsyncWriteExt}, net::TcpStream, };