featherChat/warzone/crates/warzone-wasm/src/lib.rs

//! WASM bridge: exposes warzone-protocol to JavaScript.
//!
//! Gives the web client the EXACT same crypto as the CLI:
//! X25519, ChaCha20-Poly1305, X3DH, Double Ratchet.

use wasm_bindgen::prelude::*;

use warzone_protocol::identity::{IdentityKeyPair, PublicIdentity, Seed};
use warzone_protocol::message::{ReceiptType, WireMessage};
use warzone_protocol::prekey::{
    generate_signed_pre_key, PreKeyBundle,
};
use warzone_protocol::ratchet::RatchetState;
use warzone_protocol::x3dh;
use x25519_dalek::PublicKey;

// ── Identity ──

#[wasm_bindgen]
pub struct WasmIdentity {
    seed_bytes: [u8; 32],
    #[wasm_bindgen(skip)]
    pub identity: IdentityKeyPair,
    #[wasm_bindgen(skip)]
    pub pub_id: PublicIdentity,
    // Pre-key secrets (generated once, reused for decrypt)
    spk_secret_bytes: [u8; 32],
    bundle_cache: Option<Vec<u8>>,
}

#[wasm_bindgen]
impl WasmIdentity {
    #[wasm_bindgen(constructor)]
    pub fn new() -> WasmIdentity {
        let seed = Seed::generate();
        Self::from_seed(seed)
    }

    pub fn from_hex_seed(hex_seed: &str) -> Result<WasmIdentity, JsValue> {
        let bytes = hex::decode(hex_seed).map_err(|e| JsValue::from_str(&e.to_string()))?;
        if bytes.len() != 32 { return Err(JsValue::from_str("seed must be 32 bytes")); }
        let mut seed_bytes = [0u8; 32];
        seed_bytes.copy_from_slice(&bytes);
        Ok(Self::from_seed(Seed::from_bytes(seed_bytes)))
    }

    pub fn fingerprint(&self) -> String { self.pub_id.fingerprint.to_string() }
    pub fn seed_hex(&self) -> String { hex::encode(self.seed_bytes) }
    pub fn fingerprint_hex(&self) -> String { self.pub_id.fingerprint.to_hex() }

    pub fn mnemonic(&self) -> String {
        Seed::from_bytes(self.seed_bytes).to_mnemonic()
    }

    /// Get the pre-key bundle as bincode bytes (for server registration).
    /// The bundle is generated once and cached. The SPK secret is stored internally.
    pub fn bundle_bytes(&mut self) -> Result<Vec<u8>, JsValue> {
        if let Some(ref cached) = self.bundle_cache {
            return Ok(cached.clone());
        }
        let bundle = self.generate_bundle_internal()
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        let bytes = bincode::serialize(&bundle)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        self.bundle_cache = Some(bytes.clone());
        Ok(bytes)
    }

    /// Get the SPK secret as hex (for persistence in localStorage).
    pub fn spk_secret_hex(&self) -> String {
        hex::encode(self.spk_secret_bytes)
    }

    /// Restore the SPK secret from hex (loaded from localStorage).
    pub fn set_spk_secret_hex(&mut self, hex: &str) -> Result<(), JsValue> {
        let bytes = hex::decode(hex).map_err(|e| JsValue::from_str(&e.to_string()))?;
        if bytes.len() != 32 { return Err(JsValue::from_str("SPK secret must be 32 bytes")); }
        self.spk_secret_bytes.copy_from_slice(&bytes);
        Ok(())
    }
}

impl WasmIdentity {
    fn from_seed(seed: Seed) -> Self {
        let seed_bytes = seed.0;
        let identity = seed.derive_identity();
        let pub_id = identity.public_identity();

        // Generate pre-keys ONCE
        let (spk_secret, _) = generate_signed_pre_key(&identity, 1);
        let spk_secret_bytes = spk_secret.to_bytes();

        WasmIdentity {
            seed_bytes,
            identity,
            pub_id,
            spk_secret_bytes,
            bundle_cache: None,
        }
    }

