Skip to content

fix: address 9 audit findings across correctness, safety, and performance #460

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jlucaso1 merged 4 commits into from
Mar 29, 2026
+120 −85
Merged

fix: address 9 audit findings across correctness, safety, and performance #460

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
b1acfcb
fix: address 9 audit findings across correctness, safety, and perform…
jlucaso1 Mar 28, 2026
b950542
fix: address review findings in flush durability, cache eviction, and…
jlucaso1 Mar 28, 2026
782515c
fix: flush signal cache per-store to eliminate dirty-set race
jlucaso1 Mar 29, 2026
a65d796
fix: dispatch Disconnected event after cleanup, not before
jlucaso1 Mar 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions src/client.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1255,7 +1255,7 @@ impl Client {
Ok(crate::transport::TransportEvent::DataReceived(data)) => {
// Update dead-socket timer (WA Web: deadSocketTimer reset)
self.last_data_received_ms.store(
wacore::time::now_millis() as u64,
wacore::time::now_millis().max(0) as u64,
Ordering::Relaxed,
);

Expand Down Expand Up @@ -1308,8 +1308,7 @@ impl Client {
}
},
Ok(crate::transport::TransportEvent::Disconnected) | Err(_) => {
self.cleanup_connection_state().await;
if !self.expected_disconnect.load(Ordering::Relaxed) {
if !self.expected_disconnect.load(Ordering::Relaxed) {
self.core.event_bus.dispatch(&Event::Disconnected(crate::types::events::Disconnected));
debug!("Transport disconnected unexpectedly.");
Comment on lines +1322 to 1323

@chatgpt-codex-connector chatgpt-codex-connector Bot Mar 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Clear connection state before firing Disconnected handlers

Event::Disconnected is now dispatched before cleanup_connection_state() runs, so synchronous handlers can observe stale connection state (is_connected, transport/noise handles, caches) and make incorrect decisions (for example, skipping reconnect logic because the client still appears connected during the callback). This regression comes from removing the in-loop cleanup call without preserving the prior cleanup-before-dispatch ordering for unexpected disconnects.

Useful? React with 👍 / 👎.

return Err(anyhow::anyhow!("Transport disconnected unexpectedly"));
Expand Down Expand Up @@ -3318,7 +3317,7 @@ impl Client {

// WA Web: callStanza → deadSocketTimer.onOrBefore(deadSocketTime, socketId)
self.last_data_sent_ms
.store(wacore::time::now_millis() as u64, Ordering::Relaxed);
.store(wacore::time::now_millis().max(0) as u64, Ordering::Relaxed);

Ok(())
}
Expand Down
6 changes: 3 additions & 3 deletions src/handlers/message.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,9 +58,9 @@ impl StanzaHandler for MessageHandler {
let tx = client
.message_queues
.get_with_by_ref(&chat_id, async {
// Create a channel with backpressure
// Increased capacity to handle high message rates without blocking
let (tx, rx) = async_channel::bounded::<Arc<Node>>(10000);
// Bounded capacity provides backpressure to prevent unbounded memory growth.
// 500 is enough for burst handling while limiting per-chat memory.
let (tx, rx) = async_channel::bounded::<Arc<Node>>(500);

let client_for_worker = client.clone();

Expand Down
16 changes: 10 additions & 6 deletions src/message.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,10 @@ impl Client {
/// Increments the retry count for a message and returns the new count.
/// Returns `None` if max retries have been reached.
///
/// Uses get + insert for portability across cache backends.
/// Note: get-then-insert has a theoretical TOCTOU window since
/// `spawn_retry_receipt` detaches. In practice, retries for the same
/// message are rare and a double-send is benign (recipients deduplicate
/// by message ID).
async fn increment_retry_count(&self, cache_key: &str) -> Option<u8> {
let current = self.message_retry_counts.get(&cache_key.to_string()).await;
match current {
Expand Down Expand Up @@ -537,10 +540,10 @@ impl Client {
session_enc_nodes.len()
);

// Skip session processing for group senders (@c.us, @g.us, @broadcast)
// Groups don't use 1:1 Signal Protocol sessions
let is_group_sender = sender_encryption_jid.server.contains(".us")
|| sender_encryption_jid.server.contains("broadcast");
// Skip session processing for group/broadcast JIDs — they use sender keys, not 1:1 sessions.
let is_group_sender = sender_encryption_jid.is_group()
|| sender_encryption_jid.is_broadcast_list()
|| sender_encryption_jid.is_status_broadcast();

let (
session_decrypted_successfully,
Expand Down Expand Up @@ -1034,7 +1037,7 @@ impl Client {
enc_nodes: &[&wacore_binary::node::Node],
info: &MessageInfo,
_sender_encryption_jid: &Jid,
_decrypt_fail_mode: crate::types::events::DecryptFailMode,
decrypt_fail_mode: crate::types::events::DecryptFailMode,
) -> Result<(), DecryptionError> {
if enc_nodes.is_empty() {
return Ok(());
Expand Down Expand Up @@ -1114,6 +1117,7 @@ impl Client {
"No sender key state for group message [msg:{}] from {}: {}. Sending retry receipt.",
info.id, info.source.sender, msg
);
self.dispatch_undecryptable_event(info, decrypt_fail_mode);
self.spawn_retry_receipt(info, RetryReason::NoSession);
}
Err(e) => {
Expand Down
8 changes: 4 additions & 4 deletions src/send.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -871,15 +871,15 @@ impl Client {
}

let force_skdm = {
use wacore::libsignal::protocol::SenderKeyStore;
use wacore::libsignal::store::sender_key_name::SenderKeyName;
let mut device_guard = device_store_arc.write().await;
let sender_address = own_sending_jid.to_protocol_address();
let sender_key_name =
SenderKeyName::new(to_str.clone(), sender_address.to_string());

let key_exists = device_guard
.load_sender_key(&sender_key_name)
let device_guard = device_store_arc.read().await;
let key_exists = self
.signal_cache
.get_sender_key(&sender_key_name, &*device_guard.backend)
.await?
.is_some();

Expand Down
18 changes: 9 additions & 9 deletions wacore/src/protocol/keepalive.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ pub fn ms_since(timestamp_ms: u64) -> Option<u64> {
if timestamp_ms == 0 {
return None;
}
let now = crate::time::now_millis() as u64;
let now = crate::time::now_millis().max(0) as u64;
Some(now.saturating_sub(timestamp_ms))
}

Expand Down Expand Up @@ -61,14 +61,14 @@ mod tests {

#[test]
fn ms_since_recent() {
let now_ms = crate::time::now_millis() as u64;
let now_ms = crate::time::now_millis().max(0) as u64;
let elapsed = ms_since(now_ms).unwrap();
assert!(elapsed < 100, "should be near-zero, got {elapsed}ms");
}

#[test]
fn ms_since_stale() {
let thirty_sec_ago = (crate::time::now_millis() as u64).saturating_sub(30_000);
let thirty_sec_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(30_000);
let elapsed = ms_since(thirty_sec_ago).unwrap();
assert!(
(29_000..=31_000).contains(&elapsed),
Expand All @@ -85,33 +85,33 @@ mod tests {

#[test]
fn dead_socket_received_after_send() {
let t = crate::time::now_millis() as u64;
let t = crate::time::now_millis().max(0) as u64;
assert!(!is_dead_socket(t, t + 1));
}

#[test]
fn dead_socket_sent_recently() {
let now = crate::time::now_millis() as u64;
let now = crate::time::now_millis().max(0) as u64;
assert!(!is_dead_socket(now, 0));
}

#[test]
fn dead_socket_sent_long_ago_no_reply() {
let thirty_ago = (crate::time::now_millis() as u64).saturating_sub(30_000);
let thirty_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(30_000);
assert!(is_dead_socket(thirty_ago, 0));
}

#[test]
fn dead_socket_sent_long_ago_old_reply() {
let thirty_ago = (crate::time::now_millis() as u64).saturating_sub(30_000);
let thirty_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(30_000);
let thirty_one_ago = thirty_ago.saturating_sub(1_000);
assert!(is_dead_socket(thirty_ago, thirty_one_ago));
}

#[test]
fn dead_socket_sent_long_ago_recent_reply() {
let thirty_ago = (crate::time::now_millis() as u64).saturating_sub(30_000);
let one_ago = (crate::time::now_millis() as u64).saturating_sub(1_000);
let thirty_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(30_000);
let one_ago = (crate::time::now_millis().max(0) as u64).saturating_sub(1_000);
assert!(!is_dead_socket(thirty_ago, one_ago));
}

Expand Down
151 changes: 99 additions & 52 deletions wacore/src/store/signal_cache.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -295,72 +295,119 @@ impl SignalStoreCache {
// === Flush ===

/// Flush all dirty state to the backend in a single batch.
/// Acquires all 3 mutexes to ensure consistency (matches WhatsApp Web's pattern).
///
/// Sessions are serialized here (not on every store_session call).
/// Dirty sets are only cleared after ALL writes succeed.
/// Uses a snapshot-then-release pattern: serialize dirty data under the lock,
/// release locks, then write to the backend. This avoids blocking all
/// encrypt/decrypt operations for the duration of I/O.
///
/// Dirty sets are only cleared after ALL writes succeed, preserving retry
/// semantics on partial failure.
pub async fn flush(&self, backend: &dyn SignalStore) -> Result<()> {
let mut sessions = self.sessions.lock().await;
let mut identities = self.identities.lock().await;
let mut sender_keys = self.sender_keys.lock().await;

// Snapshot dirty/deleted sets WITHOUT draining — preserve on failure
let session_dirty: Vec<_> = sessions.dirty.iter().cloned().collect();
let session_deleted: Vec<_> = sessions.deleted.iter().cloned().collect();
let identity_dirty: Vec<_> = identities.dirty.iter().cloned().collect();
let identity_deleted: Vec<_> = identities.deleted.iter().cloned().collect();
let sender_key_dirty: Vec<_> = sender_keys.dirty.iter().cloned().collect();

// Persist dirty sessions — serialize only here, not on every store_session
for address in &session_dirty {
if let Some(Some(record)) = sessions.cache.get(address.as_ref()) {
let bytes = record
.serialize()
.map_err(|e| anyhow::anyhow!("session serialize for {address}: {e}"))?;
backend.put_session(address, &bytes).await?;
// Phase 1: snapshot + serialize under lock, then release.
let (session_writes, session_delete_keys, session_dirty_keys) = {
let state = self.sessions.lock().await;
let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect();
let deleted_keys: Vec<_> = state.deleted.iter().cloned().collect();
let mut writes = Vec::with_capacity(dirty_keys.len());
for address in &dirty_keys {
if let Some(Some(record)) = state.cache.get(address.as_ref()) {
let bytes = record
.serialize()
.map_err(|e| anyhow::anyhow!("session serialize for {address}: {e}"))?;
writes.push((address.clone(), bytes));
}
}
(writes, deleted_keys, dirty_keys)
};

let (identity_writes, identity_delete_keys, identity_dirty_keys) = {
let state = self.identities.lock().await;
let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect();
let deleted_keys: Vec<_> = state.deleted.iter().cloned().collect();
let mut writes = Vec::with_capacity(dirty_keys.len());
for address in &dirty_keys {
if let Some(Some(data)) = state.cache.get(address.as_ref()) {
let key: [u8; 32] = data.as_ref().try_into().map_err(|_| {
anyhow::anyhow!(
"Corrupted identity key for {address}: expected 32 bytes, got {}",
data.len()
)
})?;
writes.push((address.clone(), key));
}
}
(writes, deleted_keys, dirty_keys)
};

let (sender_key_ops, sender_key_dirty_keys) = {
let state = self.sender_keys.lock().await;
let dirty_keys: Vec<_> = state.dirty.iter().cloned().collect();
let mut ops: Vec<(Arc<str>, Option<Vec<u8>>)> = Vec::with_capacity(dirty_keys.len());
for name in &dirty_keys {
match state.cache.get(name.as_ref()) {
Some(Some(record)) => {
let bytes = record
.serialize()
.map_err(|e| anyhow::anyhow!("sender key serialize for {name}: {e}"))?;
ops.push((name.clone(), Some(bytes)));
}
Some(None) => {
ops.push((name.clone(), None));
}
None => {}
}
}
(ops, dirty_keys)
};

// Phase 2: write to backend without holding any locks.
for (address, bytes) in &session_writes {
backend.put_session(address, bytes).await?;
}
for address in &session_deleted {
for address in &session_delete_keys {
backend.delete_session(address).await?;
}

for address in &identity_dirty {
if let Some(Some(data)) = identities.cache.get(address.as_ref()) {
let key: [u8; 32] = data.as_ref().try_into().map_err(|_| {
anyhow::anyhow!(
"Corrupted identity key for {address}: expected 32 bytes, got {}",
data.len()
)
})?;
backend.put_identity(address, key).await?;
}
for (address, key) in &identity_writes {
backend.put_identity(address, *key).await?;
}
for address in &identity_deleted {
for address in &identity_delete_keys {
backend.delete_identity(address).await?;
}

for name in &sender_key_dirty {
match sender_keys.cache.get(name.as_ref()) {
Some(Some(record)) => {
let bytes = record
.serialize()
.map_err(|e| anyhow::anyhow!("sender key serialize for {name}: {e}"))?;
backend.put_sender_key(name, &bytes).await?;
}
Some(None) => {
// Deleted via delete_sender_key — propagate to backend
backend.delete_sender_key(name).await?;
}
None => {}
for (name, bytes_opt) in &sender_key_ops {
match bytes_opt {
Some(bytes) => backend.put_sender_key(name, bytes).await?,
None => backend.delete_sender_key(name).await?,
}
}

// All writes succeeded — clear dirty sets (matches WA Web's clearDirty())
sessions.dirty.clear();
sessions.deleted.clear();
identities.dirty.clear();
identities.deleted.clear();
sender_keys.dirty.clear();
// Phase 3: all writes succeeded — remove only the flushed keys from dirty sets.
// New mutations that occurred during Phase 2 remain in the dirty sets.
{
let mut state = self.sessions.lock().await;
for key in &session_dirty_keys {
state.dirty.remove(key);

@chatgpt-codex-connector chatgpt-codex-connector Bot Mar 28, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Keep keys dirty when modified during flush

In SignalStoreCache::flush, phase 3 unconditionally removes every snapshotted key from dirty, but phase 2 performs backend I/O without holding the store locks. If the same key is updated after phase 1 snapshotting and before this cleanup, that newer update does not add a second marker (because dirty is a set), so remove clears the only dirty flag for the newer value. The backend then only has the stale snapshot written earlier, and the later mutation may never be retried, which can silently lose session/identity/sender-key state across restarts.

Useful? React with 👍 / 👎.

}
for key in &session_delete_keys {
state.deleted.remove(key);
}
}
{
let mut state = self.identities.lock().await;
for key in &identity_dirty_keys {
state.dirty.remove(key);
}
for key in &identity_delete_keys {
state.deleted.remove(key);
}
}
{
let mut state = self.sender_keys.lock().await;
for key in &sender_key_dirty_keys {
state.dirty.remove(key);
}
}

Ok(())
Comment on lines 306 to 387

@coderabbitai coderabbitai Bot Mar 29, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick | 🔵 Trivial

Note partial-failure semantics for future reference.

If, say, sessions flush succeeds but identities flush fails, the sessions dirty set is cleared before the error is returned. On retry, only identities (and sender_keys) will be re-flushed since sessions are already persisted and no longer dirty. This is correct behavior since the session writes did succeed.

This differs slightly from the PR description's "clearing dirty sets only after all writes succeed" (which implies a global all-or-nothing), but per-store clearing is the more practical approach given the independent store design.

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@wacore/src/store/signal_cache.rs` around lines 306 - 387, The current flush
behavior clears each store's dirty set after that store's writes succeed, which
yields per-store partial-failure semantics (e.g., sessions cleared even if
identities later fail) and differs from the PR text claiming atomic "clear dirty
sets only after all writes succeed"; update the code/docs to match intent:
either (A) change the PR description to state per-store clearing semantics, or
(B) modify flush to only clear any dirty/deleted sets after all three sections
succeed by moving the removals out of the per-store blocks and performing them
after all backend calls complete; refer to the flush method and the per-store
states sessions, identities, and sender_keys and their state.dirty/state.deleted
manipulations when making the change.

}
Expand Down
Loading