Skip to main content

source2_demo/parser/demo/writer/
run.rs

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