    fn generate_bundle_internal(&self) -> Result<PreKeyBundle, String> {
        // Recreate SPK from stored secret
        let spk_secret = x25519_dalek::StaticSecret::from(self.spk_secret_bytes);
        let spk_public = PublicKey::from(&spk_secret);

        // Sign the SPK public key
        use ed25519_dalek::Signer;
        let signature = self.identity.signing.sign(spk_public.as_bytes());

        let spk = warzone_protocol::prekey::SignedPreKey {
            id: 1,
            public_key: *spk_public.as_bytes(),
            signature: signature.to_bytes().to_vec(),
            timestamp: js_sys::Date::now() as i64 / 1000,
        };

        // No OTPKs for web client (can't store secrets for them reliably).
        // initiate() will skip DH4 when one_time_pre_key is None.
        // This is safe — OTPKs are an anti-replay optimization, not required.
        Ok(PreKeyBundle {
            identity_key: *self.pub_id.signing.as_bytes(),
            identity_encryption_key: *self.pub_id.encryption.as_bytes(),
            signed_pre_key: spk,
            one_time_pre_key: None,
        })
    }
}

// ── Session ──

#[wasm_bindgen]
pub struct WasmSession {
    ratchet: RatchetState,
}

#[wasm_bindgen]
impl WasmSession {
    pub fn initiate(
        identity: &WasmIdentity,
        their_bundle_bytes: &[u8],
    ) -> Result<WasmSession, JsValue> {
        let bundle: PreKeyBundle = bincode::deserialize(their_bundle_bytes)
            .map_err(|e| JsValue::from_str(&format!("bundle: {}", e)))?;
        let result = x3dh::initiate(&identity.identity, &bundle)
            .map_err(|e| JsValue::from_str(&format!("X3DH: {}", e)))?;
        let their_spk = PublicKey::from(bundle.signed_pre_key.public_key);
        Ok(WasmSession {
            ratchet: RatchetState::init_alice(result.shared_secret, their_spk),
        })
    }

    pub fn encrypt_key_exchange(
        &mut self,
        identity: &WasmIdentity,
        their_bundle_bytes: &[u8],
        plaintext: &str,
    ) -> Result<Vec<u8>, JsValue> {
        self.encrypt_key_exchange_with_id(identity, their_bundle_bytes, plaintext, &uuid::Uuid::new_v4().to_string())
    }

    pub fn encrypt_key_exchange_with_id(
        &mut self,
        identity: &WasmIdentity,
        their_bundle_bytes: &[u8],
        plaintext: &str,
        msg_id: &str,
    ) -> Result<Vec<u8>, JsValue> {
        let bundle: PreKeyBundle = bincode::deserialize(their_bundle_bytes)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        let result = x3dh::initiate(&identity.identity, &bundle)
            .map_err(|e| JsValue::from_str(&format!("X3DH: {}", e)))?;

        let encrypted = self.ratchet.encrypt(plaintext.as_bytes())
            .map_err(|e| JsValue::from_str(&format!("encrypt: {}", e)))?;

        let wire = WireMessage::KeyExchange {
            id: msg_id.to_string(),
            sender_fingerprint: identity.pub_id.fingerprint.to_string(),
            sender_identity_encryption_key: *identity.pub_id.encryption.as_bytes(),
            ephemeral_public: *result.ephemeral_public.as_bytes(),
            used_one_time_pre_key_id: result.used_one_time_pre_key_id,
            ratchet_message: encrypted,
        };
        bincode::serialize(&wire).map_err(|e| JsValue::from_str(&e.to_string()))
    }

    pub fn encrypt(&mut self, identity: &WasmIdentity, plaintext: &str) -> Result<Vec<u8>, JsValue> {
        self.encrypt_with_id(identity, plaintext, &uuid::Uuid::new_v4().to_string())
    }

