Signal Protocol Implementation

wacore/libsignal/src/
├── protocol/
│   ├── session_cipher.rs      # Encryption/decryption for 1:1 messages
│   ├── group_cipher.rs        # Encryption/decryption for group messages
│   ├── ratchet.rs            # Double Ratchet implementation
│   ├── sender_keys.rs        # Sender Key protocol for groups
│   └── state/                # Session state management
└── crypto/
    ├── aes_cbc.rs            # AES-256-CBC for message content
    ├── aes_gcm.rs            # AES-GCM for media encryption
    └── hash.rs               # HKDF and HMAC primitives

// Alice initiates the session (sender)
pub fn initialize_alice_session<R: Rng + CryptoRng>(
    parameters: &AliceSignalProtocolParameters,
    csprng: &mut R,
) -> Result<SessionState>

// Bob receives the session (recipient)
pub fn initialize_bob_session(
    parameters: &BobSignalProtocolParameters
) -> Result<SessionState>

// From wacore/libsignal/src/protocol/session_cipher.rs:65-183
pub async fn message_encrypt(
    ptext: &[u8],
    remote_address: &ProtocolAddress,
    session_store: &mut dyn SessionStore,
    identity_store: &mut dyn IdentityKeyStore,
) -> Result<CiphertextMessage>

// From wacore/libsignal/src/protocol/session_cipher.rs:292-363
pub async fn message_decrypt_signal<R: Rng + CryptoRng>(
    ciphertext: &SignalMessage,
    remote_address: &ProtocolAddress,
    session_store: &mut dyn SessionStore,
    identity_store: &mut dyn IdentityKeyStore,
    csprng: &mut R,
) -> Result<Vec<u8>>

pub struct ChainKey {
    key: [u8; 32],
    index: u32,
}

impl ChainKey {
    pub fn step_with_message_keys(self) -> Result<(MessageKeyGenerator, ChainKey)> {
        let message_key_gen = MessageKeyGenerator::new(self.key, self.index);
        let next_chain_key = self.next_chain_key()?;
        Ok((message_key_gen, next_chain_key))
    }
}

// 1:1 chain keys (ratchet/keys.rs)
pub fn next_chain_key(&self) -> crate::protocol::Result<Self> {
    Ok(Self {
        key: self.calculate_base_material(Self::CHAIN_KEY_SEED),
        index: self.index.checked_add(1).ok_or_else(|| {
            SignalProtocolError::InvalidState(
                "next_chain_key",
                "chain key index overflow (u32::MAX)".to_string(),
            )
        })?,
    })
}

// Group sender chain keys (sender_keys.rs)
let new_iteration = self.iteration.checked_add(1).ok_or_else(|| {
    SignalProtocolError::InvalidState(
        "sender_chain_key_next",
        "Sender chain is too long".into(),
    )
})?;

const MAX_FORWARD_JUMPS: usize = 25000;

if jump > MAX_FORWARD_JUMPS {
    return Err(SignalProtocolError::InvalidMessage(
        original_message_type,
        "message from too far into the future",
    ));
}

// Build device list — local registry first, network on miss
let mut recipient_cached = self.get_devices_from_registry(&recipient_bare).await;
if recipient_cached.is_none() {
    let _ = self.get_user_devices(std::slice::from_ref(&to)).await;
    recipient_cached = self.get_devices_from_registry(&recipient_bare).await;
}

// Filter hosted devices, exclude sender, dedup for self-DMs
all_dm_jids.retain(|j| !j.is_hosted());
all_dm_jids.retain(|j| !is_sender);
// HashSet dedup for self-DM overlap

fn partition_dm_devices(
    all_devices: Vec<Jid>,
    own_jid: &Jid,
    own_lid: Option<&Jid>,
) -> (Vec<Jid>, Vec<Jid>)

fn is_exact_dm_sender_device(device_jid: &Jid, own_jid: &Jid, own_lid: Option<&Jid>) -> bool {
    (device_jid.is_same_user_as(own_jid) && device_jid.device == own_jid.device)
        || own_lid.is_some_and(|lid|
            device_jid.is_same_user_as(lid) && device_jid.device == lid.device
        )
}

pub fn matches_user_or_lid(&self, user: &Jid, lid: Option<&Jid>) -> bool {
    self.is_same_user_as(user) || lid.is_some_and(|l| self.is_same_user_as(l))
}

