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

source2_demo/string_table/
rewrite.rs