    pub fn encrypt_with_id(&mut self, identity: &WasmIdentity, plaintext: &str, msg_id: &str) -> Result<Vec<u8>, JsValue> {
        let encrypted = self.ratchet.encrypt(plaintext.as_bytes())
            .map_err(|e| JsValue::from_str(&format!("encrypt: {}", e)))?;
        let wire = WireMessage::Message {
            id: msg_id.to_string(),
            sender_fingerprint: identity.pub_id.fingerprint.to_string(),
            ratchet_message: encrypted,
        };
        bincode::serialize(&wire).map_err(|e| JsValue::from_str(&e.to_string()))
    }

    pub fn save(&self) -> Result<String, JsValue> {
        let bytes = bincode::serialize(&self.ratchet).map_err(|e| JsValue::from_str(&e.to_string()))?;
        Ok(base64::Engine::encode(&base64::engine::general_purpose::STANDARD, &bytes))
    }

    pub fn restore(data: &str) -> Result<WasmSession, JsValue> {
        let bytes = base64::Engine::decode(&base64::engine::general_purpose::STANDARD, data)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        let ratchet: RatchetState = bincode::deserialize(&bytes)
            .map_err(|e| JsValue::from_str(&e.to_string()))?;
        Ok(WasmSession { ratchet })
    }
}

// ── Receipt creation ──

/// Create a Receipt wire message (plaintext, not encrypted).
/// `receipt_type`: "delivered" or "read".
/// Returns bincode-serialized bytes.
#[wasm_bindgen]
pub fn create_receipt(
    sender_fingerprint: &str,
    message_id: &str,
    receipt_type: &str,
) -> Result<Vec<u8>, JsValue> {
    let rt = match receipt_type {
        "delivered" => ReceiptType::Delivered,
        "read" => ReceiptType::Read,
        _ => return Err(JsValue::from_str("receipt_type must be 'delivered' or 'read'")),
    };
    let wire = WireMessage::Receipt {
        sender_fingerprint: sender_fingerprint.to_string(),
        message_id: message_id.to_string(),
        receipt_type: rt,
    };
    bincode::serialize(&wire).map_err(|e| JsValue::from_str(&e.to_string()))
}

// ── Self-test (verifies full encrypt/decrypt cycle within WASM) ──

