use bitbuffer::{
BitReadBuffer, BitReadStream, BitWrite, BitWriteSized, BitWriteStream, LittleEndian,
};
use num_traits::{PrimInt, Unsigned};
use serde::{Deserialize, Serialize};
use snap::raw::{decompress_len, Decoder};
use crate::demo::lzss::decompress;
use crate::demo::packet::stringtable::{
ExtraData, FixedUserDataSize, StringTable, StringTableEntry,
};
use crate::demo::parser::{Encode, ParseBitSkip};
use crate::{Parse, ParseError, ParserState, ReadResult, Result, Stream};
use std::borrow::Cow;
use std::cmp::min;
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
#[serde(bound(deserialize = "'a: 'static"))]
pub struct CreateStringTableMessage<'a> {
pub table: StringTable<'a>,
}
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct StringTableMeta {
pub max_entries: u16,
pub fixed_userdata_size: Option<FixedUserDataSize>,
}
impl From<&StringTable<'_>> for StringTableMeta {
fn from(table: &StringTable) -> Self {
StringTableMeta {
max_entries: table.max_entries,
fixed_userdata_size: table.fixed_user_data_size,
}
}
}
impl<'a> Parse<'a> for CreateStringTableMessage<'a> {
fn parse(stream: &mut Stream<'a>, state: &ParserState) -> Result<Self> {
let name = stream.read()?;
let max_entries: u16 = stream.read()?;
let encode_bits = log_base2(max_entries);
let entity_count: u16 = stream.read_sized(encode_bits as usize + 1)?;
let length = if state.protocol_version > 23 {
read_var_int(stream)?
} else {
stream.read_sized(20)?
};
let fixed_userdata_size = stream.read()?;
let compressed = stream.read()?;
let mut table_data = stream.read_bits(length as usize)?;
if compressed {
let decompressed_size: u32 = table_data.read()?;
let compressed_size: u32 = table_data.read()?;
if !(4..=10 * 1024 * 1024).contains(&compressed_size) {
return Err(ParseError::InvalidDemo(
"Invalid compressed string table size",
));
}
if decompressed_size > 100 * 1024 * 1024 {
return Err(ParseError::InvalidDemo(
"Invalid decompressed string table size",
));
}
let magic = table_data.read_string(Some(4))?;
match magic.as_ref() {
"SNAP" => {
let compressed_data = table_data.read_bytes(compressed_size as usize - 4)?;
let mut decoder = Decoder::new();
let decompressed_size_from_header = decompress_len(&compressed_data)?;
if decompressed_size_from_header != decompressed_size as usize {
return Err(ParseError::UnexpectedDecompressedSize {
expected: decompressed_size,
size: decompressed_size_from_header as u32,
});
}
let mut decompressed_data = vec![0; decompressed_size_from_header];
decoder
.decompress(&compressed_data, &mut decompressed_data)
.map_err(ParseError::from)?;
let buffer = BitReadBuffer::new_owned(decompressed_data, LittleEndian);
table_data = BitReadStream::new(buffer);
}
"LZSS" => {
let compressed_data = table_data.read_bytes(compressed_size as usize - 4)?;
let mut decompressed_data = Vec::with_capacity(decompressed_size as usize);
decompress(&compressed_data, &mut decompressed_data);
if decompressed_data.len() != decompressed_size as usize {
return Err(ParseError::UnexpectedDecompressedSize {
expected: decompressed_size,
size: decompressed_data.len() as u32,
});
}
let buffer = BitReadBuffer::new_owned(decompressed_data, LittleEndian);
table_data = BitReadStream::new(buffer);
}
_ => {
return Err(ParseError::UnexpectedCompressionType(magic.into_owned()));
}
}
}
let table_meta = StringTableMeta {
max_entries,
fixed_userdata_size,
};
let entries = parse_string_table_update(&mut table_data, &table_meta, entity_count)?;
let table = StringTable {
entries,
max_entries,
fixed_user_data_size: fixed_userdata_size,
client_entries: None,
compressed,
name,
};
Ok(CreateStringTableMessage { table })
}
}
impl<'a> ParseBitSkip<'a> for CreateStringTableMessage<'a> {
fn parse_skip(stream: &mut Stream<'a>, state: &ParserState) -> Result<()> {
let _: String = stream.read()?;
let max_entries: u16 = stream.read()?;
let encode_bits = log_base2(max_entries);
let _: u16 = stream.read_sized(encode_bits as usize + 1)?;
let length = if state.protocol_version > 23 {
read_var_int(stream)?
} else {
stream.read_sized(20)?
};
let _: Option<FixedUserDataSize> = stream.read()?;
let _: bool = stream.read()?;
stream.skip_bits(length as usize).map_err(ParseError::from)
}
}
impl Encode for CreateStringTableMessage<'_> {
fn encode(
&self,
stream: &mut BitWriteStream<LittleEndian>,
_state: &ParserState,
) -> Result<()> {
let table = &self.table;
table.name.write(stream)?;
table.max_entries.write(stream)?;
let encode_bits = log_base2(table.max_entries) as usize;
(table.entries.len() as u16).write_sized(stream, encode_bits + 1)?;
stream.reserve_int::<ParseError, _>(40, |stream| {
table.fixed_user_data_size.is_some().write(stream)?;
if let Some(fixed_size) = table.fixed_user_data_size {
fixed_size.write(stream)?;
}
false.write(stream)?;
let start = stream.bit_len();
let table_meta = table.get_table_meta();
write_string_table_update(&table.entries, stream, &table_meta)?;
let end = stream.bit_len();
Ok(encode_var_int_fixed((end - start) as u32))
})?;
Ok(())
}
}
#[test]
fn test_create_string_table_roundtrip() {
let state = ParserState::new(24, |_| false, false);
crate::test_roundtrip_encode(
CreateStringTableMessage {
table: StringTable {
name: "table1".into(),
entries: vec![],
max_entries: 16,
fixed_user_data_size: None,
client_entries: None,
compressed: false,
},
},
&state,
);
crate::test_roundtrip_encode(
CreateStringTableMessage {
table: StringTable {
name: "table1".into(),
entries: vec![
(
0,
StringTableEntry {
text: Some("foo".into()),
extra_data: None,
},
),
(
1,
StringTableEntry {
text: Some("bar".into()),
extra_data: None,
},
),
],
max_entries: 16,
fixed_user_data_size: Some(FixedUserDataSize { size: 12, bits: 4 }),
client_entries: None,
compressed: false,
},
},
&state,
);
}
#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)]
#[serde(bound(deserialize = "'a: 'static"))]
pub struct UpdateStringTableMessage<'a> {
pub entries: Vec<(u16, StringTableEntry<'a>)>,
pub table_id: u8,
}
impl<'a> Parse<'a> for UpdateStringTableMessage<'a> {
fn parse(stream: &mut Stream<'a>, state: &ParserState) -> Result<Self> {
let table_id = stream.read_sized(5)?;
let changed: u16 = if stream.read()? { stream.read()? } else { 1 };
let length: u32 = stream.read_int(20)?;
let mut data = stream.read_bits(length as usize)?;
let entries = match state.string_tables.get(table_id as usize) {
Some(table) => parse_string_table_update(&mut data, table, changed),
None => return Err(ParseError::StringTableNotFound(table_id)),
}?;
Ok(UpdateStringTableMessage { entries, table_id })
}
}
impl<'a> ParseBitSkip<'a> for UpdateStringTableMessage<'a> {
fn parse_skip(stream: &mut Stream<'a>, _state: &ParserState) -> Result<()> {
let _: u8 = stream.read_sized(5)?;
let _: u16 = if stream.read()? { stream.read()? } else { 1 };
let length: u32 = stream.read_int(20)?;
stream.skip_bits(length as usize).map_err(ParseError::from)
}
}
impl Encode for UpdateStringTableMessage<'_> {
fn encode(&self, stream: &mut BitWriteStream<LittleEndian>, state: &ParserState) -> Result<()> {
self.table_id.write_sized(stream, 5)?;
if self.entries.len() == 1 {
false.write(stream)?;
} else {
true.write(stream)?;
(self.entries.len() as u16).write(stream)?;
}
match state.string_tables.get(self.table_id as usize) {
Some(table) => Ok(stream.reserve_length(20, |stream| {
write_string_table_update(&self.entries, stream, table)
})?),
None => Err(ParseError::StringTableNotFound(self.table_id)),
}
}
}
#[test]
fn test_update_string_table_roundtrip() {
let mut state = ParserState::new(24, |_| false, false);
state.string_tables = vec![StringTableMeta {
max_entries: 16,
fixed_userdata_size: None,
}];
crate::test_roundtrip_encode(
UpdateStringTableMessage {
entries: vec![],
table_id: 0,
},
&state,
);
crate::test_roundtrip_encode(
UpdateStringTableMessage {
entries: vec![(
2,
StringTableEntry {
text: Some("foo".into()),
extra_data: None,
},
)],
table_id: 0,
},
&state,
);
crate::test_roundtrip_encode(
UpdateStringTableMessage {
entries: vec![
(
2,
StringTableEntry {
text: Some("foo".into()),
extra_data: None,
},
),
(
3,
StringTableEntry {
text: Some("bar".into()),
extra_data: None,
},
),
],
table_id: 0,
},
&state,
);
}
struct TableEntries<'a> {
entries: Vec<(u16, StringTableEntry<'a>)>,
history: Vec<u16>,
}
impl<'a> TableEntries<'a> {
pub fn new(count: usize) -> Self {
TableEntries {
entries: Vec::with_capacity(min(count, 128)),
history: Vec::with_capacity(32),
}
}
pub fn push(&mut self, entry: (u16, StringTableEntry<'a>)) {
if self.history.len() > 31 {
self.history.remove(0);
}
let entry_index = self.entries.len();
self.entries.push(entry);
self.history.push(entry_index as u16);
}
pub fn get_history(&self, index: usize) -> Option<&StringTableEntry<'a>> {
self.history
.get(index)
.and_then(|entry_index| self.entries.get(*entry_index as usize))
.map(|entry| &entry.1)
}
pub fn into_entries(self) -> Vec<(u16, StringTableEntry<'a>)> {
self.entries
}
pub fn find_best_history(&self, text: &str) -> Option<(usize, usize)> {
let mut best_index = None;
let mut best_count = 0;
for (history_index, entry_index) in self.history.iter().enumerate() {
if let Some((_, entry)) = self.entries.get(*entry_index as usize) {
let similar = min(31, count_similar_characters(entry.text(), text));
if similar >= 3 && similar > best_count {
best_index = Some(history_index);
best_count = similar;
}
}
}
best_index.map(|index| (index, best_count))
}
}
fn count_similar_characters(a: &str, b: &str) -> usize {
for (i, (a, b)) in a.bytes().zip(b.bytes()).enumerate() {
if a != b {
return i;
}
}
min(a.len(), b.len())
}
pub fn parse_string_table_update<'a>(
stream: &mut Stream<'a>,
table_meta: &StringTableMeta,
entry_count: u16,
) -> ReadResult<Vec<(u16, StringTableEntry<'a>)>> {
let entry_bits = log_base2(table_meta.max_entries);
let mut entries = TableEntries::new(entry_count as usize);
let mut last_entry: i16 = -1;
for _ in 0..entry_count {
let index = if stream.read()? {
last_entry.saturating_add(1) as u16
} else {
stream.read_sized(entry_bits as usize)?
};
last_entry = index as i16;
let entry = read_table_entry(stream, table_meta, &entries)?;
entries.push((index, entry));
}
Ok(entries.into_entries())
}
pub fn write_string_table_update(
entries: &[(u16, StringTableEntry)],
stream: &mut BitWriteStream<LittleEndian>,
table_meta: &StringTableMeta,
) -> ReadResult<()> {
let entry_bits = log_base2(table_meta.max_entries);
let mut last_entry: i16 = -1;
let mut history = TableEntries::new(entries.len());
for (index, entry) in entries.iter() {
let index = *index as i16;
if index == (last_entry + 1) {
true.write(stream)?;
} else {
false.write(stream)?;
index.write_sized(stream, entry_bits as usize)?;
}
last_entry = index;
write_table_entry(entry, stream, table_meta, &history)?;
history.push((index as u16, entry.clone()));
}
Ok(())
}
#[test]
fn test_table_update_roundtrip() {
fn entry_roundtrip(
entries: Vec<(u16, StringTableEntry)>,
max_entries: u16,
fixed_bits: Option<u8>,
) {
let table_meta = StringTableMeta {
max_entries,
fixed_userdata_size: fixed_bits.map(|bits| FixedUserDataSize { size: 0, bits }),
};
let mut data = Vec::new();
let pos = {
let mut write = BitWriteStream::new(&mut data, LittleEndian);
write_string_table_update(&entries, &mut write, &table_meta).unwrap();
write.bit_len()
};
let mut read = BitReadStream::new(BitReadBuffer::new(&data, LittleEndian));
assert_eq!(
entries,
parse_string_table_update(&mut read, &table_meta, entries.len() as u16).unwrap()
);
assert_eq!(pos, read.pos());
}
entry_roundtrip(
vec![(
3,
StringTableEntry {
text: None,
extra_data: None,
},
)],
8,
None,
);
entry_roundtrip(
vec![
(
0,
StringTableEntry {
text: Some("bar".into()),
extra_data: None,
},
),
(
1,
StringTableEntry {
text: Some("foo".into()),
extra_data: None,
},
),
(
5,
StringTableEntry {
text: Some("asd".into()),
extra_data: None,
},
),
],
16,
None,
);
entry_roundtrip(
vec![
(
1,
StringTableEntry {
text: Some("foo".into()),
extra_data: None,
},
),
(
2,
StringTableEntry {
text: Some("asd".into()),
extra_data: None,
},
),
],
16,
None,
);
entry_roundtrip(
vec![(
1,
StringTableEntry {
text: Some("foo".into()),
extra_data: None,
},
)],
16,
None,
);
}
fn read_table_entry<'a>(
stream: &mut Stream<'a>,
table_meta: &StringTableMeta,
history: &TableEntries,
) -> ReadResult<StringTableEntry<'a>> {
let text = if stream.read()? {
if stream.read()? {
let index: u16 = stream.read_sized(5)?;
let bytes_to_copy: u32 = stream.read_sized(5)?;
let rest_of_string: Cow<str> = stream.read()?;
Some(
match history
.get_history(index as usize)
.and_then(|entry| entry.text.as_ref())
{
Some(text) => Cow::Owned(String::from_utf8({
text.bytes()
.take(bytes_to_copy as usize)
.chain(rest_of_string.bytes())
.collect()
})?),
None => rest_of_string, },
)
} else {
Some(stream.read()?)
}
} else {
None
};
let extra_data = if stream.read()? {
Some(match table_meta.fixed_userdata_size {
Some(size) => stream.read_bits(size.bits as usize)?,
None => {
let bytes: u16 = stream.read_sized(14)?;
stream.read_bits(bytes as usize * 8)?
}
})
} else {
None
}
.map(ExtraData::new);
Ok(StringTableEntry { text, extra_data })
}
fn write_table_entry(
entry: &StringTableEntry,
stream: &mut BitWriteStream<LittleEndian>,
table_meta: &StringTableMeta,
history: &TableEntries,
) -> ReadResult<()> {
entry.text.is_some().write(stream)?;
if let Some(text) = entry.text.as_deref() {
let history_item = history.find_best_history(text);
history_item.is_some().write(stream)?;
if let Some((history_index, history_count)) = history_item {
history_index.write_sized(stream, 5)?;
history_count.write_sized(stream, 5)?;
let diff_bytes = &text.as_bytes()[history_count..];
stream.write_bytes(diff_bytes)?;
0u8.write(stream)?; } else {
text.write(stream)?;
}
}
entry.extra_data.is_some().write(stream)?;
if let Some(extra_data) = entry.extra_data.as_ref() {
match table_meta.fixed_userdata_size {
Some(size) => {
extra_data.data.write_sized(stream, size.bits as usize)?;
}
None => {
extra_data.byte_len.write_sized(stream, 14)?;
extra_data
.data
.write_sized(stream, extra_data.byte_len as usize * 8)?;
}
}
}
Ok(())
}
#[test]
fn test_table_entry_roundtrip() {
fn entry_roundtrip(entry: StringTableEntry, fixed_bits: Option<u8>) {
let table_meta = StringTableMeta {
max_entries: 0,
fixed_userdata_size: fixed_bits.map(|bits| FixedUserDataSize { size: 0, bits }),
};
let mut data = Vec::new();
let pos = {
let history = TableEntries::new(1);
let mut write = BitWriteStream::new(&mut data, LittleEndian);
write_table_entry(&entry, &mut write, &table_meta, &history).unwrap();
write.bit_len()
};
let mut read = BitReadStream::new(BitReadBuffer::new(&data, LittleEndian));
assert_eq!(
entry,
read_table_entry(&mut read, &table_meta, &TableEntries::new(0)).unwrap()
);
assert_eq!(pos, read.pos());
}
entry_roundtrip(
StringTableEntry {
text: None,
extra_data: None,
},
None,
);
entry_roundtrip(
StringTableEntry {
text: Some("foo".into()),
extra_data: None,
},
None,
);
entry_roundtrip(
StringTableEntry {
text: None,
extra_data: Some(ExtraData::new(BitReadStream::new(
BitReadBuffer::new_owned(vec![0x55], LittleEndian),
))),
},
None,
);
entry_roundtrip(
StringTableEntry {
text: None,
extra_data: Some(ExtraData::new(BitReadStream::new(
BitReadBuffer::new_owned(vec![0x55; 128], LittleEndian),
))),
},
None,
);
entry_roundtrip(
StringTableEntry {
text: None,
extra_data: Some(ExtraData::new(BitReadStream::new(
BitReadBuffer::new_owned(vec![0x55; 4], LittleEndian),
))),
},
Some(4 * 8),
);
}
pub fn read_var_int(stream: &mut Stream) -> ReadResult<u32> {
let mut result: u32 = 0;
for i in (0..35u32).step_by(7) {
let byte: u8 = stream.read()?;
result |= (byte as u32 & 0x7F) << i;
if (byte >> 7) == 0 {
break;
}
}
Ok(result)
}
pub fn write_var_int(mut int: u32, stream: &mut BitWriteStream<LittleEndian>) -> ReadResult<()> {
while int > 0x7F {
let byte: u8 = int as u8 & 0x7F;
(byte | 0x80).write(stream)?;
int >>= 7;
}
(int as u8).write(stream)
}
pub fn encode_var_int_fixed(mut int: u32) -> u64 {
let mut out = 0;
for i in 0..4 {
let byte: u8 = int as u8 & 0x7F;
out |= ((byte | 0x80) as u64) << (i * 8);
int >>= 7;
}
out |= (int as u64) << 32;
out
}
#[test]
fn test_var_int_roundtrip() {
fn var_int_roundtrip(int: u32) {
let mut data = Vec::new();
let pos = {
let mut write = BitWriteStream::new(&mut data, LittleEndian);
write_var_int(int, &mut write).unwrap();
write.bit_len()
};
let mut read = BitReadStream::new(BitReadBuffer::new(&data, LittleEndian));
assert_eq!(int, read_var_int(&mut read).unwrap());
assert_eq!(pos, read.pos());
}
var_int_roundtrip(0);
var_int_roundtrip(1);
var_int_roundtrip(10);
var_int_roundtrip(55);
var_int_roundtrip(355);
var_int_roundtrip(12354);
var_int_roundtrip(123125412);
}
#[test]
fn test_var_int_fixed_roundtrip() {
fn var_int_roundtrip(int: u32) {
let mut data = Vec::new();
let pos = {
let mut write = BitWriteStream::new(&mut data, LittleEndian);
let encoded = encode_var_int_fixed(int);
encoded.write_sized(&mut write, 40).unwrap();
write.bit_len()
};
assert_eq!(40, pos);
let mut read = BitReadStream::new(BitReadBuffer::new(&data, LittleEndian));
assert_eq!(int, read_var_int(&mut read).unwrap());
assert_eq!(pos, read.pos());
}
var_int_roundtrip(0);
var_int_roundtrip(1);
var_int_roundtrip(10);
var_int_roundtrip(55);
var_int_roundtrip(355);
var_int_roundtrip(12354);
var_int_roundtrip(123125412);
}
pub fn log_base2<T: PrimInt + Unsigned>(num: T) -> u32 {
(std::mem::size_of::<T>() as u32 * 8 - 1).saturating_sub(num.leading_zeros())
}