source2_demo/string_table/
rewrite.rs

1use crate::error::ParserError;
2use crate::proto::{CSvcMsgCreateStringTable, CSvcMsgUpdateStringTable};
3use crate::reader::{BitsReader, SliceReader};
4use crate::writer::{BitsWriter, BitstreamWriter};
5
6/// A mutable string table entry update passed to demo rewriters.
7#[derive(Clone, Debug)]
8pub struct StringTableEntryUpdate {
9    index: i32,
10    key: Option<String>,
11    value: Option<Vec<u8>>,
12    value_compressed: bool,
13}
14
15impl StringTableEntryUpdate {
16    pub(crate) fn new(index: i32, key: Option<String>, value: Option<Vec<u8>>) -> Self {
17        Self {
18            index,
19            key,
20            value,
21            value_compressed: false,
22        }
23    }
24
25    pub(crate) fn new_with_compression(
26        index: i32,
27        key: Option<String>,
28        value: Option<Vec<u8>>,
29        value_compressed: bool,
30    ) -> Self {
31        Self {
32            index,
33            key,
34            value,
35            value_compressed,
36        }
37    }
38
39    pub(crate) fn into_parts(self) -> (i32, Option<String>, Option<Vec<u8>>) {
40        (self.index, self.key, self.value)
41    }
42
43    /// Returns the entry index in the table.
44    pub fn index(&self) -> i32 {
45        self.index
46    }
47
48    /// Returns the entry key, if this update includes one.
49    pub fn key(&self) -> Option<&str> {
50        self.key.as_deref()
51    }
52
53    /// Replaces the entry key.
54    pub fn set_key(&mut self, key: impl Into<String>) {
55        self.key = Some(key.into());
56    }
57
58    /// Clears the entry key from this update.
59    pub fn clear_key(&mut self) {
60        self.key = None;
61    }
62
63    /// Returns the entry value bytes, if present.
64    pub fn value(&self) -> Option<&[u8]> {
65        self.value.as_deref()
66    }
67
68    /// Returns mutable entry value bytes, if present.
69    pub fn value_mut(&mut self) -> Option<&mut Vec<u8>> {
70        self.value.as_mut()
71    }
72
73    /// Replaces the entry value bytes.
74    pub fn set_value(&mut self, value: impl Into<Vec<u8>>) {
75        self.value = Some(value.into());
76    }
77
78    /// Clears the entry value from this update.
79    pub fn clear_value(&mut self) {
80        self.value = None;
81    }
82}
83
84#[derive(Clone, Copy)]
85pub(crate) struct PackedStringTableFormat {
86    user_data_fixed_size: bool,
87    user_data_size: i32,
88    flags: u32,
89    var_int_bit_counts: bool,
90}
91
92impl PackedStringTableFormat {
93    pub(crate) fn from_create_message(msg: &CSvcMsgCreateStringTable) -> Self {
94        Self {
95            user_data_fixed_size: msg.user_data_fixed_size(),
96            user_data_size: msg.user_data_size(),
97            flags: msg.flags() as u32,
98            var_int_bit_counts: msg.using_varint_bitcounts(),
99        }
100    }
101
102    pub(crate) fn from_table(table: &crate::StringTable) -> Self {
103        Self {
104            user_data_fixed_size: table.user_data_fixed_size,
105            user_data_size: table.user_data_size,
106            flags: table.flags,
107            var_int_bit_counts: table.var_int_bit_counts,
108        }
109    }
110}
111
112#[derive(Clone)]
113pub(crate) struct PackedStringTableState {
114    format: PackedStringTableFormat,
115    keys: Vec<String>,
116}
117
118impl PackedStringTableState {
119    pub(crate) fn new(format: PackedStringTableFormat) -> Self {
120        Self {
121            format,
122            keys: vec![String::new(); 32],
123        }
124    }
125
126    pub(crate) fn from_table(table: &crate::StringTable) -> Self {
127        Self {
128            format: PackedStringTableFormat::from_table(table),
129            keys: table.keys.borrow().clone(),
130        }
131    }
132
133    pub(crate) fn rewrite<F>(
134        &mut self,
135        data: &[u8],
136        num_entries: i32,
137        mut rewrite: F,
138    ) -> Result<Option<Vec<u8>>, ParserError>
139    where
140        F: FnMut(&mut StringTableEntryUpdate) -> Result<(), ParserError>,
141    {
142        let entries = self.decode_entries(data, num_entries)?;
143        let mut changed = false;
144        let mut rewritten = Vec::with_capacity(entries.len());
145
146        for entry in entries {
147            let before_key = entry.key.clone();
148            let before_value = entry.value.clone();
149            let mut entry = entry;
150            rewrite(&mut entry)?;
151            changed |= entry.key != before_key || entry.value != before_value;
152            rewritten.push(entry);
153        }
154
155        if changed {
156            Ok(Some(encode_entries(&rewritten, self.format)?))
157        } else {
158            Ok(None)
159        }
160    }
161
162    fn decode_entries(
163        &mut self,
164        data: &[u8],
165        num_entries: i32,
166    ) -> Result<Vec<StringTableEntryUpdate>, ParserError> {
167        let mut reader = SliceReader::new(data);
168        let mut index = -1;
169        let mut delta_pos = 0;
170        let mut entries = Vec::with_capacity(num_entries.max(0) as usize);
171
172        for _ in 0..num_entries {
173            reader.refill();
174
175            index += 1;
176            if !reader.read_bool() {
177                index += reader.read_var_u32() as i32 + 1;
178            }
179
180            let key = reader.read_bool().then(|| {
181                let delta_zero = if delta_pos > 32 { delta_pos & 31 } else { 0 };
182                let key = if reader.read_bool() {
183                    let pos = (delta_zero + reader.read_bits_unchecked(5) as usize) & 31;
184                    let size = reader.read_bits_unchecked(5) as usize;
185
186                    if delta_pos < pos || self.keys[pos].len() < size {
187                        reader.read_cstring()
188                    } else {
189                        self.keys[pos][..size].to_string() + &reader.read_cstring()
190                    }
191                } else {
192                    reader.read_cstring()
193                };
194                self.keys[delta_pos & 31].clone_from(&key);
195                delta_pos += 1;
196                key
197            });
198
199            let mut value_compressed = false;
200            let value = reader.read_bool().then(|| {
201                let bit_size = if self.format.user_data_fixed_size {
202                    self.format.user_data_size as u32
203                } else {
204                    if (self.format.flags & 0x1) != 0 {
205                        value_compressed = reader.read_bool();
206                    }
207                    if self.format.var_int_bit_counts {
208                        reader.read_ubit_var() * 8
209                    } else {
210                        reader.read_bits_unchecked(17) * 8
211                    }
212                };
213
214                let bytes = reader.read_bits_as_bytes(bit_size);
215                if value_compressed {
216                    snap::raw::Decoder::new()
217                        .decompress_vec(&bytes)
218                        .unwrap_or(bytes)
219                } else {
220                    bytes
221                }
222            });
223
224            entries.push(StringTableEntryUpdate::new_with_compression(
225                index,
226                key,
227                value,
228                value_compressed,
229            ));
230        }
231
232        Ok(entries)
233    }
234}
235
236pub(crate) fn rewrite_create_string_table<F>(
237    msg: &mut CSvcMsgCreateStringTable,
238    state: &mut PackedStringTableState,
239    rewrite: F,
240) -> Result<bool, ParserError>
241where
242    F: FnMut(&mut StringTableEntryUpdate) -> Result<(), ParserError>,
243{
244    let data = if msg.data_compressed() {
245        snap::raw::Decoder::new().decompress_vec(msg.string_data())?
246    } else {
247        msg.string_data().to_vec()
248    };
249
250    let Some(rewritten) = state.rewrite(&data, msg.num_entries(), rewrite)? else {
251        return Ok(false);
252    };
253
254    msg.uncompressed_size = Some(rewritten.len() as i32);
255    if msg.data_compressed() {
256        msg.string_data = Some(snap::raw::Encoder::new().compress_vec(&rewritten)?);
257        msg.data_compressed = Some(true);
258    } else {
259        msg.string_data = Some(rewritten);
260    }
261    Ok(true)
262}
263
264pub(crate) fn rewrite_update_string_table<F>(
265    msg: &mut CSvcMsgUpdateStringTable,
266    state: &mut PackedStringTableState,
267    rewrite: F,
268) -> Result<bool, ParserError>
269where
270    F: FnMut(&mut StringTableEntryUpdate) -> Result<(), ParserError>,
271{
272    let Some(rewritten) = state.rewrite(msg.string_data(), msg.num_changed_entries(), rewrite)?
273    else {
274        return Ok(false);
275    };
276
277    msg.string_data = Some(rewritten);
278    Ok(true)
279}
280
281pub(crate) fn rewrite_demo_string_table_items<F>(
282    items: &mut [crate::proto::c_demo_string_tables::ItemsT],
283    mut rewrite: F,
284) -> Result<bool, ParserError>
285where
286    F: FnMut(&mut StringTableEntryUpdate) -> Result<(), ParserError>,
287{
288    let mut changed = false;
289
290    for (index, item) in items.iter_mut().enumerate() {
291        let mut entry =
292            StringTableEntryUpdate::new(index as i32, item.str.clone(), item.data.clone());
293        let before_key = entry.key.clone();
294        let before_value = entry.value.clone();
295        rewrite(&mut entry)?;
296
297        if entry.key != before_key || entry.value != before_value {
298            let (_, key, value) = entry.into_parts();
299            item.str = key;
300            item.data = value;
301            changed = true;
302        }
303    }
304
305    Ok(changed)
306}
307
308fn encode_entries(
309    entries: &[StringTableEntryUpdate],
310    format: PackedStringTableFormat,
311) -> Result<Vec<u8>, ParserError> {
312    let mut out = Vec::new();
313    let mut writer = BitstreamWriter::new(&mut out);
314    let mut previous_index = -1;
315
316    for entry in entries {
317        if entry.index == previous_index + 1 {
318            writer.write_bit(true)?;
319        } else {
320            writer.write_bit(false)?;
321            writer.write_var_u32((entry.index - previous_index - 2) as u32)?;
322        }
323        previous_index = entry.index;
324
325        if let Some(key) = entry.key.as_deref() {
326            writer.write_bit(true)?;
327            writer.write_bit(false)?;
328            writer.write_cstring(key)?;
329        } else {
330            writer.write_bit(false)?;
331        }
332
333        if let Some(value) = entry.value.as_deref() {
334            writer.write_bit(true)?;
335
336            if format.user_data_fixed_size {
337                let expected_bytes = (format.user_data_size as usize).div_ceil(8);
338                if value.len() != expected_bytes {
339                    return Err(ParserError::IoError(format!(
340                        "fixed-size string table entry expected {expected_bytes} bytes, got {}",
341                        value.len()
342                    )));
343                }
344                writer.write_bits_as_bytes(value, format.user_data_size as u32)?;
345            } else {
346                let compressed;
347                let (value, is_compressed) = if (format.flags & 0x1) != 0 {
348                    compressed = snap::raw::Encoder::new().compress_vec(value)?;
349                    if compressed.len() < value.len() || entry.value_compressed {
350                        (compressed.as_slice(), true)
351                    } else {
352                        (value, false)
353                    }
354                } else {
355                    (value, false)
356                };
357
358                if (format.flags & 0x1) != 0 {
359                    writer.write_bit(is_compressed)?;
360                }
361                if format.var_int_bit_counts {
362                    writer.write_ubit_var(value.len() as u32)?;
363                } else {
364                    writer.write_bits(17, value.len() as u64)?;
365                }
366                writer.write_bits_as_bytes(value, (value.len() * 8) as u32)?;
367            }
368        } else {
369            writer.write_bit(false)?;
370        }
371    }
372
373    writer.flush()?;
374    drop(writer);
375    Ok(out)
376}
source2_demo/string_table/rewrite.rs

source2_demo/string_table/
rewrite.rs