#[wasm_bindgen]
pub fn self_test() -> Result<String, JsValue> {
    // Check randomness works
    let mut rng_test = [0u8; 8];
    rand::RngCore::fill_bytes(&mut rand::rngs::OsRng, &mut rng_test);
    let rng_hex = hex::encode(rng_test);

    // Alice
    let alice_seed = Seed::generate();
    let alice_id = alice_seed.derive_identity();
    let alice_pub = alice_id.public_identity();

    // Bob
    let bob_seed = Seed::generate();
    let bob_id = bob_seed.derive_identity();
    let bob_pub = bob_id.public_identity();

    // Bob's pre-key bundle
    let (bob_spk_secret, bob_spk) = generate_signed_pre_key(&bob_id, 1);
    let bob_spk_secret_bytes = bob_spk_secret.to_bytes();
    let bob_bundle = PreKeyBundle {
        identity_key: *bob_pub.signing.as_bytes(),
        identity_encryption_key: *bob_pub.encryption.as_bytes(),
        signed_pre_key: bob_spk,
        one_time_pre_key: None,
    };
    let _bob_bundle_bytes = bincode::serialize(&bob_bundle)
        .map_err(|e| JsValue::from_str(&e.to_string()))?;

    // Alice initiates X3DH and encrypts
    let x3dh_result = x3dh::initiate(&alice_id, &bob_bundle)
        .map_err(|e| JsValue::from_str(&format!("X3DH initiate: {}", e)))?;
    let their_spk = PublicKey::from(bob_bundle.signed_pre_key.public_key);
    let mut alice_ratchet = RatchetState::init_alice(x3dh_result.shared_secret, their_spk);
    let encrypted = alice_ratchet.encrypt(b"hello from WASM self-test")
        .map_err(|e| JsValue::from_str(&format!("encrypt: {}", e)))?;

    // Clone encrypted for later use (wire takes ownership)
    let encrypted_clone = encrypted.clone();

    let _wire = WireMessage::KeyExchange {
        id: uuid::Uuid::new_v4().to_string(),
        sender_fingerprint: alice_pub.fingerprint.to_string(),
        sender_identity_encryption_key: *alice_pub.encryption.as_bytes(),
        ephemeral_public: *x3dh_result.ephemeral_public.as_bytes(),
        used_one_time_pre_key_id: x3dh_result.used_one_time_pre_key_id,
        ratchet_message: encrypted,
    };
    // Step-by-step Bob-side decrypt (NOT using decrypt_wire_message)
    let alice_shared_hex = hex::encode(x3dh_result.shared_secret);

    // Bob: X3DH respond
    let bob_shared = x3dh::respond(
        &bob_id, &bob_spk_secret, None,
        &alice_pub.encryption, &x3dh_result.ephemeral_public,
    ).map_err(|e| JsValue::from_str(&format!("X3DH respond: {}", e)))?;
    let bob_shared_hex = hex::encode(bob_shared);
    let shared_match = alice_shared_hex == bob_shared_hex;

    // Bob: init ratchet
    // Need a fresh copy of spk_secret (bob_spk_secret was moved into respond)
    let bob_spk_secret2 = x25519_dalek::StaticSecret::from(bob_spk_secret_bytes);
    let mut bob_ratchet = RatchetState::init_bob(bob_shared, bob_spk_secret2);

    // Bob: decrypt
    let decrypt_result = bob_ratchet.decrypt(&encrypted_clone);
    let decrypt_text = match &decrypt_result {
        Ok(plain) => String::from_utf8_lossy(plain).to_string(),
        Err(e) => format!("DECRYPT_ERROR: {}", e),
    };

    Ok(format!(
        "rng={}, shared_match={}, alice_shared={}..., bob_shared={}..., decrypt='{}', PASS={}",
        rng_hex, shared_match, &alice_shared_hex[..16], &bob_shared_hex[..16],
        decrypt_text, decrypt_text == "hello from WASM self-test"
    ))
}

// ── Decrypt ──

/// Debug: dump what the WASM identity's bundle looks like (for comparing with CLI).
#[wasm_bindgen]
pub fn debug_bundle_info(identity: &mut WasmIdentity) -> Result<String, JsValue> {
    let bundle_bytes = identity.bundle_bytes()?;
    let bundle: PreKeyBundle = bincode::deserialize(&bundle_bytes)
        .map_err(|e| JsValue::from_str(&e.to_string()))?;

    let spk_pub_hex = hex::encode(bundle.signed_pre_key.public_key);
    let ik_hex = hex::encode(bundle.identity_key);
    let iek_hex = hex::encode(bundle.identity_encryption_key);
    let spk_secret_hex = identity.spk_secret_hex();

    // Verify SPK matches
    let spk_secret = x25519_dalek::StaticSecret::from(identity.spk_secret_bytes);
    let derived_pub = PublicKey::from(&spk_secret);
    let matches = *derived_pub.as_bytes() == bundle.signed_pre_key.public_key;

    Ok(format!(
        "bundle_size={}, ik={}, iek={}, spk_pub={}, spk_secret={}, spk_matches={}",
        bundle_bytes.len(), &ik_hex[..16], &iek_hex[..16], &spk_pub_hex[..16], &spk_secret_hex[..16], matches
    ))
}

