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#[derive(Debug, Clone, Default)]
29pub struct HashUpdateResult {
30 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 pub fn update_hash_from_records(&mut self, records: &[wa::SyncdRecord]) {
91 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 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 #[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 let snapshot_mac = vec![0x11u8; 32];
300 let mut blob1 = vec![0u8; 16]; blob1.extend_from_slice(&[0x22u8; 32]); 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 let mut expected_mac =
333 CryptographicMac::new("HmacSha256", &key).expect("HmacSha256 is a valid algorithm");
334 expected_mac.update(&snapshot_mac); expected_mac.update(&[0x22u8; 32]); expected_mac.update(&[0x33u8; 32]); expected_mac.update(&42u64.to_be_bytes()); expected_mac.update(b"regular_high"); 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}