pub use wacore::appstate::Mutation;
pub use wacore::appstate_sync::{AppStateProcessor, AppStateSyncDriver, AppStateSyncError};
#[cfg(test)]
mod tests {
use super::*;
use async_lock::Mutex;
use async_trait::async_trait;
use prost::Message;
use std::collections::HashMap;
use std::sync::Arc;
use wacore::appstate::WAPATCH_INTEGRITY;
use wacore::appstate::hash::HashState;
use wacore::appstate::hash::generate_content_mac;
use wacore::appstate::keys::expand_app_state_keys;
use wacore::appstate::patch_decode::{PatchList, WAPatchName};
use wacore::appstate::processor::AppStateMutationMAC;
use wacore::libsignal::crypto::aes_256_cbc_encrypt_into;
use wacore::store::error::Result as StoreResult;
use wacore::store::traits::{
AppStateSyncKey, AppSyncStore, DeviceListRecord, DeviceStore, LidPnMappingEntry,
ProtocolStore, SignalStore,
};
use wacore_binary::jid::Jid;
use waproto::whatsapp as wa;
type MockMacMap = Arc<Mutex<HashMap<(String, Vec<u8>), Vec<u8>>>>;
#[derive(Default, Clone)]
struct MockBackend {
versions: Arc<Mutex<HashMap<String, HashState>>>,
macs: MockMacMap,
keys: Arc<Mutex<HashMap<Vec<u8>, AppStateSyncKey>>>,
latest_key_id: Arc<Mutex<Option<Vec<u8>>>>,
}
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl SignalStore for MockBackend {
async fn put_identity(&self, _: &str, _: [u8; 32]) -> StoreResult<()> {
Ok(())
}
async fn load_identity(&self, _: &str) -> StoreResult<Option<Vec<u8>>> {
Ok(None)
}
async fn delete_identity(&self, _: &str) -> StoreResult<()> {
Ok(())
}
async fn get_session(&self, _: &str) -> StoreResult<Option<Vec<u8>>> {
Ok(None)
}
async fn put_session(&self, _: &str, _: &[u8]) -> StoreResult<()> {
Ok(())
}
async fn delete_session(&self, _: &str) -> StoreResult<()> {
Ok(())
}
async fn store_prekey(&self, _: u32, _: &[u8], _: bool) -> StoreResult<()> {
Ok(())
}
async fn load_prekey(&self, _: u32) -> StoreResult<Option<Vec<u8>>> {
Ok(None)
}
async fn remove_prekey(&self, _: u32) -> StoreResult<()> {
Ok(())
}
async fn get_max_prekey_id(&self) -> StoreResult<u32> {
Ok(0)
}
async fn store_signed_prekey(&self, _: u32, _: &[u8]) -> StoreResult<()> {
Ok(())
}
async fn load_signed_prekey(&self, _: u32) -> StoreResult<Option<Vec<u8>>> {
Ok(None)
}
async fn load_all_signed_prekeys(&self) -> StoreResult<Vec<(u32, Vec<u8>)>> {
Ok(vec![])
}
async fn remove_signed_prekey(&self, _: u32) -> StoreResult<()> {
Ok(())
}
async fn put_sender_key(&self, _: &str, _: &[u8]) -> StoreResult<()> {
Ok(())
}
async fn get_sender_key(&self, _: &str) -> StoreResult<Option<Vec<u8>>> {
Ok(None)
}
async fn delete_sender_key(&self, _: &str) -> StoreResult<()> {
Ok(())
}
}
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl AppSyncStore for MockBackend {
async fn get_sync_key(&self, key_id: &[u8]) -> StoreResult<Option<AppStateSyncKey>> {
Ok(self.keys.lock().await.get(key_id).cloned())
}
async fn set_sync_key(&self, key_id: &[u8], key: AppStateSyncKey) -> StoreResult<()> {
self.keys.lock().await.insert(key_id.to_vec(), key);
*self.latest_key_id.lock().await = Some(key_id.to_vec());
Ok(())
}
async fn get_version(&self, name: &str) -> StoreResult<HashState> {
Ok(self
.versions
.lock()
.await
.get(name)
.cloned()
.unwrap_or_default())
}
async fn set_version(&self, name: &str, state: HashState) -> StoreResult<()> {
self.versions.lock().await.insert(name.to_string(), state);
Ok(())
}
async fn put_mutation_macs(
&self,
name: &str,
_version: u64,
mutations: &[AppStateMutationMAC],
) -> StoreResult<()> {
let mut macs = self.macs.lock().await;
for m in mutations {
macs.insert((name.to_string(), m.index_mac.clone()), m.value_mac.clone());
}
Ok(())
}
async fn get_mutation_mac(
&self,
name: &str,
index_mac: &[u8],
) -> StoreResult<Option<Vec<u8>>> {
Ok(self
.macs
.lock()
.await
.get(&(name.to_string(), index_mac.to_vec()))
.cloned())
}
async fn delete_mutation_macs(&self, _: &str, _: &[Vec<u8>]) -> StoreResult<()> {
Ok(())
}
async fn get_latest_sync_key_id(&self) -> StoreResult<Option<Vec<u8>>> {
Ok(self.latest_key_id.lock().await.clone())
}
}
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl ProtocolStore for MockBackend {
async fn get_skdm_recipients(&self, _: &str) -> StoreResult<Vec<Jid>> {
Ok(vec![])
}
async fn add_skdm_recipients(&self, _: &str, _: &[Jid]) -> StoreResult<()> {
Ok(())
}
async fn clear_skdm_recipients(&self, _: &str) -> StoreResult<()> {
Ok(())
}
async fn get_lid_mapping(&self, _: &str) -> StoreResult<Option<LidPnMappingEntry>> {
Ok(None)
}
async fn get_pn_mapping(&self, _: &str) -> StoreResult<Option<LidPnMappingEntry>> {
Ok(None)
}
async fn put_lid_mapping(&self, _: &LidPnMappingEntry) -> StoreResult<()> {
Ok(())
}
async fn get_all_lid_mappings(&self) -> StoreResult<Vec<LidPnMappingEntry>> {
Ok(vec![])
}
async fn save_base_key(&self, _: &str, _: &str, _: &[u8]) -> StoreResult<()> {
Ok(())
}
async fn has_same_base_key(&self, _: &str, _: &str, _: &[u8]) -> StoreResult<bool> {
Ok(false)
}
async fn delete_base_key(&self, _: &str, _: &str) -> StoreResult<()> {
Ok(())
}
async fn update_device_list(&self, _: DeviceListRecord) -> StoreResult<()> {
Ok(())
}
async fn get_devices(&self, _: &str) -> StoreResult<Option<DeviceListRecord>> {
Ok(None)
}
async fn mark_forget_sender_key(&self, _: &str, _: &str) -> StoreResult<()> {
Ok(())
}
async fn consume_forget_marks(&self, _: &str) -> StoreResult<Vec<String>> {
Ok(vec![])
}
async fn get_tc_token(
&self,
_: &str,
) -> StoreResult<Option<wacore::store::traits::TcTokenEntry>> {
Ok(None)
}
async fn put_tc_token(
&self,
_: &str,
_: &wacore::store::traits::TcTokenEntry,
) -> StoreResult<()> {
Ok(())
}
async fn delete_tc_token(&self, _: &str) -> StoreResult<()> {
Ok(())
}
async fn get_all_tc_token_jids(&self) -> StoreResult<Vec<String>> {
Ok(vec![])
}
async fn delete_expired_tc_tokens(&self, _: i64) -> StoreResult<u32> {
Ok(0)
}
async fn store_sent_message(&self, _: &str, _: &str, _: &[u8]) -> StoreResult<()> {
Ok(())
}
async fn take_sent_message(&self, _: &str, _: &str) -> StoreResult<Option<Vec<u8>>> {
Ok(None)
}
async fn delete_expired_sent_messages(&self, _: i64) -> StoreResult<u32> {
Ok(0)
}
}
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
impl DeviceStore for MockBackend {
async fn save(&self, _: &wacore::store::Device) -> StoreResult<()> {
Ok(())
}
async fn load(&self) -> StoreResult<Option<wacore::store::Device>> {
Ok(Some(wacore::store::Device::new()))
}
async fn exists(&self) -> StoreResult<bool> {
Ok(true)
}
async fn create(&self) -> StoreResult<i32> {
Ok(1)
}
}
fn create_encrypted_mutation(
op: wa::syncd_mutation::SyncdOperation,
index_mac: &[u8],
plaintext: &[u8],
keys: &wacore::appstate::keys::ExpandedAppStateKeys,
key_id_bytes: &[u8],
) -> wa::SyncdMutation {
let iv = vec![0u8; 16];
let mut ciphertext = Vec::new();
aes_256_cbc_encrypt_into(plaintext, &keys.value_encryption, &iv, &mut ciphertext)
.expect("AES-CBC encryption should succeed with valid inputs");
let mut value_with_iv = iv;
value_with_iv.extend_from_slice(&ciphertext);
let value_mac = generate_content_mac(op, &value_with_iv, key_id_bytes, &keys.value_mac);
let mut value_blob = value_with_iv;
value_blob.extend_from_slice(&value_mac);
wa::SyncdMutation {
operation: Some(op as i32),
record: Some(wa::SyncdRecord {
index: Some(wa::SyncdIndex {
blob: Some(index_mac.to_vec()),
}),
value: Some(wa::SyncdValue {
blob: Some(value_blob),
}),
key_id: Some(wa::KeyId {
id: Some(key_id_bytes.to_vec()),
}),
}),
}
}
#[tokio::test]
async fn test_process_patch_list_handles_set_overwrite_correctly() {
let backend = Arc::new(MockBackend::default());
let processor =
AppStateProcessor::new(backend.clone(), Arc::new(crate::runtime_impl::TokioRuntime));
let collection_name = WAPatchName::Regular;
let index_mac = vec![1; 32];
let key_id_bytes = b"test_key_id".to_vec();
let master_key = [7u8; 32];
let keys = expand_app_state_keys(&master_key);
let sync_key = AppStateSyncKey {
key_data: master_key.to_vec(),
..Default::default()
};
backend
.set_sync_key(&key_id_bytes, sync_key)
.await
.expect("test backend should accept sync key");
let original_plaintext = wa::SyncActionData {
value: Some(wa::SyncActionValue {
timestamp: Some(1000),
..Default::default()
}),
..Default::default()
}
.encode_to_vec();
let original_mutation = create_encrypted_mutation(
wa::syncd_mutation::SyncdOperation::Set,
&index_mac,
&original_plaintext,
&keys,
&key_id_bytes,
);
let mut initial_state = HashState {
version: 1,
..Default::default()
};
let (hash_result, res) =
initial_state.update_hash(std::slice::from_ref(&original_mutation), |_, _| Ok(None));
assert!(res.is_ok() && !hash_result.has_missing_remove);
backend
.set_version(collection_name.as_str(), initial_state.clone())
.await
.expect("test backend should accept app state version");
let original_value_blob = original_mutation
.record
.expect("mutation should have record")
.value
.expect("record should have value")
.blob
.expect("value should have blob");
let original_value_mac = original_value_blob[original_value_blob.len() - 32..].to_vec();
backend
.put_mutation_macs(
collection_name.as_str(),
1,
&[AppStateMutationMAC {
index_mac: index_mac.clone(),
value_mac: original_value_mac.clone(),
}],
)
.await
.expect("test backend should accept mutation MACs");
let new_plaintext = wa::SyncActionData {
value: Some(wa::SyncActionValue {
timestamp: Some(2000),
..Default::default()
}),
..Default::default()
}
.encode_to_vec();
let overwrite_mutation = create_encrypted_mutation(
wa::syncd_mutation::SyncdOperation::Set,
&index_mac,
&new_plaintext,
&keys,
&key_id_bytes,
);
let patch_list = PatchList {
name: collection_name,
has_more_patches: false,
patches: vec![wa::SyncdPatch {
mutations: vec![overwrite_mutation.clone()],
version: Some(wa::SyncdVersion { version: Some(2) }),
key_id: Some(wa::KeyId {
id: Some(key_id_bytes),
}),
..Default::default()
}],
snapshot: None,
snapshot_ref: None,
error: None,
};
let result = processor.process_patch_list(patch_list, false).await;
assert!(
result.is_ok(),
"Processing the patch should succeed, but it failed: {:?}",
result.err()
);
let (_, final_state, _) = result.expect("process_patch_list should succeed");
let mut expected_state = initial_state.clone();
let new_value_blob = overwrite_mutation
.record
.expect("mutation should have record")
.value
.expect("record should have value")
.blob
.expect("value should have blob");
let new_value_mac = new_value_blob[new_value_blob.len() - 32..].to_vec();
WAPATCH_INTEGRITY.subtract_then_add_in_place(
&mut expected_state.hash,
&[original_value_mac],
&[new_value_mac],
);
assert_eq!(
final_state.hash, expected_state.hash,
"The final LTHash is incorrect, meaning the overwrite was not handled properly."
);
assert_eq!(
final_state.version, 2,
"The version should be updated to that of the patch."
);
}
}