diff --git a/src/bin/reencode.rs b/src/bin/reencode.rs index eb80475..053f27d 100644 --- a/src/bin/reencode.rs +++ b/src/bin/reencode.rs @@ -8,10 +8,14 @@ use main_error::MainError; use tf_demo_parser::demo::data::DemoTick; use tf_demo_parser::demo::header::Header; use tf_demo_parser::demo::message::{setconvar::SetConVarMessage, Message, NetTickMessage}; +use tf_demo_parser::demo::packet::datatable::{DataTablePacket, SendTable}; use tf_demo_parser::demo::packet::stop::StopPacket; +use tf_demo_parser::demo::packet::stringtable::{ExtraData, StringTable, StringTableEntry}; +use tf_demo_parser::demo::parser::state::StaticBaseline; +use std::collections::HashSet; use tf_demo_parser::demo::packet::{Packet, PacketType}; use tf_demo_parser::demo::parser::{DemoHandler, Encode, RawPacketStream}; -use tf_demo_parser::{Demo, ParseError}; +use tf_demo_parser::{Demo, ParseError, Stream}; fn main() -> Result<(), MainError> { #[cfg(feature = "trace")] @@ -55,6 +59,13 @@ fn main() -> Result<(), MainError> { let mut encode_handler = DemoHandler::default(); encode_handler.handle_header(&header); + // instancebaseline is created before the tables arrive, so we need + // them up front. + let mut baselines = match find_data_tables(&header, stream.clone())? { + Some(tables) => Baselines::new(&header, &tables)?, + None => None, + }; + let mut has_stop = false; let mut last_tick = DemoTick::default(); @@ -67,30 +78,27 @@ fn main() -> Result<(), MainError> { let mut encode_packet = packet.clone(); match &mut encode_packet { Packet::DataTables(tables_packet) => { - for table in tables_packet.tables.iter_mut() { - for prop in table.props.iter_mut() { - match (table.name.as_str(), prop.name.as_str()) { - ("DT_ObjectDispenser", "\"healing_array\"") => { - prop.element_count = Some(101); - } - ("DT_Team", "\"player_array\"") => { - prop.element_count = Some(101); - } - ("DT_TFTeam", "\"team_object_array\"") => { - prop.element_count = Some(606); - } - _ => {} - } + fixup_tables(tables_packet); + } + Packet::StringTables(tables_packet) => { + if let Some(baselines) = baselines.as_mut() { + for table in tables_packet.tables.iter_mut() { + baselines.rewrite_table(table); } } } - Packet::Signon(message_packet) | Packet::Message(message_packet) if strip_pov => { - message_packet.meta.view_angles = Default::default(); - message_packet.messages.iter_mut().for_each(|msg| { - if let Message::ServerInfo(info) = msg { + Packet::Signon(message_packet) | Packet::Message(message_packet) => { + if strip_pov { + message_packet.meta.view_angles = Default::default(); + } + for msg in message_packet.messages.iter_mut() { + if let Some(baselines) = baselines.as_mut() { + baselines.rewrite_message(msg); + } + if strip_pov && let Message::ServerInfo(info) = msg { info.stv = true; } - }); + } } _ => {} } @@ -121,6 +129,148 @@ fn main() -> Result<(), MainError> { Ok(()) } +/// Raise the MAX_PLAYERS arrays from the pre-100-player sizes, which a current +/// client rejects with "RecvProp type doesn't match server type". +/// +/// Also changes the encoding: the length field is log2(element_count) + 1 bits. +/// See `Baselines`. +fn fixup_tables(tables: &mut DataTablePacket) -> bool { + let mut changed = false; + for table in tables.tables.iter_mut() { + for prop in table.props.iter_mut() { + let size = match (table.name.as_str(), prop.name.as_str()) { + ("DT_ObjectDispenser", "\"healing_array\"") => 101, + ("DT_Team", "\"player_array\"") => 101, + ("DT_TFTeam", "\"team_object_array\"") => 606, + _ => continue, + }; + if prop.element_count != Some(size) { + prop.element_count = Some(size); + changed = true; + } + } + } + changed +} + +/// Read forward to the DataTables packet. +fn find_data_tables(header: &Header, stream: Stream) -> Result, MainError> { + let mut packets = RawPacketStream::new(stream); + let mut handler = DemoHandler::default(); + handler.handle_header(header); + while let Some(packet) = packets.next(&handler.state_handler)? { + if let Packet::DataTables(tables) = &packet { + return Ok(Some(tables.clone())); + } + handler.handle_packet(packet)?; + } + Ok(None) +} + +/// Re-encodes instance baselines to match the tables `fixup_tables` rewrote. +/// +/// Left stale, SourceTV demos decode team data wrongly and POV demos overflow +/// the baseline merge buffer. +struct Baselines { + old: Vec, + new: Vec, + /// Classes whose props encode differently. The rest are left untouched. + affected: HashSet, + /// Ids of the instancebaseline tables. UpdateStringTable names its table + /// by id, in creation order. + baseline_tables: HashSet, + created: u8, +} + +impl Baselines { + fn new(header: &Header, tables: &DataTablePacket) -> Result, MainError> { + let mut fixed = tables.clone(); + if !fixup_tables(&mut fixed) { + return Ok(None); + } + + let mut before = DemoHandler::default(); + before.handle_header(header); + before.handle_packet(Packet::DataTables(tables.clone()))?; + + let mut after = DemoHandler::default(); + after.handle_header(header); + after.handle_packet(Packet::DataTables(fixed))?; + + let old = before.state_handler.send_tables.clone(); + let new = after.state_handler.send_tables.clone(); + let affected = old + .iter() + .zip(new.iter()) + .enumerate() + .filter(|(_, (o, n))| o.flattened_props != n.flattened_props) + .map(|(id, _)| id as u32) + .collect(); + + Ok(Some(Baselines { + old, + new, + affected, + baseline_tables: HashSet::new(), + created: 0, + })) + } + + fn rewrite_message(&mut self, message: &mut Message) { + match message { + Message::CreateStringTable(create) => { + if create.table.name.as_ref() == "instancebaseline" { + self.baseline_tables.insert(self.created); + } + self.created = self.created.wrapping_add(1); + self.rewrite_table(&mut create.table); + } + Message::UpdateStringTable(update) => { + if self.baseline_tables.contains(&update.table_id) { + for (_, entry) in update.entries.iter_mut() { + self.rewrite_entry(entry); + } + } + } + _ => {} + } + } + + fn rewrite_table(&mut self, table: &mut StringTable) { + if table.name.as_ref() != "instancebaseline" { + return; + } + for (_, entry) in table.entries.iter_mut() { + self.rewrite_entry(entry); + } + } + + fn rewrite_entry(&mut self, entry: &mut StringTableEntry) { + let Some(class_id) = entry.text.as_ref().and_then(|t| t.trim().parse::().ok()) else { + return; + }; + if !self.affected.contains(&class_id) { + return; + } + let (Some(from), Some(to), Some(extra)) = ( + self.old.get(class_id as usize), + self.new.get(class_id as usize), + entry.extra_data.as_ref(), + ) else { + return; + }; + + let baseline = StaticBaseline { + class_id: (class_id as u16).into(), + raw: extra.data.to_owned(), + }; + // A blob we cannot read is better left alone than half written. + if let Ok(raw) = baseline.reencode(from, to) { + entry.extra_data = Some(ExtraData::new(raw)); + } + } +} + fn header_fixup(header: &mut Header, mut packets: RawPacketStream) -> Result<(), MainError> { let mut ticks = 0; let mut handler = DemoHandler::default();