source2_demo/parser/demo/writer/
mod.rs

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

source2_demo/parser/demo/writer/
mod.rs
