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

source2_demo/parser/demo/writer/
mod.rs
