fix(prekey): remove the consumed one-time prekey atomically with the session flush

src/features/signal.rs

@@ -112,6 +112,15 @@ impl<'a> Signal<'a> {
         )
         .await?;
 
+        // A pkmsg consumed prekey is reported, not deleted by the decrypt; buffer
+        // it so the flush below removes it atomically with the promoted session.
+        if let Some(prekey_id) = decrypted.consumed_prekey_id {
+            adapter
+                .pre_key_store
+                .buffer_consumed_prekey(prekey_id, &signal_addr)
+                .await;
+        }
+
         drop(_guard);
         self.client.flush_signal_cache().await?;
 

src/message.rs

@@ -1845,6 +1845,16 @@ impl Client {
 
             match decrypt_res {
                 Ok(decrypted) => {
+                    // Buffer the prekey this pkmsg consumed: message_decrypt promoted
+                    // the session into the (volatile) cache but no longer deletes the
+                    // prekey itself. The post-loop flush deletes it only once that
+                    // session is durable, keeping a crash from orphaning the prekey.
+                    if let Some(prekey_id) = decrypted.consumed_prekey_id {
+                        adapter
+                            .pre_key_store
+                            .buffer_consumed_prekey(prekey_id, &signal_address)
+                            .await;
+                    }
                     if decrypted.identity_change == IdentityChange::ReplacedExisting
                         && !local_identity_reacted
                     {
@@ -1950,6 +1960,12 @@ impl Client {
                                     info.id,
                                     address
                                 );
+                                if let Some(prekey_id) = decrypted.consumed_prekey_id {
+                                    adapter
+                                        .pre_key_store
+                                        .buffer_consumed_prekey(prekey_id, &signal_address)
+                                        .await;
+                                }
                                 // Normally NewOrUnchanged here (the untrusted
                                 // identity was deleted+flushed before the retry),
                                 // but mirror the main-decode gate so a concurrent
@@ -2631,6 +2647,12 @@ impl Client {
                     info.id,
                     info.source.sender
                 );
+                if let Some(prekey_id) = decrypted.consumed_prekey_id {
+                    adapter
+                        .pre_key_store
+                        .buffer_consumed_prekey(prekey_id, signal_address)
+                        .await;
+                }
                 let padded_plaintext = decrypted.plaintext;
                 match self
                     .clone()

src/store/signal_adapter.rs

@@ -223,13 +223,33 @@ impl PreKeyStore for PreKeyAdapter {
             .map_err(signal_err("backend"))
     }
     async fn remove_pre_key(&mut self, prekey_id: PreKeyId) -> Result<(), SignalProtocolError> {
+        // Plain immediate-removal primitive. The inbound pkmsg path does NOT route
+        // through here: message_decrypt reports the consumed prekey and the receive
+        // path buffers it via buffer_consumed_prekey so the durable delete is
+        // atomic with the session flush (matching WAWebSignalProtocolStoreUnifiedApi).
         let device = self.0.device.read().await;
-        WacorePreKeyStore::remove_prekey(&*device, prekey_id.into())
+        device
+            .backend
+            .remove_prekey(prekey_id.into())
             .await
             .map_err(signal_err("backend"))
     }
 }
 
+impl PreKeyAdapter {
+    /// Buffer a consumed one-time prekey for deletion on the next cache flush,
+    /// keyed by the session address whose pkmsg promotion consumed it. Called by
+    /// the inbound receive path after `message_decrypt` reports the consumed
+    /// prekey: the promoted session is still volatile in the cache, so the prekey
+    /// must only be deleted once that session is durably flushed.
+    pub async fn buffer_consumed_prekey(&self, prekey_id: PreKeyId, address: &ProtocolAddress) {
+        self.0
+            .cache
+            .remove_prekey(prekey_id.into(), address.as_str())
+            .await;
+    }
+}
+
 #[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
 #[cfg_attr(not(target_arch = "wasm32"), async_trait)]
 impl SignedPreKeyStore for SignedPreKeyAdapter {
@@ -290,3 +310,82 @@ impl wacore::libsignal::protocol::SenderKeyStore for SenderKeyAdapter {
         self.0.cache.sender_key_lock(sender_key_name).await
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::store::Device;
+    use wacore::store::in_memory::InMemoryBackend;
+
+    const PREKEY_ID: u32 = 7777;
+
+    /// The inbound decrypt path consumes a one-time prekey and buffers it via
+    /// `buffer_consumed_prekey`. It must NOT delete the prekey from the backend
+    /// synchronously: the promoted session is still volatile at that point, so an
+    /// eager backend delete would lose both on a crash. The removal must only be
+    /// committed during the session-bearing cache flush.
+    #[tokio::test]
+    async fn buffer_consumed_prekey_defers_backend_delete_to_flush() {
+        let backend: Arc<dyn crate::store::Backend> = Arc::new(InMemoryBackend::new());
+        backend
+            .store_prekey(PREKEY_ID, b"durable-prekey", false)
+            .await
+            .unwrap();
+
+        let device = Arc::new(RwLock::new(Device::new(backend.clone())));
+        let cache = Arc::new(SignalStoreCache::new());
+        let adapter = SignalProtocolStoreAdapter::new(device, cache.clone());
+
+        let addr = ProtocolAddress::new("bob".to_string(), 1.into());
+        // The real path stores the promoted session before buffering the prekey.
+        cache
+            .put_session(
+                &addr,
+                wacore::libsignal::protocol::SessionRecord::new_fresh(),
+            )
+            .await;
+        adapter
+            .pre_key_store
+            .buffer_consumed_prekey(PREKEY_ID.into(), &addr)
+            .await;
+
+        // Still durable: the removal was only buffered, not written to the backend.
+        assert!(
+            backend.load_prekey(PREKEY_ID).await.unwrap().is_some(),
+            "buffer_consumed_prekey must not delete from the backend before flush"
+        );
+
+        // The flush commits the session AND the buffered prekey removal together.
+        cache.flush(backend.as_ref()).await.unwrap();
+        assert!(
+            backend.load_prekey(PREKEY_ID).await.unwrap().is_none(),
+            "flush must commit the buffered prekey removal"
+        );
+    }
+
+    /// The plain `remove_pre_key` primitive (not used by the inbound consume path)
+    /// removes immediately from the backend.
+    #[tokio::test]
+    async fn remove_pre_key_deletes_immediately() {
+        let backend: Arc<dyn crate::store::Backend> = Arc::new(InMemoryBackend::new());
+        backend
+            .store_prekey(PREKEY_ID, b"durable-prekey", false)
+            .await
+            .unwrap();
+
+        let device = Arc::new(RwLock::new(Device::new(backend.clone())));
+        let cache = Arc::new(SignalStoreCache::new());
+        let mut adapter = SignalProtocolStoreAdapter::new(device, cache.clone());
+
+        adapter
+            .pre_key_store
+            .remove_pre_key(PREKEY_ID.into())
+            .await
+            .unwrap();
+
+        assert!(
+            backend.load_prekey(PREKEY_ID).await.unwrap().is_none(),
+            "remove_pre_key must delete from the backend immediately"
+        );
+    }
+}

wacore/libsignal/src/protocol/session_cipher.rs

@@ -71,6 +71,13 @@ use crate::protocol::{
 pub struct DecryptionResult {
     pub plaintext: Vec<u8>,
     pub identity_change: IdentityChange,
+    /// The one-time pre-key a pkmsg consumed, if any. The decrypt does NOT delete
+    /// it: removing the prekey is the caller's responsibility, and only once the
+    /// promoted session is itself durable. A crash with the prekey already gone
+    /// but the session still volatile makes a redelivered pkmsg undecryptable, so
+    /// the caller buffers this id and deletes it alongside the session flush.
+    /// `None` for a SignalMessage decrypt or a pkmsg that reused an existing session.
+    pub consumed_prekey_id: Option<PreKeyId>,
 }
 
 pub async fn message_encrypt(
@@ -295,13 +302,13 @@ pub async fn message_decrypt_prekey<R: Rng + CryptoRng>(
 
     let (plaintext, pre_key_used, identity_change) = result?;
 
-    if let Some(pre_key_id) = pre_key_used {
-        pre_key_store.remove_pre_key(pre_key_id).await?;
-    }
-
+    // The consumed prekey is reported up, not deleted here: the promoted session
+    // is still volatile in the caller's cache, so the prekey must only be removed
+    // once that session is durable (see DecryptionResult::consumed_prekey_id).
     Ok(DecryptionResult {
         plaintext,
         identity_change,
+        consumed_prekey_id: pre_key_used,
     })
 }
 
@@ -406,6 +413,7 @@ pub async fn message_decrypt_signal<R: Rng + CryptoRng>(
     Ok(DecryptionResult {
         plaintext,
         identity_change,
+        consumed_prekey_id: None,
     })
 }