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
187impl<T> DemoRewriter for Rc<RefCell<T>>
188where
189    T: DemoRewriter,
190{
191    fn interests(&self) -> RewriteInterests {
192        self.borrow().interests()
193    }
194
195    fn rewrite_demo_message(
196        &mut self,
197        ctx: &Context,
198        tick: u32,
199        msg_type: EDemoCommands,
200        payload: &[u8],
201    ) -> Result<MessageRewrite, ParserError> {
202        self.borrow_mut()
203            .rewrite_demo_message(ctx, tick, msg_type, payload)
204    }
205
206    fn rewrite_packet_message(
207        &mut self,
208        ctx: &Context,
209        tick: u32,
210        msg_type: i32,
211        payload: &[u8],
212    ) -> Result<MessageRewrite, ParserError> {
213        self.borrow_mut()
214            .rewrite_packet_message(ctx, tick, msg_type, payload)
215    }
216
217    fn rewrite_packet_messages(
218        &mut self,
219        ctx: &Context,
220        tick: u32,
221        messages: &mut Vec<PacketMessage>,
222    ) -> Result<(), ParserError> {
223        self.borrow_mut()
224            .rewrite_packet_messages(ctx, tick, messages)
225    }
226
227    fn rewrite_demo_string_tables(
228        &mut self,
229        ctx: &Context,
230        tick: u32,
231        message: &mut CDemoStringTables,
232    ) -> Result<MessageRewrite, ParserError> {
233        self.borrow_mut()
234            .rewrite_demo_string_tables(ctx, tick, message)
235    }
236
237    fn rewrite_string_table_entry(
238        &mut self,
239        ctx: &Context,
240        tick: u32,
241        table_name: &str,
242        entry: &mut StringTableEntryUpdate,
243    ) -> Result<(), ParserError> {
244        self.borrow_mut()
245            .rewrite_string_table_entry(ctx, tick, table_name, entry)
246    }
247
248    fn rewrite_svc_create_string_table(
249        &mut self,
250        ctx: &Context,
251        tick: u32,
252        message: &mut CSvcMsgCreateStringTable,
253    ) -> Result<MessageRewrite, ParserError> {
254        self.borrow_mut()
255            .rewrite_svc_create_string_table(ctx, tick, message)
256    }
257
258    fn rewrite_svc_update_string_table(
259        &mut self,
260        ctx: &Context,
261        tick: u32,
262        message: &mut CSvcMsgUpdateStringTable,
263    ) -> Result<MessageRewrite, ParserError> {
264        self.borrow_mut()
265            .rewrite_svc_update_string_table(ctx, tick, message)
266    }
267
268    fn replace_entity_field(
269        &mut self,
270        ctx: &Context,
271        event: EntityEvents,
272        entity: &Entity,
273        field_name: &str,
274        value: &FieldValue,
275    ) -> Option<FieldValue> {
276        self.borrow_mut()
277            .replace_entity_field(ctx, event, entity, field_name, value)
278    }
279
280    fn should_rewrite_entity(
281        &mut self,
282        ctx: &Context,
283        event: EntityEvents,
284        entity: &Entity,
285    ) -> bool {
286        self.borrow_mut().should_rewrite_entity(ctx, event, entity)
287    }
288}
289
290/// Rewrites a prost message by decoding, mutating, and re-encoding it.
291pub fn rewrite_protobuf_message<M, F>(
292    msg: &[u8],
293    mut rewrite: F,
294) -> Result<MessageRewrite, ParserError>
295where
296    M: Message + Default,
297    F: FnMut(&mut M) -> Result<MessageRewrite, ParserError>,
298{
299    let mut decoded = M::decode(msg)?;
300    match rewrite(&mut decoded)? {
301        MessageRewrite::Keep => Ok(MessageRewrite::Keep),
302        MessageRewrite::Drop => Ok(MessageRewrite::Drop),
303        MessageRewrite::Replace(bytes) => Ok(MessageRewrite::Replace(bytes)),
304        MessageRewrite::Rewrite => Ok(MessageRewrite::Replace(decoded.encode_to_vec())),
305    }
306}
307
308/// Packet message passed to packet-list rewriters.
309#[derive(Clone, Debug)]
310pub struct PacketMessage {
311    /// Inner packet message type.
312    pub msg_type: i32,
313    /// Encoded inner packet message payload.
314    pub payload: Vec<u8>,
315}
316
317impl PacketMessage {
318    /// Creates a packet message from a message type and encoded payload.
319    pub fn new(msg_type: i32, payload: impl Into<Vec<u8>>) -> Self {
320        Self {
321            msg_type,
322            payload: payload.into(),
323        }
324    }
325
326    /// Creates a packet message by encoding a protobuf payload.
327    pub fn encoded<M>(msg_type: i32, message: &M) -> Self
328    where
329        M: Message,
330    {
331        Self::new(msg_type, message.encode_to_vec())
332    }
333
334    /// Decodes the packet message payload as a protobuf message.
335    pub fn decode<M>(&self) -> Result<M, ParserError>
336    where
337        M: Message + Default,
338    {
339        Ok(M::decode(self.payload.as_slice())?)
340    }
341
342    /// Replaces the packet message payload with an encoded protobuf message.
343    pub fn replace_with<M>(&mut self, message: &M)
344    where
345        M: Message,
346    {
347        self.payload = message.encode_to_vec();
348    }
349}
source2_demo/parser/demo/writer/rewriter.rs

source2_demo/parser/demo/writer/
rewriter.rs
