source2_demo/parser/demo/writer/
rewriter.rs

1use crate::entity::field::FieldValue;
2use crate::entity::{Entity, EntityEvents};
3use crate::error::ParserError;
4use crate::parser::Context;
5use crate::proto::{
6    CDemoStringTables, CSvcMsgCreateStringTable, CSvcMsgUpdateStringTable, EDemoCommands, Message,
7};
8use crate::string_table::StringTableEntryUpdate;
9use std::cell::RefCell;
10use std::rc::Rc;
11
12bitflags::bitflags! {
13    /// Bitflags for declaring rewrite interests.
14    ///
15    /// Use these flags in [`DemoRewriter::interests`] to specify which rewrite
16    /// callbacks your rewriter wants to receive. This lets
17    /// [`DemoWriter`](super::DemoWriter) skip expensive decoding paths no
18    /// registered rewriter needs.
19    #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
20    pub struct RewriteInterests: u32 {
21        /// Interest in outer demo command payload rewrites.
22        const DEMO_MESSAGE = 1 << 0;
23        /// Interest in individual packet message payload rewrites.
24        const PACKET_MESSAGE = 1 << 1;
25        /// Interest in mutating the packet message list after per-message rewrites.
26        const PACKET_MESSAGES = 1 << 2;
27        /// Interest in outer `CDemoStringTables` rewrites.
28        const DEMO_STRING_TABLES = 1 << 3;
29        /// Interest in decoded string table entry updates.
30        const STRING_TABLE_ENTRIES = 1 << 4;
31        /// Interest in `svc_CreateStringTable` rewrites.
32        const SVC_CREATE_STRING_TABLE = 1 << 5;
33        /// Interest in `svc_UpdateStringTable` rewrites.
34        const SVC_UPDATE_STRING_TABLE = 1 << 6;
35        /// Interest in entity field replacement and entity rewrite filtering.
36        const ENTITY_FIELDS = 1 << 7;
37    }
38}
39
40/// Outcome for a message rewrite operation.
41pub enum MessageRewrite {
42    /// Leave the message unchanged.
43    Keep,
44    /// Encode the mutated decoded message and replace the payload.
45    ///
46    /// Use this after changing a decoded protobuf message in place.
47    Rewrite,
48    /// Replace the message payload with the provided bytes.
49    Replace(Vec<u8>),
50    /// Drop the message entirely.
51    Drop,
52}
53
54/// Trait for handling demo rewrites.
55///
56/// Implement this trait directly for custom behavior, or use the `#[rewriter]`
57/// attribute macro to generate it from annotated methods.
58#[allow(unused_variables)]
59pub trait DemoRewriter {
60    /// Declares which rewrite callbacks this rewriter is interested in.
61    ///
62    /// Return an empty [`RewriteInterests`] to leave the demo unchanged.
63    fn interests(&self) -> RewriteInterests {
64        RewriteInterests::empty()
65    }
66
67    /// Rewrites an outer demo command payload.
68    ///
69    /// Requires [`RewriteInterests::DEMO_MESSAGE`] to be set. The payload is
70    /// passed as bytes; return [`MessageRewrite::Replace`] to emit new bytes.
71    fn rewrite_demo_message(
72        &mut self,
73        ctx: &Context,
74        tick: u32,
75        msg_type: EDemoCommands,
76        payload: &[u8],
77    ) -> Result<MessageRewrite, ParserError> {
78        Ok(MessageRewrite::Keep)
79    }
80
81    /// Rewrites one packet message payload.
82    ///
83    /// Requires [`RewriteInterests::PACKET_MESSAGE`] to be set. The payload is
84    /// passed as bytes; return [`MessageRewrite::Replace`] to emit new bytes.
85    fn rewrite_packet_message(
86        &mut self,
87        ctx: &Context,
88        tick: u32,
89        msg_type: i32,
90        payload: &[u8],
91    ) -> Result<MessageRewrite, ParserError> {
92        Ok(MessageRewrite::Keep)
93    }
94
95    /// Mutates the packet message list after per-message rewrites.
96    ///
97    /// Requires [`RewriteInterests::PACKET_MESSAGES`] to be set. Messages
98    /// inserted here are output-only; they are not processed for writer
99    /// metadata state.
100    fn rewrite_packet_messages(
101        &mut self,
102        ctx: &Context,
103        tick: u32,
104        messages: &mut Vec<PacketMessage>,
105    ) -> Result<(), ParserError> {
106        Ok(())
107    }
108
109    /// Rewrites an outer demo string table message after it has been decoded.
110    ///
111    /// Requires [`RewriteInterests::DEMO_STRING_TABLES`] to be set.
112    fn rewrite_demo_string_tables(
113        &mut self,
114        ctx: &Context,
115        tick: u32,
116        message: &mut CDemoStringTables,
117    ) -> Result<MessageRewrite, ParserError> {
118        Ok(MessageRewrite::Keep)
119    }
120
121    /// Rewrites one decoded string table entry update.
122    ///
123    /// Requires [`RewriteInterests::STRING_TABLE_ENTRIES`] to be set.
124    fn rewrite_string_table_entry(
125        &mut self,
126        ctx: &Context,
127        tick: u32,
128        table_name: &str,
129        entry: &mut StringTableEntryUpdate,
130    ) -> Result<(), ParserError> {
131        Ok(())
132    }
133
134    /// Rewrites a decoded `svc_CreateStringTable` message.
135    ///
136    /// Requires [`RewriteInterests::SVC_CREATE_STRING_TABLE`] to be set.
137    fn rewrite_svc_create_string_table(
138        &mut self,
139        ctx: &Context,
140        tick: u32,
141        message: &mut CSvcMsgCreateStringTable,
142    ) -> Result<MessageRewrite, ParserError> {
143        Ok(MessageRewrite::Keep)
144    }
145
146    /// Rewrites a decoded `svc_UpdateStringTable` message.
147    ///
148    /// Requires [`RewriteInterests::SVC_UPDATE_STRING_TABLE`] to be set.
149    fn rewrite_svc_update_string_table(
150        &mut self,
151        ctx: &Context,
152        tick: u32,
153        message: &mut CSvcMsgUpdateStringTable,
154    ) -> Result<MessageRewrite, ParserError> {
155        Ok(MessageRewrite::Keep)
156    }
157
158    /// Returns a replacement value for a decoded entity field.
159    ///
160    /// Requires [`RewriteInterests::ENTITY_FIELDS`] to be set. The first
161    /// registered rewriter that returns `Some` wins.
162    fn replace_entity_field(
163        &mut self,
164        ctx: &Context,
165        event: EntityEvents,
166        entity: &Entity,
167        field_name: &str,
168        value: &FieldValue,
169    ) -> Option<FieldValue> {
170        None
171    }
172
173    /// Decides whether an entity should enter the decoded field rewrite path.
174    ///
175    /// Requires [`RewriteInterests::ENTITY_FIELDS`] to be set. If any
176    /// registered rewriter returns `false`, the entity is not field-rewritten.
177    fn should_rewrite_entity(
178        &mut self,
179        ctx: &Context,
180        event: EntityEvents,
181        entity: &Entity,
182    ) -> bool {
183        true
184    }
185
186    /// Decides whether decoded state should be retained for an entity.
187    ///
188    /// Requires [`RewriteInterests::ENTITY_FIELDS`] to be set. Returning
189    /// `false` allows the writer to skip decoding fields when the entity also
190    /// does not enter the rewrite path.
191    fn should_track_entity(&mut self, ctx: &Context, event: EntityEvents, entity: &Entity) -> bool {
192        true
193    }
194}
195
196impl<T> DemoRewriter for Rc<RefCell<T>>
197where
198    T: DemoRewriter,
199{
200    fn interests(&self) -> RewriteInterests {
201        self.borrow().interests()
202    }
203
204    fn rewrite_demo_message(
205        &mut self,
206        ctx: &Context,
207        tick: u32,
208        msg_type: EDemoCommands,
209        payload: &[u8],
210    ) -> Result<MessageRewrite, ParserError> {
211        self.borrow_mut()
212            .rewrite_demo_message(ctx, tick, msg_type, payload)
213    }
214
215    fn rewrite_packet_message(
216        &mut self,
217        ctx: &Context,
218        tick: u32,
219        msg_type: i32,
220        payload: &[u8],
221    ) -> Result<MessageRewrite, ParserError> {
222        self.borrow_mut()
223            .rewrite_packet_message(ctx, tick, msg_type, payload)
224    }
225
226    fn rewrite_packet_messages(
227        &mut self,
228        ctx: &Context,
229        tick: u32,
230        messages: &mut Vec<PacketMessage>,
231    ) -> Result<(), ParserError> {
232        self.borrow_mut()
233            .rewrite_packet_messages(ctx, tick, messages)
234    }
235
236    fn rewrite_demo_string_tables(
237        &mut self,
238        ctx: &Context,
239        tick: u32,
240        message: &mut CDemoStringTables,
241    ) -> Result<MessageRewrite, ParserError> {
242        self.borrow_mut()
243            .rewrite_demo_string_tables(ctx, tick, message)
244    }
245
246    fn rewrite_string_table_entry(
247        &mut self,
248        ctx: &Context,
249        tick: u32,
250        table_name: &str,
251        entry: &mut StringTableEntryUpdate,
252    ) -> Result<(), ParserError> {
253        self.borrow_mut()
254            .rewrite_string_table_entry(ctx, tick, table_name, entry)
255    }
256
257    fn rewrite_svc_create_string_table(
258        &mut self,
259        ctx: &Context,
260        tick: u32,
261        message: &mut CSvcMsgCreateStringTable,
262    ) -> Result<MessageRewrite, ParserError> {
263        self.borrow_mut()
264            .rewrite_svc_create_string_table(ctx, tick, message)
265    }
266
267    fn rewrite_svc_update_string_table(
268        &mut self,
269        ctx: &Context,
270        tick: u32,
271        message: &mut CSvcMsgUpdateStringTable,
272    ) -> Result<MessageRewrite, ParserError> {
273        self.borrow_mut()
274            .rewrite_svc_update_string_table(ctx, tick, message)
275    }
276
277    fn replace_entity_field(
278        &mut self,
279        ctx: &Context,
280        event: EntityEvents,
281        entity: &Entity,
282        field_name: &str,
283        value: &FieldValue,
284    ) -> Option<FieldValue> {
285        self.borrow_mut()
286            .replace_entity_field(ctx, event, entity, field_name, value)
287    }
288
289    fn should_rewrite_entity(
290        &mut self,
291        ctx: &Context,
292        event: EntityEvents,
293        entity: &Entity,
294    ) -> bool {
295        self.borrow_mut().should_rewrite_entity(ctx, event, entity)
296    }
297
298    fn should_track_entity(&mut self, ctx: &Context, event: EntityEvents, entity: &Entity) -> bool {
299        self.borrow_mut().should_track_entity(ctx, event, entity)
300    }
301}
302
303/// Rewrites a prost message by decoding, mutating, and re-encoding it.
304pub fn rewrite_protobuf_message<M, F>(
305    msg: &[u8],
306    mut rewrite: F,
307) -> Result<MessageRewrite, ParserError>
308where
309    M: Message + Default,
310    F: FnMut(&mut M) -> Result<MessageRewrite, ParserError>,
311{
312    let mut decoded = M::decode(msg)?;
313    match rewrite(&mut decoded)? {
314        MessageRewrite::Keep => Ok(MessageRewrite::Keep),
315        MessageRewrite::Drop => Ok(MessageRewrite::Drop),
316        MessageRewrite::Replace(bytes) => Ok(MessageRewrite::Replace(bytes)),
317        MessageRewrite::Rewrite => Ok(MessageRewrite::Replace(decoded.encode_to_vec())),
318    }
319}
320
321/// Packet message passed to packet-list rewriters.
322#[derive(Clone, Debug)]
323pub struct PacketMessage {
324    /// Inner packet message type.
325    pub msg_type: i32,
326    /// Encoded inner packet message payload.
327    pub payload: Vec<u8>,
328}
329
330impl PacketMessage {
331    /// Creates a packet message from a message type and encoded payload.
332    pub fn new(msg_type: i32, payload: impl Into<Vec<u8>>) -> Self {
333        Self {
334            msg_type,
335            payload: payload.into(),
336        }
337    }
338
339    /// Creates a packet message by encoding a protobuf payload.
340    pub fn encoded<M>(msg_type: i32, message: &M) -> Self
341    where
342        M: Message,
343    {
344        Self::new(msg_type, message.encode_to_vec())
345    }
346
347    /// Decodes the packet message payload as a protobuf message.
348    pub fn decode<M>(&self) -> Result<M, ParserError>
349    where
350        M: Message + Default,
351    {
352        Ok(M::decode(self.payload.as_slice())?)
353    }
354
355    /// Replaces the packet message payload with an encoded protobuf message.
356    pub fn replace_with<M>(&mut self, message: &M)
357    where
358        M: Message,
359    {
360        self.payload = message.encode_to_vec();
361    }
362}
source2_demo/parser/demo/writer/rewriter.rs

source2_demo/parser/demo/writer/
rewriter.rs
