Skip to main content

source2_demo/parser/demo/writer/
run.rs

1use super::packet::PacketDataRewrite;
2use super::*;
3use crate::parser::demo::DemoCommands;
4use crate::proto::{
5    CDemoClassInfo, CDemoFullPacket, CDemoPacket, CDemoSendTables, CDemoStringTables,
6    EDemoCommands, Message,
7};
8use crate::reader::{BitsReader, MessageReader};
9use std::io::{Seek, Write};
10
11impl<'a, R, W> DemoWriter<'a, R, W>
12where
13    R: BitsReader + MessageReader,
14    W: Write + Seek,
15{
16    /// Parses the demo while writing the rewritten stream.
17    pub fn run(&mut self) -> Result<(), ParserError> {
18        self.parser.reader.seek(0);
19        let header = self.parser.reader.read_bytes(16);
20        self.writer.write_all(&header)?;
21        self.bytes_written = header.len() as u64;
22
23        'messages: while let Some(message) = self.read_next_raw_message()? {
24            self.update_context_tick(message.tick);
25
26            let mut payload = None;
27            if self.has_rewriters(RewriteInterests::DEMO_MESSAGE) {
28                let mut decoded = Self::decode_raw_payload(&message)?;
29                let ctx = &self.parser.context;
30                for rewriter in self.rewriters.iter_mut().filter(|rewriter| {
31                    rewriter
32                        .interests()
33                        .contains(RewriteInterests::DEMO_MESSAGE)
34                }) {
35                    match rewriter.rewrite_demo_message(
36                        ctx,
37                        message.tick,
38                        message.msg_type,
39                        decoded.as_slice(),
40                    )? {
41                        MessageRewrite::Drop => continue 'messages,
42                        MessageRewrite::Replace(bytes) => decoded = bytes,
43                        MessageRewrite::Keep | MessageRewrite::Rewrite => {}
44                    }
45                }
46                payload = Some(decoded);
47            }
48
49            match message.msg_type {
50                EDemoCommands::DemPacket | EDemoCommands::DemSignonPacket => {
51                    if !self.needs_packet_scan() {
52                        self.write_message_or_raw(&message, payload.as_deref())?;
53                        continue;
54                    }
55
56                    let can_write_raw = payload.is_none();
57                    let payload = Self::materialize_payload(&mut payload, &message)?;
58                    let mut packet = CDemoPacket::decode(payload.as_slice())?;
59                    let rewritten = self.rewrite_packet_data_if_changed(
60                        message.tick,
61                        packet.data(),
62                        self.needs_packet_state(),
63                    )?;
64                    match rewritten {
65                        PacketDataRewrite::Unchanged if can_write_raw => {
66                            self.write_raw_demo_message(
67                                message.msg_type,
68                                message.tick,
69                                message.raw_payload.as_slice(),
70                                message.compressed,
71                            )?;
72                            continue;
73                        }
74                        PacketDataRewrite::Unchanged => {}
75                        PacketDataRewrite::Changed(rewritten) => packet.data = Some(rewritten),
76                    }
77                    let packet_bytes = packet.encode_to_vec();
78                    self.write_demo_message_with_compression(
79                        message.msg_type,
80                        message.tick,
81                        &packet_bytes,
82                        message.compressed,
83                    )?;
84                }
85                EDemoCommands::DemFullPacket => {
86                    if !self.needs_packet_scan() && !self.needs_demo_string_table_scan() {
87                        self.write_message_or_raw(&message, payload.as_deref())?;
88                        continue;
89                    }
90
91                    let can_write_raw = payload.is_none();
92                    let payload = Self::materialize_payload(&mut payload, &message)?;
93                    let mut full = CDemoFullPacket::decode(payload.as_slice())?;
94                    let should_process = self.parser.context.last_full_packet_tick == u32::MAX
95                        || self.parser.skip_deltas;
96                    let mut changed = false;
97
98                    if let Some(table) = full.string_table.take() {
99                        if self.needs_demo_string_table_scan() {
100                            if let Some((rewritten, table_changed)) =
101                                self.rewrite_string_tables(message.tick, table)?
102                            {
103                                if should_process && self.needs_demo_string_table_state() {
104                                    self.parser.dem_string_tables(rewritten.clone())?;
105                                }
106                                full.string_table = Some(rewritten);
107                                changed |= table_changed;
108                            } else {
109                                changed = true;
110                            }
111                        } else {
112                            full.string_table = Some(table);
113                        }
114                    }
115
116                    if let Some(packet) = full.packet.as_mut() {
117                        if self.needs_packet_scan() {
118                            let rewritten = self.rewrite_packet_data_if_changed(
119                                message.tick,
120                                packet.data(),
121                                should_process && self.needs_packet_state(),
122                            )?;
123                            if let PacketDataRewrite::Changed(rewritten) = rewritten {
124                                packet.data = Some(rewritten);
125                                changed = true;
126                            }
127                        }
128                    }
129
130                    self.parser.context.last_full_packet_tick = message.tick;
131                    if can_write_raw && !changed {
132                        self.write_raw_demo_message(
133                            message.msg_type,
134                            message.tick,
135                            message.raw_payload.as_slice(),
136                            message.compressed,
137                        )?;
138                        continue;
139                    }
140                    let full_bytes = full.encode_to_vec();
141                    self.write_demo_message_with_compression(
142                        message.msg_type,
143                        message.tick,
144                        &full_bytes,
145                        message.compressed,
146                    )?;
147                }
148                EDemoCommands::DemSendTables if self.needs_class_metadata() => {
149                    let payload = Self::materialize_payload(&mut payload, &message)?;
150                    let msg = CDemoSendTables::decode(payload.as_slice())?;
151                    self.parser.dem_send_tables(msg)?;
152                    self.write_demo_message_with_compression(
153                        message.msg_type,
154                        message.tick,
155                        &payload,
156                        message.compressed,
157                    )?;
158                }
159                EDemoCommands::DemClassInfo if self.needs_class_metadata() => {
160                    let payload = Self::materialize_payload(&mut payload, &message)?;
161                    let msg = CDemoClassInfo::decode(payload.as_slice())?;
162                    self.parser.dem_class_info(msg)?;
163                    self.write_demo_message_with_compression(
164                        message.msg_type,
165                        message.tick,
166                        &payload,
167                        message.compressed,
168                    )?;
169                }
170                EDemoCommands::DemStringTables => {
171                    if !self.needs_demo_string_table_scan() {
172                        self.write_message_or_raw(&message, payload.as_deref())?;
173                        continue;
174                    }
175
176                    let payload = Self::materialize_payload(&mut payload, &message)?;
177                    let mut msg = CDemoStringTables::decode(payload.as_slice())?;
178                    let mut changed = false;
179                    if self.rewrites_entity_fields() {
180                        changed |= self.rewrite_instance_baselines(&mut msg)?;
181                    }
182                    changed |= self.rewrite_demo_string_table_entries(message.tick, &mut msg)?;
183                    let mut out_payload: Option<Vec<u8>> = None;
184                    let ctx = &self.parser.context;
185                    for rewriter in self.rewriters.iter_mut().filter(|rewriter| {
186                        rewriter
187                            .interests()
188                            .contains(RewriteInterests::DEMO_STRING_TABLES)
189                    }) {
190                        match rewriter.rewrite_demo_string_tables(ctx, message.tick, &mut msg)? {
191                            MessageRewrite::Drop => continue 'messages,
192                            MessageRewrite::Keep => {}
193                            MessageRewrite::Rewrite => {
194                                changed = false;
195                                out_payload = Some(msg.encode_to_vec());
196                            }
197                            MessageRewrite::Replace(bytes) => {
198                                msg = CDemoStringTables::decode(bytes.as_slice())?;
199                                changed = false;
200                                out_payload = Some(bytes);
201                            }
202                        }
203                    }
204                    let payload = out_payload.unwrap_or_else(|| {
205                        if changed {
206                            msg.encode_to_vec()
207                        } else {
208                            payload
209                        }
210                    });
211                    if self.needs_demo_string_table_state() {
212                        self.parser.dem_string_tables(msg)?;
213                    }
214                    self.write_demo_message_with_compression(
215                        message.msg_type,
216                        message.tick,
217                        &payload,
218                        message.compressed,
219                    )?;
220                }
221                _ => {
222                    self.write_message_or_raw(&message, payload.as_deref())?;
223                }
224            }
225        }
226
227        self.finalize_file_info_offset()?;
228        Ok(())
229    }
230
231    fn update_context_tick(&mut self, tick: u32) {
232        self.parser.context.previous_tick = self.parser.context.tick;
233        self.parser.context.tick = tick;
234    }
235}