pub struct PreparedDmStanza {
    pub node: Node,
    /// Locally computed phash from the sent device set. Not sent on the
    /// wire (WA Web only sends phash for groups). Used by the caller to
    /// compare against the server's ACK phash for device-list drift detection.
    pub phash: Option<String>,
}

// src/client/lid_pn.rs
pub(crate) async fn migrate_signal_sessions_on_lid_discovery(&self, pn: &str, lid: &str) {
    for device_id in 0..=99u16 {
        // Read from signal_cache (authoritative over backend)
        // If PN session exists and no LID session → move session to LID via cache
        // If both exist → delete the stale PN session from cache
        // Identity keys are migrated independently of sessions
    }
    // Flush migrated state to backend so it survives restarts
    self.signal_cache.flush(backend.as_ref()).await;
}

// src/client/sessions.rs
pub(crate) async fn establish_primary_phone_session_immediate(&self) -> Result<()> {
    // Checks LID session → logs PN-only state → defers migration to message path
}

// Decryption path (src/message.rs) — normalize before group_decrypt
let sender_for_sk = info.source.sender.to_non_ad();
let sender_address = sender_for_sk.to_protocol_address();
let sender_key_name = SenderKeyName::from_jid(&info.source.chat, &sender_address);

// SKDM storage path (src/message.rs) — normalize before process_sender_key_distribution_message
let sender_bare = sender_jid.to_non_ad();
let sender_address = sender_bare.to_protocol_address();
let sender_key_name = SenderKeyName::from_jid(&group_jid, &sender_address);

// From wacore/libsignal/src/protocol/group_cipher.rs:283-336
pub async fn create_sender_key_distribution_message<R: Rng + CryptoRng>(
    sender_key_name: &SenderKeyName,
    sender_key_store: &mut dyn SenderKeyStore,
    csprng: &mut R,
) -> Result<SenderKeyDistributionMessage>

// From wacore/libsignal/src/protocol/group_cipher.rs:53-116
pub async fn group_encrypt<R: Rng + CryptoRng>(
    sender_key_store: &mut dyn SenderKeyStore,
    sender_key_name: &SenderKeyName,
    plaintext: &[u8],
    csprng: &mut R,
) -> Result<SenderKeyMessage>

// From wacore/libsignal/src/protocol/group_cipher.rs:162-250
pub async fn group_decrypt(
    skm_bytes: &[u8],
    sender_key_store: &dyn SenderKeyStore,
    sender_key_name: &SenderKeyName,
) -> Result<Vec<u8>>

// src/message.rs — simplified flow
async fn handle_unknown_device_sync(&self, info: &Arc<MessageInfo>) {
    let user_jid = info.source.sender.to_non_ad();
    if !self.pending_device_sync.add(user_jid.clone()).await {
        return; // already queued, dedup
    }
    if info.is_offline {
        return; // batched for deferred flush
    }
    // Online: immediate sync
    self.invalidate_device_cache(&user_jid.user).await;
    self.get_user_devices(&[user_jid]).await.ok();
}

// wacore/src/protocol/retry.rs
pub fn should_drop_unknown_device_retry(keys_present: bool, device_known: bool) -> bool {
    !keys_present && !device_known
}

pub struct PreparedGroupStanza {
    pub node: Node,
    /// Devices that received SKDM in this stanza. Empty when no SKDM was distributed.
    pub skdm_devices: Vec<Jid>,
}

pub(crate) struct SenderKeyDeviceCache {
    inner: Cache<String, Arc<SenderKeyDeviceMap>>,
}

// Atomic get-or-init: concurrent callers for the same group
// share the single database read result
let cached_map = self
    .sender_key_device_cache
    .get_or_init(group_jid, async {
        let db_rows = pm.get_sender_key_devices(group_jid).await.unwrap_or_default();
        Arc::new(SenderKeyDeviceMap::from_db_rows(&db_rows))
    })
    .await;

pub(crate) struct SenderKeyDeviceMap {
    /// user → (device_id → has_key)
    devices: HashMap<Arc<str>, HashMap<u16, bool>>,
    /// Users with at least one has_key=false device
    forgotten_users: HashSet<Arc<str>>,
}

// src/send.rs — simplified phash validation flow

