Merge pull request 'reencode: rewrite instance baselines after the array size fixup' (#11) from pet0r/demostf-parser:fix-baseline-reencode into master

Reviewed-on: https://codeberg.org/demostf/parser/pulls/11

src/bin/reencode.rs

@@ -8,10 +8,15 @@ 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::borrow::Cow;
+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 +60,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 +79,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, &handler.string_table_names);
+                        }
+                        if strip_pov && let Message::ServerInfo(info) = msg {
                             info.stv = true;
                         }
-                    });
+                    }
                 }
                 _ => {}
             }
@@ -121,6 +130,135 @@ 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<Option<DataTablePacket>, 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<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>, 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 }))
+    }
+
+    /// `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, table_names: &[Cow<str>]) {
+        match message {
+            Message::CreateStringTable(create) => self.rewrite_table(&mut create.table),
+            Message::UpdateStringTable(update) => {
+                let name = table_names.get(update.table_id as usize);
+                if name.is_some_and(|name| name == "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));
+        }
+    }
+}
+
 fn header_fixup(header: &mut Header, mut packets: RawPacketStream) -> Result<(), MainError> {
     let mut ticks = 0;
     let mut handler = DemoHandler::default();

src/demo/parser/state.rs

@@ -70,6 +70,26 @@ impl StaticBaseline {
         )?;
         Ok(props)
     }
+
+    /// Re-encode the baseline from one send table to another.
+    ///
+    /// Baselines are packed against a specific table, so any edit that changes
+    /// how a prop encodes leaves them readable at the wrong offsets.
+    #[cfg(feature = "write")]
+    pub fn reencode(&self, from: &SendTable, to: &SendTable) -> Result<Stream<'static>> {
+        use bitbuffer::{BitReadBuffer, BitReadStream, BitWriteStream, LittleEndian};
+
+        let props = self.parse(from)?;
+        let mut bytes = Vec::new();
+        {
+            let mut out = BitWriteStream::new(&mut bytes, LittleEndian);
+            PacketEntitiesMessage::write_update(&props, &mut out, to, 0u32.into())?;
+        }
+        Ok(BitReadStream::new(BitReadBuffer::new_owned(
+            bytes,
+            LittleEndian,
+        )))
+    }
 }
 
 impl ParserState {
@@ -424,3 +444,72 @@ impl From<BaselineEntity> for PacketEntity {
         }
     }
 }
+
+#[test]
+#[cfg(feature = "write")]
+fn test_static_baseline_reencode() {
+    use crate::demo::packet::datatable::SendTable;
+    use crate::demo::sendprop::{SendPropDefinition, SendPropParseDefinition, SendPropValue};
+    use bitbuffer::{BitReadBuffer, BitReadStream, BitWriteStream, LittleEndian};
+
+    fn table(first_prop_bits: u8) -> SendTable {
+        SendTable {
+            name: SendTableName::from("table"),
+            needs_decoder: false,
+            flattened_props: vec![
+                SendPropDefinition {
+                    identifier: SendPropIdentifier::new("table", "first"),
+                    parse_definition: SendPropParseDefinition::UnsignedInt {
+                        changes_often: false,
+                        bit_count: first_prop_bits,
+                    },
+                },
+                SendPropDefinition {
+                    identifier: SendPropIdentifier::new("table", "second"),
+                    parse_definition: SendPropParseDefinition::UnsignedInt {
+                        changes_often: false,
+                        bit_count: 8,
+                    },
+                },
+            ],
+        }
+    }
+
+    let from = table(8);
+    let to = table(16);
+    let props = vec![
+        SendProp {
+            index: 0,
+            identifier: SendPropIdentifier::new("table", "first"),
+            value: SendPropValue::Integer(200),
+        },
+        SendProp {
+            index: 1,
+            identifier: SendPropIdentifier::new("table", "second"),
+            value: SendPropValue::Integer(37),
+        },
+    ];
+
+    let mut bytes = Vec::new();
+    {
+        let mut out = BitWriteStream::new(&mut bytes, LittleEndian);
+        PacketEntitiesMessage::write_update(&props, &mut out, &from, 0u32.into()).unwrap();
+    }
+    let baseline = StaticBaseline::new(
+        ClassId::from(0u16),
+        BitReadStream::new(BitReadBuffer::new_owned(bytes, LittleEndian)),
+    );
+
+    assert_eq!(baseline.parse(&from).unwrap(), props);
+
+    // the old blob no longer decodes to the same values
+    if let Ok(misread) = baseline.parse(&to) {
+        assert_ne!(misread, props);
+    }
+
+    let fixed = StaticBaseline::new(
+        ClassId::from(0u16),
+        baseline.reencode(&from, &to).unwrap(),
+    );
+    assert_eq!(fixed.parse(&to).unwrap(), props);
+}

src/demo/sendprop.rs

@@ -386,7 +386,7 @@ impl BitWrite<LittleEndian> for SendPropFlags {
 }
 
 #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub enum FloatDefinition {
     Coord,
     CoordMP,
@@ -429,7 +429,7 @@ impl FloatDefinition {
 }
 
 #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct SendPropDefinition {
     pub identifier: SendPropIdentifier,
     pub parse_definition: SendPropParseDefinition,
@@ -448,7 +448,7 @@ impl TryFrom<&RawSendPropDefinition> for SendPropDefinition {
 }
 
 #[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub enum SendPropParseDefinition {
     NormalVarInt {
         changes_often: bool,