tf_demo_parser/demo/message/
packetentities.rs

1use bitbuffer::{BitRead, BitReadSized, BitWrite, BitWriteSized, BitWriteStream, LittleEndian};
2use serde::{Deserialize, Serialize};
3use serde_repr::{Deserialize_repr, Serialize_repr};
4use std::borrow::Cow;
5
6use crate::demo::message::stringtable::log_base2;
7use crate::demo::packet::datatable::{ClassId, SendTable};
8use crate::demo::parser::{Encode, ParseBitSkip};
9use crate::demo::sendprop::{SendProp, SendPropIdentifier, SendPropValue};
10use crate::{Parse, ParseError, ParserState, ReadResult, Result, Stream};
11use parse_display::{Display, FromStr};
12use std::cmp::{min, Ordering};
13use std::collections::HashSet;
14
15use crate::demo::data::ServerTick;
16use itertools::Either;
17use std::fmt;
18#[cfg(feature = "trace")]
19use tracing::trace;
20
21#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
22#[derive(
23    Debug,
24    Copy,
25    Clone,
26    PartialEq,
27    Eq,
28    Hash,
29    Serialize,
30    Deserialize,
31    Display,
32    Ord,
33    PartialOrd,
34    FromStr,
35    Default,
36)]
37pub struct EntityId(u32);
38
39impl From<u32> for EntityId {
40    fn from(num: u32) -> Self {
41        EntityId(num)
42    }
43}
44
45impl From<EntityId> for u32 {
46    fn from(id: EntityId) -> Self {
47        id.0
48    }
49}
50
51impl From<usize> for EntityId {
52    fn from(num: usize) -> Self {
53        EntityId(num as u32)
54    }
55}
56
57impl From<EntityId> for usize {
58    fn from(id: EntityId) -> Self {
59        id.0 as usize
60    }
61}
62
63impl PartialEq<u32> for EntityId {
64    fn eq(&self, other: &u32) -> bool {
65        self.0 == *other
66    }
67}
68
69impl PartialOrd<u32> for EntityId {
70    fn partial_cmp(&self, other: &u32) -> Option<Ordering> {
71        self.0.partial_cmp(other)
72    }
73}
74
75#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
76#[derive(
77    BitRead, BitWrite, Clone, Copy, Debug, PartialEq, Eq, Serialize_repr, Deserialize_repr,
78)]
79#[discriminant_bits = 2]
80#[repr(u8)]
81pub enum UpdateType {
82    Preserve = 0b00,
83    Leave = 0b01,
84    Enter = 0b10,
85    Delete = 0b11,
86}
87
88#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
89#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
90pub struct PacketEntity {
91    pub server_class: ClassId,
92    pub entity_index: EntityId,
93    pub props: Vec<SendProp>,
94    pub in_pvs: bool,
95    pub update_type: UpdateType,
96    pub serial_number: u32,
97    pub delay: Option<f32>,
98    pub delta: Option<ServerTick>,
99    pub baseline_index: usize,
100}
101
102impl fmt::Display for PacketEntity {
103    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
104        writeln!(f, "{}({}) {{", self.entity_index, self.server_class)?;
105        for child in self.props.iter() {
106            writeln!(f, "\t{}", child)?;
107        }
108        write!(f, "}}")
109    }
110}
111
112impl PacketEntity {
113    fn mut_prop_by_identifier(&mut self, index: &SendPropIdentifier) -> Option<&mut SendProp> {
114        self.props.iter_mut().find(|prop| prop.identifier == *index)
115    }
116
117    pub fn get_prop_by_identifier(
118        &self,
119        index: &SendPropIdentifier,
120        parser_state: &ParserState,
121    ) -> Option<SendProp> {
122        self.props(parser_state)
123            .find(|prop| prop.identifier == *index)
124    }
125
126    pub fn apply_update(&mut self, props: &[SendProp]) {
127        for prop in props {
128            match self.mut_prop_by_identifier(&prop.identifier) {
129                Some(existing_prop) => existing_prop.value = prop.value.clone(),
130                None => self.props.push(prop.clone()),
131            }
132        }
133    }
134
135    pub fn get_prop_by_name(
136        &self,
137        table_name: &str,
138        name: &str,
139        parser_state: &ParserState,
140    ) -> Option<SendProp> {
141        let identifier = SendPropIdentifier::new(table_name, name);
142        self.get_prop_by_identifier(&identifier, parser_state)
143    }
144
145    pub fn get_baseline_props<'a>(&self, parser_state: &'a ParserState) -> Cow<'a, [SendProp]> {
146        parser_state
147            .get_baseline(
148                self.baseline_index,
149                self.entity_index,
150                self.server_class,
151                &parser_state.send_tables[usize::from(self.server_class)],
152                self.delta.is_some(),
153            )
154            .unwrap_or_default()
155    }
156
157    pub fn props<'a>(
158        &'a self,
159        parser_state: &'a ParserState,
160    ) -> impl Iterator<Item = SendProp> + 'a {
161        if self.update_type == UpdateType::Enter {
162            let mut found_props = HashSet::<SendPropIdentifier>::new();
163            let props = self.props.iter().cloned();
164            #[allow(clippy::unnecessary_to_owned)]
165            let baseline_props = self
166                .get_baseline_props(parser_state)
167                .into_owned()
168                .into_iter();
169            Either::Left(props.chain(baseline_props).filter(move |prop| {
170                let found = found_props.contains(&prop.identifier);
171                found_props.insert(prop.identifier);
172                !found
173            }))
174        } else {
175            Either::Right(self.props.iter().cloned())
176        }
177    }
178}
179
180fn read_bit_var<'a, T: BitReadSized<'a, LittleEndian>>(stream: &mut Stream<'a>) -> ReadResult<T> {
181    let ty: u8 = stream.read_sized(2)?;
182
183    let bits = match ty {
184        0 => 4,
185        1 => 8,
186        2 => 12,
187        _ => 32,
188    };
189    stream.read_sized(bits)
190}
191
192fn write_bit_var(var: u32, stream: &mut BitWriteStream<LittleEndian>) -> ReadResult<()> {
193    let (ty, bits): (u8, usize) = if var >= 2u32.pow(12) {
194        (3, 32)
195    } else if var >= 2u32.pow(8) {
196        (2, 12)
197    } else if var >= 2u32.pow(4) {
198        (1, 8)
199    } else {
200        (0, 4)
201    };
202
203    ty.write_sized(stream, 2)?;
204    var.write_sized(stream, bits)
205}
206
207#[test]
208fn test_bit_var_roundtrip() {
209    use bitbuffer::{BitReadBuffer, BitReadStream};
210
211    fn bit_var_normal(val: u32) {
212        let mut data = Vec::with_capacity(16);
213        let pos = {
214            let mut write = BitWriteStream::new(&mut data, LittleEndian);
215            write_bit_var(val, &mut write).unwrap();
216            write.bit_len()
217        };
218        let mut read = BitReadStream::new(BitReadBuffer::new(&data, LittleEndian));
219        assert_eq!(val, read_bit_var::<u32>(&mut read).unwrap());
220        assert_eq!(pos, read.pos());
221    }
222    bit_var_normal(0);
223    bit_var_normal(1);
224    bit_var_normal(24);
225    bit_var_normal(1234);
226    bit_var_normal(12345);
227    bit_var_normal(123456);
228    bit_var_normal(1234567);
229    bit_var_normal(12345678);
230    bit_var_normal(123456789);
231    bit_var_normal(u32::MAX);
232
233    for i in 0..31 {
234        bit_var_normal(2u32.pow(i) as u32);
235        bit_var_normal(2u32.pow(i) as u32 - 1);
236        bit_var_normal(2u32.pow(i) as u32 + 1);
237    }
238}
239
240#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
241#[derive(Debug, PartialEq, Serialize, Deserialize, Clone, Default)]
242pub struct PacketEntitiesMessage {
243    pub entities: Vec<PacketEntity>,
244    pub removed_entities: Vec<EntityId>,
245    pub max_entries: u16,
246    pub delta: Option<ServerTick>,
247    pub base_line: u8,
248    pub updated_base_line: bool,
249}
250
251fn get_send_table(state: &ParserState, class: ClassId) -> Result<&SendTable> {
252    state
253        .send_tables
254        .get(usize::from(class))
255        .ok_or(ParseError::UnknownServerClass(class))
256}
257
258fn get_entity_for_update(
259    state: &ParserState,
260    entity_index: EntityId,
261    update_type: UpdateType,
262    delta: Option<ServerTick>,
263) -> Result<PacketEntity> {
264    let class_id = *state
265        .entity_classes
266        .get(&entity_index)
267        .ok_or(ParseError::UnknownEntity(entity_index))?;
268
269    Ok(PacketEntity {
270        server_class: class_id,
271        entity_index,
272        props: Vec::with_capacity(8),
273        in_pvs: false,
274        update_type,
275        serial_number: 0,
276        delay: None,
277        delta,
278        baseline_index: 0,
279    })
280}
281
282impl Parse<'_> for PacketEntitiesMessage {
283    fn parse(stream: &mut Stream, state: &ParserState) -> Result<Self> {
284        let max_entries = stream.read_sized(11)?;
285        let delta: Option<ServerTick> = stream.read()?;
286        let base_line = stream.read_sized(1)?;
287        let updated_entries: u16 = stream.read_sized(11)?;
288        let length: u32 = stream.read_sized(20)?;
289        let updated_base_line = stream.read()?;
290
291        let mut data = stream.read_bits(length as usize)?;
292
293        let mut entities = Vec::with_capacity(min(updated_entries, 128) as usize);
294        let mut removed_entities = Vec::new();
295
296        let mut last_index: i32 = -1;
297
298        for _ in 0..updated_entries {
299            let diff: u32 = read_bit_var(&mut data)?;
300            last_index = last_index.saturating_add(diff as i32).saturating_add(1);
301            if last_index >= 2048 {
302                return Err(ParseError::InvalidDemo("invalid entity index"));
303            }
304            let entity_index = EntityId::from(last_index as u32);
305
306            let update_type = data.read()?;
307            if update_type == UpdateType::Enter {
308                let mut entity =
309                    Self::read_enter(&mut data, entity_index, state, base_line as usize, delta)?;
310                let send_table = get_send_table(state, entity.server_class)?;
311                Self::read_update(&mut data, send_table, &mut entity.props, entity_index)?;
312
313                entities.push(entity);
314            } else if update_type == UpdateType::Preserve {
315                let mut entity = get_entity_for_update(state, entity_index, update_type, delta)?;
316                let send_table = get_send_table(state, entity.server_class)?;
317
318                Self::read_update(&mut data, send_table, &mut entity.props, entity_index)?;
319                entity.in_pvs = true;
320
321                entities.push(entity);
322            } else if state.entity_classes.contains_key(&entity_index) {
323                let entity = get_entity_for_update(state, entity_index, update_type, delta)?;
324                entities.push(entity);
325            } else {
326                entities.push(PacketEntity {
327                    server_class: 0.into(),
328                    entity_index,
329                    props: vec![],
330                    in_pvs: false,
331                    update_type,
332                    serial_number: 0,
333                    delay: None,
334                    delta,
335                    baseline_index: 0,
336                });
337            }
338        }
339
340        if delta.is_some() {
341            while data.read()? {
342                removed_entities.push(data.read_sized::<u32>(11)?.into())
343            }
344        }
345
346        Ok(PacketEntitiesMessage {
347            entities,
348            removed_entities,
349            max_entries,
350            delta,
351            base_line,
352            updated_base_line,
353        })
354    }
355}
356
357impl Encode for PacketEntitiesMessage {
358    fn encode(&self, stream: &mut BitWriteStream<LittleEndian>, state: &ParserState) -> Result<()> {
359        self.max_entries.write_sized(stream, 11)?;
360        self.delta.is_some().write(stream)?;
361        if let Some(delta) = self.delta {
362            delta.write(stream)?;
363        }
364        self.base_line.write_sized(stream, 1)?;
365        self.entities.len().write_sized(stream, 11)?;
366
367        stream.reserve_int(20, |stream| {
368            self.updated_base_line.write(stream)?;
369
370            let length_start = stream.bit_len();
371
372            let mut last_index: i32 = -1;
373
374            for entity in self.entities.iter() {
375                let diff = entity.entity_index.0 as i32 - last_index - 1;
376                write_bit_var(diff as u32, stream)?;
377                last_index = entity.entity_index.0 as i32;
378
379                entity.update_type.write(stream)?;
380
381                let send_table = get_send_table(state, entity.server_class)?;
382                match entity.update_type {
383                    UpdateType::Enter => {
384                        Self::write_enter(entity, stream, state)?;
385                        Self::write_update(&entity.props, stream, send_table, entity.entity_index)?;
386                    }
387                    UpdateType::Preserve => {
388                        Self::write_update(&entity.props, stream, send_table, entity.entity_index)?;
389                    }
390                    _ => {}
391                }
392            }
393
394            if self.delta.is_some() {
395                for removed in self.removed_entities.iter() {
396                    true.write(stream)?;
397                    removed.0.write_sized(stream, 11)?;
398                }
399                false.write(stream)?;
400            }
401
402            let length_end = stream.bit_len();
403
404            Ok((length_end - length_start) as u64)
405        })
406    }
407}
408
409impl PacketEntitiesMessage {
410    fn read_enter(
411        stream: &mut Stream,
412        entity_index: EntityId,
413        state: &ParserState,
414        baseline_index: usize,
415        delta: Option<ServerTick>,
416    ) -> Result<PacketEntity> {
417        let bits = log_base2(state.server_classes.len()) + 1;
418        let class_index: ClassId = stream.read_sized::<u16>(bits as usize)?.into();
419
420        let serial = stream.read_sized(10)?;
421
422        Ok(PacketEntity {
423            server_class: class_index,
424            entity_index,
425            props: vec![],
426            in_pvs: true,
427            update_type: UpdateType::Enter,
428            serial_number: serial,
429            delay: None,
430            delta,
431            baseline_index,
432        })
433    }
434
435    fn write_enter(
436        entity: &PacketEntity,
437        stream: &mut BitWriteStream<LittleEndian>,
438        state: &ParserState,
439    ) -> Result<()> {
440        let bits = log_base2(state.server_classes.len()) + 1;
441        u16::from(entity.server_class).write_sized(stream, bits as usize)?;
442        entity.serial_number.write_sized(stream, 10)?;
443
444        Ok(())
445    }
446
447    pub fn read_update(
448        stream: &mut Stream,
449        send_table: &SendTable,
450        props: &mut Vec<SendProp>,
451        entity_index: EntityId,
452    ) -> Result<()> {
453        let mut index: i32 = -1;
454
455        #[cfg(feature = "trace")]
456        trace!(entity_index = display(entity_index), "reading update");
457        #[cfg(not(feature = "trace"))]
458        let _ = entity_index;
459
460        while stream.read()? {
461            let diff: u32 = read_bit_var(stream)?;
462            index = index.saturating_add(diff as i32).saturating_add(1);
463
464            match send_table.flattened_props.get(index as usize) {
465                Some(definition) => {
466                    let value = SendPropValue::parse(stream, &definition.parse_definition)?;
467
468                    #[cfg(feature = "trace")]
469                    trace!(
470                        entity_index = display(entity_index),
471                        index = display(index),
472                        value = debug(&value),
473                        definition = display(definition.identifier),
474                        "reading prop"
475                    );
476                    props.push(SendProp {
477                        index: index as u32,
478                        identifier: definition.identifier,
479                        value,
480                    });
481                }
482                None => {
483                    return Err(ParseError::PropIndexOutOfBounds {
484                        index,
485                        prop_count: send_table.flattened_props.len(),
486                        table: send_table.name.to_string(),
487                    });
488                }
489            }
490        }
491
492        Ok(())
493    }
494
495    pub fn write_update<'a, Props: IntoIterator<Item = &'a SendProp>>(
496        props: Props,
497        stream: &mut BitWriteStream<LittleEndian>,
498        send_table: &SendTable,
499        _entity_index: EntityId,
500    ) -> Result<()> {
501        let mut last_index: i32 = -1;
502
503        let mut props: Vec<&SendProp> = props.into_iter().collect();
504        props.sort_by(|a, b| a.index.cmp(&b.index));
505
506        for prop in props {
507            true.write(stream)?;
508
509            let index = prop.index as usize;
510
511            let definition =
512                send_table
513                    .flattened_props
514                    .get(index)
515                    .ok_or(ParseError::PropIndexOutOfBounds {
516                        index: index as i32,
517                        prop_count: send_table.flattened_props.len(),
518                        table: send_table.name.to_string(),
519                    })?;
520            write_bit_var((index as i32 - last_index - 1) as u32, stream)?;
521            last_index = index as i32;
522            prop.value.encode(stream, &definition.parse_definition)?;
523        }
524        false.write(stream)?;
525        Ok(())
526    }
527}
528
529impl ParseBitSkip<'_> for PacketEntitiesMessage {
530    fn parse_skip(stream: &mut Stream, _state: &ParserState) -> Result<()> {
531        stream.skip_bits(11)?;
532        if stream.read()? {
533            stream.skip_bits(32)?;
534        }
535        stream.skip_bits(12)?;
536        let length: u32 = stream.read_sized(20)?;
537        stream
538            .skip_bits(length as usize + 1)
539            .map_err(ParseError::from)
540    }
541}
542
543#[test]
544fn test_packet_entitier_message_roundtrip() {
545    use crate::demo::packet::datatable::{SendTable, SendTableName, ServerClass, ServerClassName};
546    use crate::demo::sendprop::{FloatDefinition, SendPropDefinition, SendPropParseDefinition};
547
548    let mut state = ParserState::new(24, |_| false, false);
549    state.server_classes = vec![
550        ServerClass {
551            id: ClassId::from(0),
552            name: ServerClassName::from("class1"),
553            data_table: SendTableName::from("table1"),
554        },
555        ServerClass {
556            id: ClassId::from(1),
557            name: ServerClassName::from("class2"),
558            data_table: SendTableName::from("table2"),
559        },
560    ];
561    state.send_tables = vec![
562        SendTable {
563            name: SendTableName::from("table1"),
564            needs_decoder: false,
565            flattened_props: vec![],
566        },
567        SendTable {
568            name: SendTableName::from("table2"),
569            needs_decoder: false,
570            flattened_props: vec![
571                SendPropDefinition {
572                    identifier: SendPropIdentifier::new("table2", "prop1"),
573                    parse_definition: SendPropParseDefinition::Int {
574                        changes_often: false,
575                        bit_count: 8,
576                    },
577                },
578                SendPropDefinition {
579                    identifier: SendPropIdentifier::new("table2", "prop2"),
580                    parse_definition: SendPropParseDefinition::String {
581                        changes_often: false,
582                    },
583                },
584                SendPropDefinition {
585                    identifier: SendPropIdentifier::new("table2", "prop3"),
586                    parse_definition: SendPropParseDefinition::Float {
587                        changes_often: false,
588                        definition: FloatDefinition::Coord,
589                    },
590                },
591            ],
592        },
593    ];
594    state
595        .entity_classes
596        .insert(EntityId::from(4u32), ClassId::from(1));
597    crate::test_roundtrip_encode(
598        PacketEntitiesMessage {
599            entities: vec![],
600            removed_entities: vec![],
601            max_entries: 0,
602            delta: None,
603            base_line: 0,
604            updated_base_line: false,
605        },
606        &state,
607    );
608    crate::test_roundtrip_encode(
609        PacketEntitiesMessage {
610            entities: vec![PacketEntity {
611                server_class: ClassId::from(0),
612                entity_index: Default::default(),
613                props: vec![],
614                in_pvs: true,
615                update_type: UpdateType::Enter,
616                serial_number: 0,
617                delay: None,
618                delta: None,
619                baseline_index: 0,
620            }],
621            removed_entities: vec![],
622            max_entries: 4,
623            delta: None,
624            base_line: 0,
625            updated_base_line: false,
626        },
627        &state,
628    );
629    crate::test_roundtrip_encode(
630        PacketEntitiesMessage {
631            entities: vec![
632                PacketEntity {
633                    server_class: ClassId::from(0),
634                    entity_index: EntityId::from(0u32),
635                    props: vec![],
636                    in_pvs: true,
637                    update_type: UpdateType::Enter,
638                    serial_number: 0,
639                    delay: None,
640                    delta: None,
641                    baseline_index: 0,
642                },
643                PacketEntity {
644                    server_class: ClassId::from(1),
645                    entity_index: EntityId::from(4u32),
646                    props: vec![
647                        SendProp {
648                            index: 0,
649                            identifier: SendPropIdentifier::new("table2", "prop1"),
650                            value: SendPropValue::Integer(4),
651                        },
652                        SendProp {
653                            index: 2,
654                            identifier: SendPropIdentifier::new("table2", "prop3"),
655                            value: SendPropValue::Float(1.0),
656                        },
657                    ],
658                    in_pvs: true,
659                    update_type: UpdateType::Preserve,
660                    serial_number: 0,
661                    delay: None,
662                    delta: None,
663                    baseline_index: 0,
664                },
665                PacketEntity {
666                    server_class: ClassId::from(1),
667                    entity_index: EntityId::from(5u32),
668                    delta: None,
669                    baseline_index: 0,
670                    props: vec![
671                        SendProp {
672                            index: 0,
673                            identifier: SendPropIdentifier::new("table2", "prop1"),
674                            value: SendPropValue::Integer(4),
675                        },
676                        SendProp {
677                            index: 2,
678                            identifier: SendPropIdentifier::new("table2", "prop3"),
679                            value: SendPropValue::Float(1.0),
680                        },
681                    ],
682                    in_pvs: true,
683                    update_type: UpdateType::Enter,
684                    serial_number: 0,
685                    delay: None,
686                },
687            ],
688            removed_entities: vec![],
689            max_entries: 4,
690            delta: None,
691            base_line: 0,
692            updated_base_line: false,
693        },
694        &state,
695    );
696}