// Group/status path: phash from stanza attribute
let our_phash = stanza.attrs().optional_string("phash").map(|s| s.into_owned());

// DM path: phash from PreparedDmStanza (not on the wire)
let dm_phash = prepared.phash;

// On DM phash mismatch:
if !jid.is_group() && !jid.is_status_broadcast() {
    client.invalidate_device_cache(&jid.user).await;
    if let Some(own_pn) = &client.persistence_manager.get_device_snapshot().pn {
        client.invalidate_device_cache(&own_pn.user).await;
    }
}

pub fn aes_256_cbc_encrypt_into(
    plaintext: &[u8],
    key: &[u8],      // 32 bytes
    iv: &[u8],       // 16 bytes
    output: &mut Vec<u8>,
) -> Result<()>

thread_local! {
    static ENCRYPTION_BUFFER: RefCell<EncryptionBuffer> = ...;
    static DECRYPTION_BUFFER: RefCell<EncryptionBuffer> = ...;
}

// Usage in session_cipher.rs:99-111
let ctext = ENCRYPTION_BUFFER.with(|buffer| {
    let mut buf_wrapper = buffer.borrow_mut();
    let buf = buf_wrapper.get_buffer();
    aes_256_cbc_encrypt_into(ptext, message_keys.cipher_key(), 
                            message_keys.iv(), buf)?;
    let result = std::mem::take(buf);
    buf.reserve(EncryptionBuffer::INITIAL_CAPACITY);
    Ok::<Vec<u8>, SignalProtocolError>(result)
})?;

pub fn derive_keys(secret_input: &[u8]) -> (RootKey, ChainKey, InitialPQRKey) {
    let mut secrets = [0; 96];
    hkdf::Hkdf::<sha2::Sha256>::new(None, secret_input)
        .expand(b"WhisperText", &mut secrets)
        .expect("valid length");
    // Split into RootKey[32], ChainKey[32], PQRKey[32]
}

// Via the Bot builder
Bot::builder()
    .with_wanted_pre_key_count(256) // smaller batches for embedded consumers
    // ...

// Or directly on a Client constructed by hand (before connect)
client.set_wanted_pre_key_count(256);

// Determine starting ID using both the persistent counter AND the store max
let max_id = backend.get_max_prekey_id().await?;
let start_id = if device_snapshot.next_pre_key_id > 0 {
    std::cmp::max(device_snapshot.next_pre_key_id, max_id + 1)
} else {
    // Migration: start from MAX(key_id) + 1
    max_id + 1
};

const MAX_PREKEY_ID: u32 = 16777215; // 2^24 - 1

// Wrap start ID into valid [1, MAX_PREKEY_ID] range
let start_id = ((raw_start as u64 - 1) % MAX_PREKEY_ID as u64) as u32 + 1;

// Each key ID in the batch is also wrapped
let pre_key_id = (((start_id as u64 - 1) + i as u64) % (MAX_PREKEY_ID as u64)) as u32 + 1;

// After upload, the persisted next_pre_key_id wraps too
let next_id = (((start_id as u64 - 1) + key_pairs_to_upload.len() as u64)
    % (MAX_PREKEY_ID as u64)) as u32 + 1;

// After restoring a session from another library
client.refresh_pre_keys().await?;

<!-- Request -->
<iq xmlns="encrypt" type="get" to="s.whatsapp.net" id="...">
  <digest/>
</iq>

<!-- Response -->
<iq from="s.whatsapp.net" id="..." type="result">
  <digest>
    <registration>[4-byte BE registration ID]</registration>
    <type>[1-byte: 5]</type>
    <identity>[32-byte identity public key]</identity>
    <skey>
      <id>[3-byte BE signed pre-key ID]</id>
      <value>[32-byte signed pre-key public]</value>
      <signature>[64-byte signature]</signature>
    </skey>
    <list>
      <id>[3-byte BE prekey ID]</id>
      ...
    </list>
    <hash>[20-byte SHA-1 hash]</hash>
  </digest>
</iq>

src/store/
├── signal.rs               # SignalStore trait impl for Device (identity, session, prekey, sender key)
├── signal_adapter.rs       # SignalProtocolStoreAdapter — cache-backed adapter bridging wacore traits to libsignal traits
└── signal_cache.rs         # Re-export of wacore::store::signal_cache::SignalStoreCache

