Skip to main content

wacore_appstate/
hash.rs

1use serde::{Deserialize, Serialize};
2use serde_big_array::BigArray;
3use std::collections::HashMap;
4use wacore_libsignal::crypto::CryptographicMac;
5use waproto::whatsapp as wa;
6
7use crate::{AppStateError, WAPATCH_INTEGRITY};
8
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct HashState {
11    pub version: u64,
12    #[serde(with = "BigArray")]
13    pub hash: [u8; 128],
14    pub index_value_map: HashMap<String, Vec<u8>>,
15}
16
17impl Default for HashState {
18    fn default() -> Self {
19        Self {
20            version: 0,
21            hash: [0; 128],
22            index_value_map: HashMap::new(),
23        }
24    }
25}
26
27/// Result of updating the hash state with mutations.
28#[derive(Debug, Clone, Default)]
29pub struct HashUpdateResult {
30    /// Whether a REMOVE mutation was missing its previous value.
31    /// This happens when the server has an entry we don't have locally.
32    /// WhatsApp Web tracks this as `hasMissingRemove` and uses it to
33    /// determine if MAC validation failures should be fatal.
34    pub has_missing_remove: bool,
35}
36
37impl HashState {
38    pub fn update_hash<F>(
39        &mut self,
40        mutations: &[wa::SyncdMutation],
41        mut get_prev_set_value_mac: F,
42    ) -> (HashUpdateResult, anyhow::Result<()>)
43    where
44        F: FnMut(&[u8], usize) -> anyhow::Result<Option<Vec<u8>>>,
45    {
46        let mut added: Vec<Vec<u8>> = Vec::with_capacity(mutations.len());
47        let mut removed: Vec<Vec<u8>> = Vec::with_capacity(mutations.len());
48        let mut result = HashUpdateResult::default();
49
50        for (i, mutation) in mutations.iter().enumerate() {
51            let op = mutation.operation.unwrap_or_default();
52            if op == wa::syncd_mutation::SyncdOperation::Set as i32
53                && let Some(record) = &mutation.record
54                && let Some(value) = &record.value
55                && let Some(blob) = &value.blob
56                && blob.len() >= 32
57            {
58                added.push(blob[blob.len() - 32..].to_vec());
59            }
60            let index_mac_opt = mutation
61                .record
62                .as_ref()
63                .and_then(|r| r.index.as_ref())
64                .and_then(|idx| idx.blob.as_ref());
65            if let Some(index_mac) = index_mac_opt {
66                match get_prev_set_value_mac(index_mac, i) {
67                    Ok(Some(prev)) => removed.push(prev),
68                    Ok(None) => {
69                        if op == wa::syncd_mutation::SyncdOperation::Remove as i32 {
70                            result.has_missing_remove = true;
71                            log::trace!(
72                                target: "AppState",
73                                "REMOVE mutation missing previous value (hasMissingRemove=true)"
74                            );
75                        }
76                    }
77                    Err(e) => return (result, Err(anyhow::anyhow!(e))),
78                }
79            }
80        }
81
82        WAPATCH_INTEGRITY.subtract_then_add_in_place(&mut self.hash, &removed, &added);
83        (result, Ok(()))
84    }
85
86    /// Update hash state from snapshot records directly (avoids cloning into SyncdMutation).
87    ///
88    /// This is an optimized version for snapshots where all operations are SET
89    /// and there are no previous values to look up.
90    pub fn update_hash_from_records(&mut self, records: &[wa::SyncdRecord]) {
91        // Collect slices directly — no Vec<u8> allocation per MAC.
92        let added: Vec<&[u8]> = records
93            .iter()
94            .filter_map(|record| {
95                record
96                    .value
97                    .as_ref()
98                    .and_then(|v| v.blob.as_ref())
99                    .filter(|blob| blob.len() >= 32)
100                    .map(|blob| &blob[blob.len() - 32..])
101            })
102            .collect();
103
104        WAPATCH_INTEGRITY.subtract_then_add_in_place(&mut self.hash, &[] as &[&[u8]], &added);
105    }
106
107    pub fn generate_snapshot_mac(&self, name: &str, key: &[u8]) -> Vec<u8> {
108        let version_be = u64_to_be(self.version);
109        let mut mac =
110            CryptographicMac::new("HmacSha256", key).expect("HmacSha256 is a valid algorithm");
111        mac.update(&self.hash);
112        mac.update(&version_be);
113        mac.update(name.as_bytes());
114        mac.finalize()
115    }
116}
117
118pub fn generate_patch_mac(patch: &wa::SyncdPatch, name: &str, key: &[u8], version: u64) -> Vec<u8> {
119    let mut mac =
120        CryptographicMac::new("HmacSha256", key).expect("HmacSha256 is a valid algorithm");
121
122    // Feed directly to HMAC without collecting into Vec<Vec<u8>>
123    if let Some(sm) = &patch.snapshot_mac {
124        mac.update(sm);
125    }
126    for m in &patch.mutations {
127        if let Some(record) = &m.record
128            && let Some(val) = &record.value
129            && let Some(blob) = &val.blob
130            && blob.len() >= 32
131        {
132            mac.update(&blob[blob.len() - 32..]);
133        }
134    }
135    mac.update(&u64_to_be(version));
136    mac.update(name.as_bytes());
137
138    mac.finalize()
139}
140
141pub fn generate_content_mac(
142    operation: wa::syncd_mutation::SyncdOperation,
143    data: &[u8],
144    key_id: &[u8],
145    key: &[u8],
146) -> Vec<u8> {
147    let op_byte = [operation as u8 + 1];
148    let key_data_length = u64_to_be((key_id.len() + 1) as u64);
149    let mac_full = {
150        let mut mac =
151            CryptographicMac::new("HmacSha512", key).expect("HmacSha512 is a valid algorithm");
152        mac.update(&op_byte);
153        mac.update(key_id);
154        mac.update(data);
155        mac.update(&key_data_length);
156        mac.finalize()
157    };
158    mac_full[..32].to_vec()
159}
160
161fn u64_to_be(val: u64) -> [u8; 8] {
162    val.to_be_bytes()
163}
164
165pub fn validate_index_mac(
166    index_json_bytes: &[u8],
167    expected_mac: &[u8],
168    key: &[u8; 32],
169) -> Result<(), AppStateError> {
170    let computed = {
171        let mut mac =
172            CryptographicMac::new("HmacSha256", key).expect("HmacSha256 is a valid algorithm");
173        mac.update(index_json_bytes);
174        mac.finalize()
175    };
176    if computed.as_slice() != expected_mac {
177        Err(AppStateError::MismatchingIndexMAC)
178    } else {
179        Ok(())
180    }
181}
182
183#[cfg(test)]
184mod tests {
185    use super::*;
186
187    fn create_mutation(
188        operation: wa::syncd_mutation::SyncdOperation,
189        index_mac: Vec<u8>,
190        value_mac: Option<Vec<u8>>,
191    ) -> wa::SyncdMutation {
192        let value_blob = value_mac.map(|mac| {
193            let mut blob = vec![0u8; 16];
194            blob.extend_from_slice(&mac);
195            blob
196        });
197
198        wa::SyncdMutation {
199            operation: Some(operation as i32),
200            record: Some(wa::SyncdRecord {
201                index: Some(wa::SyncdIndex {
202                    blob: Some(index_mac),
203                }),
204                value: value_blob.map(|b| wa::SyncdValue { blob: Some(b) }),
205                key_id: Some(wa::KeyId {
206                    id: Some(b"test_key_id".to_vec()),
207                }),
208            }),
209        }
210    }
211
212    #[test]
213    fn test_update_hash_with_set_overwrite_and_remove() {
214        const INDEX_MAC_1: &[u8] = &[1; 32];
215        const VALUE_MAC_1: &[u8] = &[10; 32];
216
217        const INDEX_MAC_2: &[u8] = &[2; 32];
218        const VALUE_MAC_2: &[u8] = &[20; 32];
219
220        const VALUE_MAC_3_OVERWRITE: &[u8] = &[30; 32];
221
222        let mut prev_macs = HashMap::<Vec<u8>, Vec<u8>>::new();
223
224        let mut state = HashState::default();
225        let initial_mutations = vec![
226            create_mutation(
227                wa::syncd_mutation::SyncdOperation::Set,
228                INDEX_MAC_1.to_vec(),
229                Some(VALUE_MAC_1.to_vec()),
230            ),
231            create_mutation(
232                wa::syncd_mutation::SyncdOperation::Set,
233                INDEX_MAC_2.to_vec(),
234                Some(VALUE_MAC_2.to_vec()),
235            ),
236        ];
237
238        let get_prev_mac_closure = |_: &[u8], _: usize| Ok(None);
239        let (hash_result, result) = state.update_hash(&initial_mutations, get_prev_mac_closure);
240        assert!(result.is_ok());
241        assert!(!hash_result.has_missing_remove);
242
243        const EMPTY: &[Vec<u8>] = &[];
244        let expected_hash_after_add = WAPATCH_INTEGRITY.subtract_then_add(
245            &[0; 128],
246            EMPTY,
247            &[VALUE_MAC_1.to_vec(), VALUE_MAC_2.to_vec()],
248        );
249        assert_eq!(state.hash.as_slice(), expected_hash_after_add.as_slice());
250
251        prev_macs.insert(INDEX_MAC_1.to_vec(), VALUE_MAC_1.to_vec());
252        prev_macs.insert(INDEX_MAC_2.to_vec(), VALUE_MAC_2.to_vec());
253
254        let update_and_remove_mutations = vec![
255            create_mutation(
256                wa::syncd_mutation::SyncdOperation::Set,
257                INDEX_MAC_1.to_vec(),
258                Some(VALUE_MAC_3_OVERWRITE.to_vec()),
259            ),
260            create_mutation(
261                wa::syncd_mutation::SyncdOperation::Remove,
262                INDEX_MAC_2.to_vec(),
263                None,
264            ),
265        ];
266
267        let get_prev_mac_closure_phase2 =
268            |index_mac: &[u8], _: usize| Ok(prev_macs.get(index_mac).cloned());
269        let (hash_result, result) =
270            state.update_hash(&update_and_remove_mutations, get_prev_mac_closure_phase2);
271        assert!(result.is_ok());
272        assert!(!hash_result.has_missing_remove);
273
274        let expected_final_hash = WAPATCH_INTEGRITY.subtract_then_add(
275            &expected_hash_after_add,
276            &[VALUE_MAC_1.to_vec(), VALUE_MAC_2.to_vec()],
277            &[VALUE_MAC_3_OVERWRITE.to_vec()],
278        );
279
280        assert_eq!(
281            state.hash.as_slice(),
282            expected_final_hash.as_slice(),
283            "The final hash state after overwrite and remove is incorrect."
284        );
285    }
286
287    /// Known-answer test for generate_patch_mac to guard byte ordering and input
288    /// concatenation.  The expected MAC was computed by feeding:
289    ///   snapshot_mac ‖ mutation1_tail(32) ‖ mutation2_tail(32)
290    ///   ‖ u64_to_be(42) ‖ b"regular_high"
291    /// into HMAC-SHA256 with key = [0xAA; 32].
292    #[test]
293    fn test_generate_patch_mac_known_answer() {
294        let key = [0xAAu8; 32];
295        let name = "regular_high";
296        let version: u64 = 42;
297
298        // Build a patch with snapshot_mac and two mutations with >=32 byte blobs.
299        let snapshot_mac = vec![0x11u8; 32];
300        let mut blob1 = vec![0u8; 16]; // 16 prefix bytes
301        blob1.extend_from_slice(&[0x22u8; 32]); // 32-byte tail taken by generate_patch_mac
302        let mut blob2 = vec![0u8; 16];
303        blob2.extend_from_slice(&[0x33u8; 32]);
304
305        let patch = wa::SyncdPatch {
306            version: Some(wa::SyncdVersion {
307                version: Some(version),
308            }),
309            snapshot_mac: Some(snapshot_mac.clone()),
310            mutations: vec![
311                wa::SyncdMutation {
312                    operation: Some(wa::syncd_mutation::SyncdOperation::Set as i32),
313                    record: Some(wa::SyncdRecord {
314                        index: None,
315                        value: Some(wa::SyncdValue { blob: Some(blob1) }),
316                        key_id: None,
317                    }),
318                },
319                wa::SyncdMutation {
320                    operation: Some(wa::syncd_mutation::SyncdOperation::Set as i32),
321                    record: Some(wa::SyncdRecord {
322                        index: None,
323                        value: Some(wa::SyncdValue { blob: Some(blob2) }),
324                        key_id: None,
325                    }),
326                },
327            ],
328            ..Default::default()
329        };
330
331        // Compute expected MAC manually using the same HMAC-SHA256 inputs.
332        let mut expected_mac =
333            CryptographicMac::new("HmacSha256", &key).expect("HmacSha256 is a valid algorithm");
334        expected_mac.update(&snapshot_mac); // snapshot_mac
335        expected_mac.update(&[0x22u8; 32]); // mutation 1 tail
336        expected_mac.update(&[0x33u8; 32]); // mutation 2 tail
337        expected_mac.update(&42u64.to_be_bytes()); // version
338        expected_mac.update(b"regular_high"); // name
339        let expected = expected_mac.finalize();
340
341        let actual = generate_patch_mac(&patch, name, &key, version);
342        assert_eq!(
343            actual, expected,
344            "generate_patch_mac output must match manual HMAC-SHA256 computation"
345        );
346    }
347}