rewrite instance baselines after the array size fixup

wasm/src/cut/mod.rs

@@ -39,7 +39,9 @@ pub fn cut(input: &[u8], options: EditOptions) -> Vec<u8> {
         let mut stream = demo.get_stream();
         let mut header = Header::read(&mut stream).unwrap();
 
-        let mut mutators = options.as_mutator(spectator_id);
+        let mut mutators = options
+            .as_mutator(spectator_id, &header, &stream)
+            .expect("invalid demo");
         let start_tick = options.cut.unwrap().from;
         let end_tick = options.cut.unwrap().to;
 
@@ -105,9 +107,9 @@ pub fn cut(input: &[u8], options: EditOptions) -> Vec<u8> {
         }
 
         // create the net ticks needed for later deltas
-        let fill_ticks = (delta_tick + 1)
-            .range_inclusive(start_state.server_tick)
-            .map(net_tick);
+        let next_tick = delta_tick + 1;
+        let tick_range = next_tick.range_inclusive(start_state.server_tick);
+        let fill_ticks = tick_range.map(net_tick);
         let fill_packets = fill_ticks.map(|msg| {
             Packet::Message(MessagePacket {
                 messages: vec![

wasm/src/lib.rs

@@ -4,8 +4,8 @@ mod cut;
 pub mod missing_preserve;
 mod mutate;
 mod options;
-mod pov;
 mod player_count;
+mod pov;
 
 use bitbuffer::{BitRead, BitWriteStream, LittleEndian};
 use tf_demo_parser::demo::header::Header;
@@ -92,12 +92,15 @@ fn no_cut(input: &[u8], options: EditOptions) -> Vec<u8> {
         let mut output_handler = DemoHandler::default();
         output_handler.handle_header(&header);
 
-        let mutators = options.as_mutator(spectator_id);
+        let mut mutators = options
+            .as_mutator(spectator_id, &header, &stream)
+            .expect("invalid demo");
 
         while let Some(mut packet) = packets.next(&handler.state_handler).unwrap() {
-            handler.handle_packet(packet.clone()).expect("failed to handle source packet");
+            handler
+                .handle_packet(packet.clone())
+                .expect("failed to handle source packet");
             mutators.mutate_packet(&mut packet, &handler.state_handler);
-            output_handler.handle_packet(packet.clone()).expect("failed to handle mutated packet");
 
             if packet.packet_type() != PacketType::ConsoleCmd
                 && packet.packet_type() != PacketType::UserCmd
@@ -106,6 +109,10 @@ fn no_cut(input: &[u8], options: EditOptions) -> Vec<u8> {
                     .encode(&mut out_stream, &output_handler.state_handler)
                     .unwrap();
             }
+
+            output_handler
+                .handle_packet(packet)
+                .expect("failed to handle mutated packet");
         }
     }
     out_buffer

wasm/src/mutate.rs

@@ -4,7 +4,7 @@ use tf_demo_parser::demo::packet::Packet;
 use tf_demo_parser::ParserState;
 
 pub trait PacketMutator {
-    fn mutate_packet(&self, packet: &mut Packet, state: &ParserState);
+    fn mutate_packet(&mut self, packet: &mut Packet, state: &ParserState);
 }
 
 pub trait MessageMutator {
@@ -20,7 +20,7 @@ pub struct PacketMessageMutator<T: MessageMutator> {
 }
 
 impl<T: MessageMutator> PacketMutator for PacketMessageMutator<T> {
-    fn mutate_packet(&self, packet: &mut Packet, state: &ParserState) {
+    fn mutate_packet(&mut self, packet: &mut Packet, state: &ParserState) {
         match packet {
             Packet::Message(msg_packet) | Packet::Signon(msg_packet) => {
                 msg_packet
@@ -34,7 +34,7 @@ impl<T: MessageMutator> PacketMutator for PacketMessageMutator<T> {
 }
 
 impl<F: Fn(&mut Packet)> PacketMutator for F {
-    fn mutate_packet(&self, packet: &mut Packet, _state: &ParserState) {
+    fn mutate_packet(&mut self, packet: &mut Packet, _state: &ParserState) {
         self(packet)
     }
 }
@@ -56,7 +56,7 @@ pub struct PacketMessageFilter<T: MessageFilter> {
 }
 
 impl<T: MessageFilter> PacketMutator for PacketMessageFilter<T> {
-    fn mutate_packet(&self, packet: &mut Packet, _state: &ParserState) {
+    fn mutate_packet(&mut self, packet: &mut Packet, _state: &ParserState) {
         match packet {
             Packet::Message(msg_packet) | Packet::Signon(msg_packet) => {
                 let messages = take(&mut msg_packet.messages);
@@ -108,8 +108,8 @@ impl MutatorList {
 }
 
 impl PacketMutator for MutatorList {
-    fn mutate_packet(&self, packet: &mut Packet, state: &ParserState) {
-        for mutator in self.mutators.iter() {
+    fn mutate_packet(&mut self, packet: &mut Packet, state: &ParserState) {
+        for mutator in self.mutators.iter_mut() {
             mutator.mutate_packet(packet, state);
         }
     }

wasm/src/options.rs

@@ -1,10 +1,12 @@
 use crate::missing_preserve::removed_invalid_deletes;
+use crate::player_count::fix_player_count;
 use crate::{clean_demo, strip_cond, unlock_pov, MutatorList};
 use serde::{Deserialize, Serialize};
 use tf_demo_parser::demo::data::DemoTick;
+use tf_demo_parser::demo::header::Header;
 use tf_demo_parser::demo::message::packetentities::EntityId;
 use tf_demo_parser::demo::message::Message;
-use crate::player_count::fix_player_count;
+use tf_demo_parser::{ParseError, Stream};
 
 #[derive(Debug, Serialize, Deserialize, Default)]
 #[non_exhaustive]
@@ -27,7 +29,12 @@ pub struct EditOptions {
 }
 
 impl EditOptions {
-    pub fn as_mutator(&self, spectator_id: EntityId) -> MutatorList {
+    pub fn as_mutator(
+        &self,
+        spectator_id: EntityId,
+        header: &Header,
+        stream: &Stream,
+    ) -> Result<MutatorList, ParseError> {
         let mut mutators = MutatorList::new();
 
         mutators.push_message_mutator(|message: &mut Message| {
@@ -56,10 +63,10 @@ impl EditOptions {
         }
 
         if self.fix_player_count {
-            fix_player_count(&mut mutators);
+            fix_player_count(&mut mutators, header, stream)?;
         }
 
-        mutators
+        Ok(mutators)
     }
 }
 

wasm/src/player_count.rs

@@ -1,57 +1,43 @@
+use std::collections::HashSet;
+
 use crate::mutate::{MutatorList, PacketMutator};
+use tf_demo_parser::demo::header::Header;
+use tf_demo_parser::demo::message::Message;
+use tf_demo_parser::demo::packet::datatable::{DataTablePacket, SendTable};
+use tf_demo_parser::demo::packet::stringtable::{ExtraData, StringTable, StringTableEntry};
 use tf_demo_parser::demo::packet::Packet;
-use tf_demo_parser::ParserState;
+use tf_demo_parser::demo::parser::state::StaticBaseline;
+use tf_demo_parser::demo::parser::{DemoHandler, RawPacketStream};
+use tf_demo_parser::{ParseError, ParserState, Stream};
 
-#[derive(Default)]
-pub struct PlayerCount {}
+pub struct PlayerCount {
+    baselines: Option<Baselines>,
+    string_table_names: Vec<String>,
+}
 
-/// property arrays that got renamed and new items
-const PLAYER_ARRAY_PROPS_RENAMED: &[(&str, u16, u16)] = &[
-    ("m_aBonusPlayerRoll_33", 33, 101),
-    ("m_ConditionData_128", 128, 131),
-];
+fn find_data_tables(header: &Header, stream: Stream) -> Result<DataTablePacket, ParseError> {
+    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(tables.clone());
+        }
+        handler.handle_packet(packet)?;
+    }
+    Err(ParseError::InvalidDemo("missing data tables"))
+}
 
-/// property arrays that just got new items
-const PLAYER_ARRAY_PROPS: &[(&str, u16)] = &[
-    ("m_iPing", 101),
-    ("m_iScore", 101),
-    ("m_iDeaths", 101),
-    ("m_bConnected", 101),
-    ("m_iTeam", 101),
-    ("m_bAlive", 101),
-    ("m_iHealth", 101),
-    ("m_iAccountID", 101),
-    ("m_bValid", 101),
-    ("m_iUserID", 101),
-    ("m_bPlayerReady", 101),
-    ("m_ePlayerWantsRematch", 101),
-    ("m_bPlayerIsPackMember", 101),
-    ("m_bPlayerDominated", 101),
-    ("m_bPlayerDominatingMe", 101),
-    ("m_iTotalScore", 101),
-    ("m_iMaxHealth", 101),
-    ("m_iMaxBuffedHealth", 101),
-    ("m_iPlayerClass", 101),
-    ("m_bArenaSpectator", 101),
-    ("m_iActiveDominations", 101),
-    ("m_flNextRespawnTime", 101),
-    ("m_iChargeLevel", 101),
-    ("m_iDamage", 101),
-    ("m_iDamageAssist", 101),
-    ("m_iDamageBoss", 101),
-    ("m_iHealing", 101),
-    ("m_iHealingAssist", 101),
-    ("m_iDamageBlocked", 101),
-    ("m_iCurrencyCollected", 101),
-    ("m_iBonusPoints", 101),
-    ("m_iPlayerLevel", 101),
-    ("m_iUpgradeRefundCredits", 101),
-    ("m_iBuybackCredits", 101),
-    ("m_iPlayerClassWhenKilled", 101),
-    ("m_iConnectionState", 101),
-    ("m_flConnectTime", 101),
-    ("m_iStreaks", 407),
-];
+impl PlayerCount {
+    pub fn new(header: &Header, stream: Stream) -> Result<PlayerCount, ParseError> {
+        let data_tables = find_data_tables(header, stream)?;
+        let baselines = Baselines::new(&header, &data_tables)?;
+        Ok(PlayerCount {
+            baselines,
+            string_table_names: Vec::new(),
+        })
+    }
+}
 
 // property values that got changed
 const PLAYER_PROPS_VALUES: &[(&str, &str, u16, u16)] = &[
@@ -60,61 +46,163 @@ const PLAYER_PROPS_VALUES: &[(&str, &str, u16, u16)] = &[
     ("DT_TFTeam", r#""team_object_array""#, 198, 606),
 ];
 
+fn new_prop_size(table: &str, prop: &str) -> Option<u16> {
+    for (changed_table, changed_prop, _, new_size) in PLAYER_PROPS_VALUES.iter().copied() {
+        if changed_table == table && changed_prop == prop {
+            return Some(new_size);
+        }
+    }
+    None
+}
+
 impl PacketMutator for PlayerCount {
-    fn mutate_packet(&self, packet: &mut Packet, _state: &ParserState) {
-        if let Packet::DataTables(data_tables) = packet {
-            for table in &mut data_tables.tables {
-                let table_name = table.name.as_str();
-                if let Some(new_length) = PLAYER_ARRAY_PROPS.iter().copied().find_map(|(name, new_length)| (name == table_name).then_some(new_length)) {
-                    let base = table.props.last().unwrap().clone();
-                    let old_length = table.props.len() as u16;
-                    for i in old_length ..= new_length {
-                        let mut prop = base.clone();
-                        prop.name = format!("{i:03}").into();
-                        table.props.push(prop);
-                    }
-                } else if let Some((old_length, new_length)) = PLAYER_ARRAY_PROPS_RENAMED.iter().copied().find_map(|(name, old_length, new_length)| table_name.ends_with(name).then_some((old_length, new_length))) {
-                    if table_name.starts_with("_ST_") {
-                        let base = table.props.last().unwrap().clone();
-                        for i in old_length.. new_length { // non-inclusive range, because for some reasons these arrays are one short
-                            let mut prop = base.clone();
-                            prop.name = format!("{i:03}").into();
-                            table.props.push(prop);
-                        }
-                    } else if table_name.starts_with("_LPT_") {
-                        table.props[0].name = format!("lengthprop{new_length}").into();
-                    }
-                    let name_prefix = table_name.rsplit_once('_').unwrap().0;
-                    table.name = format!("{name_prefix}_{}", new_length).into();
-                } else if let Some((prop_name, _, new_length)) = PLAYER_PROPS_VALUES.iter().copied().find_map(|(name, prop_name, old_length, new_length)| table_name.eq(name).then_some((prop_name, old_length, new_length))) {
-                    for prop in &mut table.props {
-                        if prop.name == prop_name {
-                            prop.element_count = new_length.into();
-                        }
-                    }
-                }
-                for prop in &mut table.props {
-                    if let Some(target_name) = &mut prop.table_name {
-                        if let Some((_, new_length)) = PLAYER_ARRAY_PROPS_RENAMED.iter().copied().find_map(|(name, old_length, new_length)| target_name.ends_with(name).then_some((old_length, new_length))) {
-                            let name_prefix = target_name.rsplit_once('_').unwrap().0;
-                            *target_name = format!("{name_prefix}_{}", new_length).into();
-                        }
-                    }
-                }
-                // let table_name = table.name.as_str();
-                // for prop in &mut table.props {
-                //     println!(
-                //         "{}.{}: {}",
-                //         table_name,
-                //         prop.name,
-                //         prop.element_count.unwrap_or_default()
-                //     );
-                // }
+    fn mutate_packet(&mut self, packet: &mut Packet, _state: &ParserState) {
+        match packet {
+            Packet::DataTables(data_tables) => {
+                fixup_tables(data_tables);
             }
+            Packet::StringTables(tables_packet) => {
+                if let Some(baselines) = self.baselines.as_mut() {
+                    for table in tables_packet.tables.iter_mut() {
+                        self.string_table_names.push(table.name.as_ref().into());
+                        baselines.rewrite_table(table);
+                    }
+                }
+            }
+            Packet::Signon(message_packet) | Packet::Message(message_packet) => {
+                for msg in message_packet.messages.iter_mut() {
+                    if let Some(baselines) = self.baselines.as_mut() {
+                        baselines.rewrite_message(msg, &self.string_table_names);
+                    }
+                }
+            }
+            _ => {}
         }
     }
 }
 
-pub fn fix_player_count(mutators: &mut MutatorList) {
-    mutators.push_packet_mutator(PlayerCount::default())
+pub fn fix_player_count(
+    mutators: &mut MutatorList,
+    header: &Header,
+    stream: &Stream,
+) -> Result<(), ParseError> {
+    PlayerCount::new(header, stream.clone())
+        .map(|player_count| mutators.push_packet_mutator(player_count))
+}
+
+/// 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() {
+            if let Some(size) = new_prop_size(table.name.as_str(), prop.name.as_str()) {
+                if prop.element_count != Some(size) {
+                    prop.element_count = Some(size);
+                    changed = true;
+                }
+            }
+        }
+    }
+    changed
+}
+
+/// 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<SendTable>,
+    new: Vec<SendTable>,
+    /// Classes whose props encode differently. The rest are left untouched.
+    affected: HashSet<u32>,
+}
+
+impl Baselines {
+    fn new(header: &Header, tables: &DataTablePacket) -> Result<Option<Self>, ParseError> {
+        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 }))
+    }
+
+    /// `table_names` is the handler's `string_table_names`, which maps the id
+    /// an UpdateStringTable carries back to the table it belongs to.
+    fn rewrite_message(&mut self, message: &mut Message, string_tables: &[String]) {
+        match message {
+            Message::CreateStringTable(create) => self.rewrite_table(&mut create.table),
+            Message::UpdateStringTable(update) => {
+                let name = string_tables
+                    .get(update.table_id as usize)
+                    .map(String::as_str);
+                if name == Some("instancebaseline") {
+                    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::<u32>().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));
+        }
+    }
 }