// SignalProtocolStoreAdapter reads/writes through the cache
#[async_trait]
impl SessionStore for SessionAdapter {
    async fn load_session(
        &self,
        address: &ProtocolAddress,
    ) -> Result<Option<SessionRecord>, SignalProtocolError> {
        // Returns cached SessionRecord object directly (no deserialization)
        // Cold load deserializes from backend bytes once and caches the object
        self.cache.get_session(&addr_str, &*device.backend).await
    }

    async fn store_session(
        &mut self,
        address: &ProtocolAddress,
        record: SessionRecord,  // Takes ownership — zero-cost move
    ) -> Result<(), SignalProtocolError> {
        // Stores the SessionRecord object in cache, marks dirty
        // Serialization happens only during flush()
        self.cache.put_session(&addr_str, record).await;
        Ok(())
    }
}

if !identity_store
    .is_trusted_identity(remote_address, &their_identity_key, 
                         Direction::Sending)
    .await?
{
    return Err(SignalProtocolError::UntrustedIdentity(
        remote_address.clone(),
    ));
}

if chain_index > counter {
    return match state.get_message_keys(their_ephemeral, counter)? {
        Some(keys) => Ok(keys),
        None => Err(SignalProtocolError::DuplicateMessage(chain_index, counter)),
    };
}

fn create_decryption_failure_log(
    remote_address: &ProtocolAddress,
    errs: &[SignalProtocolError],
    record: &SessionRecord,
    ciphertext: &SignalMessage,
) -> Result<String>

// wacore/src/store/signal_cache.rs
enum SessionEntry {
    /// Arc so peek_session (retry / LID-migration checks) bumps a refcount
    /// instead of deep-cloning the record (KBs with archived states).
    Present(Arc<SessionRecord>),
    Absent,
    /// Taken by load_session; has_session treats as present, flush/eviction skip.
    CheckedOut,
}

struct SessionStoreState {
    cache: HashMap<Arc<str>, SessionEntry>,  // Objects, not bytes
    dirty: HashSet<Arc<str>>,
    deleted: HashSet<Arc<str>>,
}

// Sender keys use the same object-caching pattern
struct SenderKeyStoreState {
    cache: HashMap<Arc<str>, Option<SenderKeyRecord>>,  // Objects, not bytes
    dirty: HashSet<Arc<str>>,
}

// wacore/libsignal/src/protocol/storage/traits.rs
pub trait SessionStore {
    async fn store_session(
        &mut self,
        address: &ProtocolAddress,
        record: SessionRecord,  // Owned — zero-cost move, no clone
    ) -> Result<()>;
}

// wacore/libsignal/src/protocol/state/session.rs
pub struct SessionRecord {
    current_session: Option<SessionState>,
    /// Wrapped in Arc so cloning is O(1). Only mutated on rare paths via Arc::make_mut.
    previous_sessions: Arc<Vec<SessionStructure>>,
}

// wacore/src/store/signal_cache.rs — ByteStoreState
fn put_dedup(&mut self, address: &str, data: &[u8]) {
    if let Some(Some(existing)) = self.cache.get(address)
        && existing.as_ref() == data
    {
        return; // Skip — data unchanged, no dirty mark
    }
    self.put(address, data);
}

fn key_for(&self, address: &str) -> Arc<str> {
    match self.cache.get_key_value(address) {
        Some((existing, _)) => existing.clone(),  // O(1) refcount bump
        None => Arc::from(address),                // Only on first insert
    }
}

pub trait JidExt {
    /// Construct a fresh ProtocolAddress for this JID.
    fn to_protocol_address(&self) -> ProtocolAddress;

    /// Signal address string: `{user}[:device]@{server}`
    /// Device part only included when device != 0.
    fn to_signal_address_string(&self) -> String;

    /// Full protocol address string: `{signal_address_string}.0`
    /// Equivalent to `to_protocol_address().to_string()` but avoids the
    /// intermediate ProtocolAddress allocation — one String instead of two.
    fn to_protocol_address_string(&self) -> String;

    /// Rewrite a reusable ProtocolAddress in place for this JID.
    /// See [Reusable hot-loop address construction](#reusable-hot-loop-address-construction).
    fn reset_protocol_address(&self, addr: &mut ProtocolAddress);
}