/// Decrypt a bincode WireMessage. `spk_secret_hex` is the signed pre-key secret
/// (stored in localStorage, generated during identity creation).
/// Returns JSON: { "sender": "fp", "text": "...", "new_session": bool, "session_data": "base64...", "message_id": "..." }
/// For Receipt messages: { "type": "receipt", "sender": "fp", "message_id": "...", "receipt_type": "delivered"|"read" }
#[wasm_bindgen]
pub fn decrypt_wire_message(
    identity_hex_seed: &str,
    spk_secret_hex: &str,
    message_bytes: &[u8],
    existing_session_b64: Option<String>,
) -> Result<String, JsValue> {
    let seed_bytes = hex::decode(identity_hex_seed)
        .map_err(|e| JsValue::from_str(&e.to_string()))?;
    let mut sb = [0u8; 32];
    sb.copy_from_slice(&seed_bytes);
    let seed = Seed::from_bytes(sb);
    let id = seed.derive_identity();

    let wire: WireMessage = bincode::deserialize(message_bytes)
        .map_err(|e| JsValue::from_str(&format!("deserialize wire: {}", e)))?;

    match wire {
        WireMessage::KeyExchange {
            id: msg_id,
            sender_fingerprint,
            sender_identity_encryption_key,
            ephemeral_public,
            used_one_time_pre_key_id: _,
            ratchet_message,
        } => {
            // Use the STORED SPK secret, not a regenerated one
            let spk_bytes = hex::decode(spk_secret_hex)
                .map_err(|e| JsValue::from_str(&e.to_string()))?;
            let mut spk_arr = [0u8; 32];
            spk_arr.copy_from_slice(&spk_bytes);
            let spk_secret = x25519_dalek::StaticSecret::from(spk_arr);

            let their_id = PublicKey::from(sender_identity_encryption_key);
            let their_eph = PublicKey::from(ephemeral_public);

            let shared = x3dh::respond(&id, &spk_secret, None, &their_id, &their_eph)
                .map_err(|e| JsValue::from_str(&format!("X3DH respond: {}", e)))?;

            let mut ratchet = RatchetState::init_bob(shared, spk_secret);
            let plain = ratchet.decrypt(&ratchet_message)
                .map_err(|e| JsValue::from_str(&format!("decrypt: {}", e)))?;

            let session_b64 = base64::Engine::encode(
                &base64::engine::general_purpose::STANDARD,
                &bincode::serialize(&ratchet).unwrap_or_default(),
            );

            Ok(serde_json::json!({
                "sender": sender_fingerprint,
                "text": String::from_utf8_lossy(&plain),
                "new_session": true,
                "session_data": session_b64,
                "message_id": msg_id,
            }).to_string())
        }
        WireMessage::Message {
            id: msg_id,
            sender_fingerprint,
            ratchet_message,
        } => {
            let session_data = existing_session_b64
                .ok_or_else(|| JsValue::from_str("no session for this peer"))?;
            let session_bytes = base64::Engine::decode(
                &base64::engine::general_purpose::STANDARD, &session_data,
            ).map_err(|e| JsValue::from_str(&e.to_string()))?;
            let mut ratchet: RatchetState = bincode::deserialize(&session_bytes)
                .map_err(|e| JsValue::from_str(&e.to_string()))?;

            let plain = ratchet.decrypt(&ratchet_message)
                .map_err(|e| JsValue::from_str(&format!("decrypt: {}", e)))?;

            let session_b64 = base64::Engine::encode(
                &base64::engine::general_purpose::STANDARD,
                &bincode::serialize(&ratchet).unwrap_or_default(),
            );

            Ok(serde_json::json!({
                "sender": sender_fingerprint,
                "text": String::from_utf8_lossy(&plain),
                "new_session": false,
                "session_data": session_b64,
                "message_id": msg_id,
            }).to_string())
        }
        WireMessage::Receipt {
            sender_fingerprint,
            message_id,
            receipt_type,
        } => {
            let rt_str = match receipt_type {
                ReceiptType::Delivered => "delivered",
                ReceiptType::Read => "read",
            };
            Ok(serde_json::json!({
                "type": "receipt",
                "sender": sender_fingerprint,
                "message_id": message_id,
                "receipt_type": rt_str,
            }).to_string())
        }
        _ => {
            // File transfer messages not yet handled in WASM
            Ok(serde_json::json!({
                "type": "unsupported",
            }).to_string())
        }
    }
}