source2_demo/parser/demo/writer/
mod.rs

1mod baseline;
2mod commands;
3mod entity;
4mod input;
5mod output;
6mod packet;
7mod packet_state;
8mod rewriter;
9mod runner;
10mod string_table;
11
12use crate::entity::field::{FieldPath, FieldValue, Serializer};
13use crate::entity::{Entity, EntityEvents};
14use crate::error::ParserError;
15use crate::parser::Parser;
16use crate::reader::{BitsReader, FieldPathCodec, MessageReader, SeekableReader, SliceReader};
17use crate::string_table::PackedStringTableState;
18use std::cell::RefCell;
19use std::io::{Seek, Write};
20use std::rc::Rc;
21
22pub(crate) use input::RawDemoMessage;
23pub use rewriter::{
24    rewrite_protobuf_message, DemoRewriter, MessageRewrite, PacketMessage, RewriteInterests,
25};
26
27const INSTANCE_BASELINE_TABLE: &str = "instancebaseline";
28const ENTITY_REWRITE_BUFFER_CAPACITY: usize = 8192;
29
30struct FieldReplacement {
31    serializer: Rc<Serializer>,
32    fp: FieldPath,
33    value: FieldValue,
34    value_start: usize,
35    value_end: usize,
36}
37
38struct DecodedEntityField {
39    fp: FieldPath,
40    name: Rc<str>,
41    value_start: usize,
42    value_end: usize,
43}
44
45struct EntityRewriteState {
46    field_path_codec: FieldPathCodec,
47    rewrite_paths: Vec<FieldPath>,
48    decoded_fields: Vec<DecodedEntityField>,
49    replacements: Vec<FieldReplacement>,
50}
51
52impl EntityRewriteState {
53    fn new() -> Self {
54        Self {
55            field_path_codec: FieldPathCodec::default(),
56            rewrite_paths: Vec::with_capacity(ENTITY_REWRITE_BUFFER_CAPACITY),
57            decoded_fields: Vec::with_capacity(ENTITY_REWRITE_BUFFER_CAPACITY),
58            replacements: Vec::with_capacity(ENTITY_REWRITE_BUFFER_CAPACITY),
59        }
60    }
61}
62
63#[derive(Default)]
64struct StringTableRewriteState {
65    tables: Vec<Option<PackedStringTableState>>,
66}
67
68impl StringTableRewriteState {
69    fn len(&self) -> usize {
70        self.tables.len()
71    }
72
73    fn ensure(&mut self, table_id: usize) {
74        if self.tables.len() <= table_id {
75            self.tables.resize_with(table_id + 1, || None);
76        }
77    }
78
79    fn set(&mut self, table_id: usize, state: PackedStringTableState) {
80        self.tables[table_id] = Some(state);
81    }
82
83    fn is_missing(&self, table_id: usize) -> bool {
84        self.tables[table_id].is_none()
85    }
86
87    fn take(&mut self, table_id: usize) -> Option<PackedStringTableState> {
88        self.tables[table_id].take()
89    }
90}
91
92/// Demo writer that reads demo messages and writes a rewritten stream.
93///
94/// The writer maintains the parser metadata needed for the registered
95/// rewrites, such as serializers, classes, string tables, baselines, and
96/// entity state. Output targets must be seekable so the demo header can be
97/// patched after writing.
98pub struct DemoWriter<'a, R, W>
99where
100    R: BitsReader + MessageReader,
101    W: Write + Seek,
102{
103    parser: Parser<'a, R>,
104    writer: W,
105    string_table_rewrite: StringTableRewriteState,
106    rewriters: Vec<Box<dyn DemoRewriter + 'a>>,
107    rewriter_interests: RewriteInterests,
108    entity_rewrite: EntityRewriteState,
109    bytes_written: u64,
110    file_info_offset: Option<u64>,
111}
112
113impl<'a, R, W> DemoWriter<'a, R, W>
114where
115    R: BitsReader + MessageReader,
116    W: Write + Seek,
117{
118    /// Creates a new demo writer from an existing parser and output target.
119    pub fn new(parser: Parser<'a, R>, writer: W) -> Self {
120        Self {
121            parser,
122            writer,
123            string_table_rewrite: StringTableRewriteState::default(),
124            rewriters: Vec::new(),
125            rewriter_interests: RewriteInterests::empty(),
126            entity_rewrite: EntityRewriteState::new(),
127            bytes_written: 0,
128            file_info_offset: None,
129        }
130    }
131
132    /// Adds an already constructed demo rewriter and returns a handle to it.
133    ///
134    /// Use this when the rewriter needs custom constructor state. Rewriters run
135    /// in registration order; message callbacks see the output of earlier
136    /// rewriters.
137    pub fn add_rewriter<T>(&mut self, rewriter: T) -> Rc<RefCell<T>>
138    where
139        T: DemoRewriter + 'a,
140    {
141        let rewriter = Rc::new(RefCell::new(rewriter));
142        self.rewriter_interests |= rewriter.borrow().interests();
143        self.rewriters.push(Box::new(rewriter.clone()));
144        rewriter
145    }
146
147    /// Registers a default demo rewriter and returns a handle to it.
148    ///
149    /// This mirrors
150    /// [`Parser::register_observer`](crate::Parser::register_observer): the
151    /// writer constructs `T::default()`, registers it, and returns an
152    /// `Rc<RefCell<T>>` so callers can inspect accumulated state after
153    /// writing.
154    ///
155    /// # Examples
156    ///
157    /// ```ignore
158    /// # use source2_demo::prelude::*;
159    /// # use source2_demo::proto::CDotaUserMsgChatMessage;
160    /// # use source2_demo::writer::*;
161    /// # use std::fs::File;
162    /// #[derive(Default)]
163    /// struct RemoveChat;
164    ///
165    /// #[rewriter]
166    /// impl RemoveChat {
167    ///     #[rewrite_packet_message]
168    ///     fn remove_chat(
169    ///         &mut self,
170    ///         _message: CDotaUserMsgChatMessage,
171    ///     ) -> Result<MessageRewrite, ParserError> {
172    ///         Ok(MessageRewrite::Drop)
173    ///     }
174    /// }
175    ///
176    /// # fn main() -> anyhow::Result<()> {
177    /// # let input = File::open("input.dem")?;
178    /// # let output = File::create("output.dem")?;
179    /// let mut writer = DemoWriter::from_reader(input, output)?;
180    /// writer.register_rewriter::<RemoveChat>();
181    /// writer.run()?;
182    /// # Ok(())
183    /// # }
184    /// ```
185    pub fn register_rewriter<T>(&mut self) -> Rc<RefCell<T>>
186    where
187        T: DemoRewriter + Default + 'a,
188    {
189        self.add_rewriter(T::default())
190    }
191
192    fn should_rewrite_entity(&mut self, event: EntityEvents, entity: &Entity) -> bool {
193        let ctx = &self.parser.context;
194        self.rewriters
195            .iter_mut()
196            .filter(|rewriter| {
197                rewriter
198                    .interests()
199                    .contains(RewriteInterests::ENTITY_FIELDS)
200            })
201            .all(|rewriter| rewriter.should_rewrite_entity(ctx, event, entity))
202    }
203
204    fn should_track_entity(&mut self, event: EntityEvents, entity: &Entity) -> bool {
205        let ctx = &self.parser.context;
206        self.rewriters
207            .iter_mut()
208            .filter(|rewriter| {
209                rewriter
210                    .interests()
211                    .contains(RewriteInterests::ENTITY_FIELDS)
212            })
213            .all(|rewriter| rewriter.should_track_entity(ctx, event, entity))
214    }
215
216    fn replace_entity_field(
217        &mut self,
218        event: EntityEvents,
219        entity: &Entity,
220        field_name: &str,
221        value: &FieldValue,
222    ) -> Option<FieldValue> {
223        let ctx = &self.parser.context;
224        self.rewriters
225            .iter_mut()
226            .filter(|rewriter| {
227                rewriter
228                    .interests()
229                    .contains(RewriteInterests::ENTITY_FIELDS)
230            })
231            .find_map(|rewriter| {
232                rewriter.replace_entity_field(ctx, event, entity, field_name, value)
233            })
234    }
235
236    fn has_rewriters(&self, interests: RewriteInterests) -> bool {
237        self.rewriter_interests.intersects(interests)
238    }
239
240    fn rewrites_entity_fields(&self) -> bool {
241        self.has_rewriters(RewriteInterests::ENTITY_FIELDS)
242    }
243
244    fn rewrites_string_table_entries(&self) -> bool {
245        self.has_rewriters(RewriteInterests::STRING_TABLE_ENTRIES)
246    }
247
248    fn needs_string_table_context(&self) -> bool {
249        self.rewrites_entity_fields() || self.rewrites_string_table_entries()
250    }
251
252    fn needs_packet_scan(&self) -> bool {
253        self.has_rewriters(RewriteInterests::PACKET_MESSAGE | RewriteInterests::PACKET_MESSAGES)
254            || self.rewrites_entity_fields()
255            || self.rewrites_string_table_entries()
256            || self.has_rewriters(
257                RewriteInterests::SVC_CREATE_STRING_TABLE
258                    | RewriteInterests::SVC_UPDATE_STRING_TABLE,
259            )
260    }
261
262    fn needs_svc_packet_scan(&self) -> bool {
263        self.needs_string_table_context()
264            || self.has_rewriters(
265                RewriteInterests::SVC_CREATE_STRING_TABLE
266                    | RewriteInterests::SVC_UPDATE_STRING_TABLE,
267            )
268    }
269
270    fn needs_packet_state(&self) -> bool {
271        self.needs_string_table_context()
272    }
273
274    fn needs_class_metadata(&self) -> bool {
275        self.rewrites_entity_fields()
276    }
277
278    fn needs_demo_string_table_scan(&self) -> bool {
279        self.needs_string_table_context()
280            || self.has_rewriters(RewriteInterests::DEMO_STRING_TABLES)
281    }
282
283    fn needs_demo_string_table_state(&self) -> bool {
284        self.rewrites_string_table_entries()
285    }
286
287    /// Returns the wrapped parser and output target.
288    pub fn into_parts(self) -> (Parser<'a, R>, W) {
289        (self.parser, self.writer)
290    }
291}
292
293impl<'a, W> DemoWriter<'a, SliceReader<'a>, W>
294where
295    W: Write + Seek,
296{
297    /// Creates a demo writer from replay bytes and an output target.
298    ///
299    /// This is a convenience wrapper around [`Parser::from_slice`] and
300    /// [`DemoWriter::new`].
301    pub fn from_slice(replay: &'a [u8], writer: W) -> Result<Self, ParserError> {
302        Ok(Self::new(Parser::from_slice(replay)?, writer))
303    }
304}
305
306impl<S, W> DemoWriter<'static, SeekableReader<S>, W>
307where
308    S: std::io::Read + std::io::Seek,
309    W: Write + Seek,
310{
311    /// Creates a demo writer from a seekable reader and an output target.
312    ///
313    /// This is a convenience wrapper around [`Parser::from_reader`] and
314    /// [`DemoWriter::new`].
315    pub fn from_reader(reader: S, writer: W) -> Result<Self, ParserError> {
316        Ok(Self::new(Parser::from_reader(reader)?, writer))
317    }
318}
source2_demo/parser/demo/writer/mod.rs

source2_demo/parser/demo/writer/
mod.rs