// In message decryption (src/message.rs) — single lock per sender device
let signal_addr_str = sender_encryption_jid.to_protocol_address_string();
let session_mutex = self.session_locks
    .get_with(signal_addr_str.clone(), async {
        Arc::new(async_lock::Mutex::new(()))
    }).await;
let _session_guard = session_mutex.lock().await;

// In peer message encryption (src/send.rs) — single lock
let signal_addr_str = encryption_jid.to_protocol_address_string();

// In DM message encryption (src/send.rs) — per-device locks for all devices
let lock_jids = self.build_session_lock_keys(&all_dm_devices).await;
let session_mutexes = self.session_mutexes_for(&lock_jids).await;
// Guards acquired in sorted order to prevent deadlocks

// In DM message encryption (src/send.rs)
let lock_jids = self.build_session_lock_keys(&all_dm_devices).await;
let session_mutexes = self.session_mutexes_for(&lock_jids).await;
// Guards acquired in sorted order to prevent deadlocks

// wacore/libsignal/src/core/address.rs
pub struct ProtocolAddress {
    buf: String,       // "{name}.{device_id}" in one buffer
    name_len: usize,   // marks where the name ends
    device_id: DeviceId,
}

impl ProtocolAddress {
    /// One-shot construction: takes ownership of `name` and appends the suffix.
    pub fn new(name: String, device_id: DeviceId) -> Self;

    /// Pre-allocated empty address for hot-loop reuse. Call `reset_with()` to fill.
    pub fn with_capacity(capacity: usize, device_id: DeviceId) -> Self;

    /// Rewrite the address in place via closure. Single write pass — no intermediate copy.
    pub fn reset_with(&mut self, write_name: impl FnOnce(&mut String));

    /// Zero-cost slice of the name portion.
    pub fn name(&self) -> &str;

    /// Zero-cost slice of the full buffer ("{name}.{device_id}").
    pub fn as_str(&self) -> &str;
}

// wacore/src/types/jid.rs
pub fn make_reusable_protocol_address() -> ProtocolAddress {
    ProtocolAddress::with_capacity(64, SIGNAL_DEVICE_ID)
}

pub trait JidExt {
    /// Rewrite a reusable ProtocolAddress in place for this JID.
    fn reset_protocol_address(&self, addr: &mut ProtocolAddress);
}

// wacore/src/send.rs — session resolution loop
let mut reusable_addr = make_reusable_protocol_address();

for device_jid in devices {
    // Rewrite the same buffer — no new allocation per device
    device_jid.reset_protocol_address(&mut reusable_addr);

    if stores.session_store.load_session(&reusable_addr).await?.is_some() {
        // Session exists — use it
    }
}

/// Sort and deduplicate by user identity (user + server).
pub fn sort_dedup_by_user(jids: &mut Vec<Jid>);

/// Sort and deduplicate by device identity (user + server + agent + device).
pub fn sort_dedup_by_device(jids: &mut Vec<Jid>);

// Before device resolution — dedup participants by user identity
sort_dedup_by_user(&mut jids_to_resolve);

// After LID conversion — dedup devices by full device identity
// Catches duplicates where both phone and LID queries resolve
// to the same device (e.g., 559980000003:33 and 100000037037034:33@lid)
sort_dedup_by_device(&mut resolved_list);

// Take ownership instead of cloning
if let Some(mut current_state) = record.take_session_state() {
    let result = decrypt_message_with_state(&mut current_state, ...);
    match result {
        Ok(ptext) => {
            record.set_session_state(current_state);
            return Ok(ptext);
        }
        Err(e) => {
            record.set_session_state(current_state);  // Restore
        }
    }
}

struct EncryptionBuffer {
    buffer: Vec<u8>,
    usage_count: usize,
}

const INITIAL_CAPACITY: usize = 1024;
const MAX_CAPACITY: usize = 16 * 1024;
const SHRINK_THRESHOLD: usize = 100;

fn get_buffer(&mut self) -> &mut Vec<u8> {
    self.usage_count += 1;
    if self.usage_count.is_multiple_of(SHRINK_THRESHOLD) {
        if self.buffer.capacity() > MAX_CAPACITY {
            self.buffer = Vec::with_capacity(INITIAL_CAPACITY);
        }
    }
    &mut self.buffer
}