khadhroony-solana-project/crates/ksp-wallet-lib/src/crypto.rs

// file: crates/ksp-wallet-lib/src/crypto.rs
// version: 5

//! In-memory cryptographic primitives shared by native `.kspwallet` V1/V2 runtime.

use chacha20poly1305::KeyInit; // rust-rules: trait-import
use chacha20poly1305::aead::Aead; // rust-rules: trait-import

/// Exact content-key and password-derived-key size in bytes for the current V1/V2 profiles.
pub(crate) const SECRET_KEY_BYTES: usize = 32;

/// Owned 32-byte secret key with redacted diagnostics and drop-time zeroization.
pub(crate) struct SecretKeyV1 {
    bytes: [u8; SECRET_KEY_BYTES],
}

impl crate::SecretKeyV1 {
    /// Takes ownership of exact 32-byte secret material.
    pub(crate) const fn from_bytes(bytes: [u8; SECRET_KEY_BYTES]) -> Self {
        return Self { bytes };
    }

    /// Generates a fresh secret key from the operating-system CSPRNG.
    pub(crate) fn random() -> ksp_core_lib::Result<Self> {
        let bytes = match random_bytes::<SECRET_KEY_BYTES>() {
            std::result::Result::Ok(bytes) => bytes,
            std::result::Result::Err(error) => return std::result::Result::Err(error),
        };
        return std::result::Result::Ok(Self { bytes });
    }

    /// Borrows secret bytes internally without allocating or cloning.
    pub(crate) const fn as_bytes(&self) -> &[u8; SECRET_KEY_BYTES] {
        return &self.bytes;
    }
}

impl std::fmt::Debug for crate::SecretKeyV1 {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        return formatter.write_str("SecretKeyV1(<redacted>)");
    }
}

impl std::ops::Drop for crate::SecretKeyV1 {
    fn drop(&mut self) {
        zeroize::Zeroize::zeroize(&mut self.bytes);
    }
}

/// Generates a fresh fixed-size byte array from the operating-system CSPRNG.
pub(crate) fn random_bytes<const LENGTH: usize>() -> ksp_core_lib::Result<[u8; LENGTH]> {
    let mut bytes = [0_u8; LENGTH];
    let fill_result = getrandom::fill(bytes.as_mut_slice());
    if fill_result.is_err() {
        zeroize::Zeroize::zeroize(&mut bytes);
        return std::result::Result::Err(randomness_error());
    }
    return std::result::Result::Ok(bytes);
}

/// Generates a fresh XChaCha20-Poly1305 nonce from the operating-system CSPRNG.
pub(crate) fn random_nonce() -> ksp_core_lib::Result<[u8; crate::KSPWALLET_V1_XCHACHA_NONCE_BYTES]> {
    return random_bytes::<{ crate::KSPWALLET_V1_XCHACHA_NONCE_BYTES }>();
}

/// Derives one V1 password wrapping key from serialized Argon2id parameters.
pub(crate) fn derive_password_key(password: &[u8], kdf: &crate::WalletKdfParametersV1) -> ksp_core_lib::Result<crate::SecretKeyV1> {
    return derive_argon2id(password, kdf.salt(), kdf.memory_kib(), kdf.iterations(), kdf.parallelism());
}

/// Derives one V2 password wrapping key from serialized Argon2id parameters.
pub(crate) fn derive_password_key_v2(password: &[u8], kdf: &crate::WalletKdfParametersV2) -> ksp_core_lib::Result<crate::SecretKeyV1> {
    return derive_argon2id(password, kdf.salt(), kdf.memory_kib(), kdf.iterations(), kdf.parallelism());
}

/// Wraps one 32-byte content key with XChaCha20-Poly1305 and caller-provided domain-separated AAD.
pub(crate) fn wrap_key(
    wrapping_key: &crate::SecretKeyV1,
    key_to_wrap: &crate::SecretKeyV1,
    nonce: &[u8; crate::KSPWALLET_V1_XCHACHA_NONCE_BYTES],
    aad: &[u8],
) -> ksp_core_lib::Result<std::vec::Vec<u8>> {
    return crate::encrypt_bytes(wrapping_key, nonce, aad, key_to_wrap.as_bytes());
}

/// Unwraps one 32-byte content key and maps every AEAD authentication failure to the generic Wallet authentication error.
pub(crate) fn unwrap_key(
    wrapping_key: &crate::SecretKeyV1,
    nonce: &[u8; crate::KSPWALLET_V1_XCHACHA_NONCE_BYTES],
    aad: &[u8],
    ciphertext: &[u8],
) -> ksp_core_lib::Result<crate::SecretKeyV1> {
    let plaintext_result = crate::decrypt_bytes(wrapping_key, nonce, aad, ciphertext);
    let mut plaintext = match plaintext_result {
        std::result::Result::Ok(plaintext) => plaintext,
        std::result::Result::Err(error) => return std::result::Result::Err(error),
    };
    if plaintext.len() != SECRET_KEY_BYTES {
        zeroize::Zeroize::zeroize(plaintext.as_mut_slice());
        return std::result::Result::Err(authentication_error());
    }
    let converted = <[u8; SECRET_KEY_BYTES]>::try_from(plaintext.as_slice());
    let key = match converted {
        std::result::Result::Ok(key) => crate::SecretKeyV1::from_bytes(key),
        std::result::Result::Err(_) => {
            zeroize::Zeroize::zeroize(plaintext.as_mut_slice());
            return std::result::Result::Err(authentication_error());
        },
    };
    zeroize::Zeroize::zeroize(plaintext.as_mut_slice());
    return std::result::Result::Ok(key);
}

/// Encrypts bounded plaintext bytes with XChaCha20-Poly1305 and caller-provided domain-separated AAD.
pub(crate) fn encrypt_bytes(
    key: &crate::SecretKeyV1,
    nonce: &[u8; crate::KSPWALLET_V1_XCHACHA_NONCE_BYTES],
    aad: &[u8],
    plaintext: &[u8],
) -> ksp_core_lib::Result<std::vec::Vec<u8>> {
    let cipher_result = chacha20poly1305::XChaCha20Poly1305::new_from_slice(key.as_bytes());
    let cipher = match cipher_result {
        std::result::Result::Ok(cipher) => cipher,
        std::result::Result::Err(_) => return std::result::Result::Err(crypto_parameter_error()),
    };
    let nonce_value = chacha20poly1305::XNonce::from(*nonce);
    let payload = chacha20poly1305::aead::Payload { msg: plaintext, aad };
    let encrypted = cipher.encrypt(&nonce_value, payload);
    return match encrypted {
        std::result::Result::Ok(ciphertext) => std::result::Result::Ok(ciphertext),
        std::result::Result::Err(_) => std::result::Result::Err(crypto_parameter_error()),
    };
}

/// Decrypts authenticated ciphertext bytes and returns a generic authentication error on tag failure.
pub(crate) fn decrypt_bytes(
    key: &crate::SecretKeyV1,
    nonce: &[u8; crate::KSPWALLET_V1_XCHACHA_NONCE_BYTES],
    aad: &[u8],
    ciphertext: &[u8],
) -> ksp_core_lib::Result<std::vec::Vec<u8>> {
    let cipher_result = chacha20poly1305::XChaCha20Poly1305::new_from_slice(key.as_bytes());
    let cipher = match cipher_result {
        std::result::Result::Ok(cipher) => cipher,
        std::result::Result::Err(_) => return std::result::Result::Err(crypto_parameter_error()),
    };
    let nonce_value = chacha20poly1305::XNonce::from(*nonce);
    let payload = chacha20poly1305::aead::Payload { msg: ciphertext, aad };
    let decrypted = cipher.decrypt(&nonce_value, payload);
    return match decrypted {
        std::result::Result::Ok(plaintext) => std::result::Result::Ok(plaintext),
        std::result::Result::Err(_) => std::result::Result::Err(authentication_error()),
    };
}

fn derive_argon2id(password: &[u8], salt: &[u8], memory_kib: u32, iterations: u32, parallelism: u32) -> ksp_core_lib::Result<crate::SecretKeyV1> {
    if password.is_empty() || password.len() > crate::KSPWALLET_V1_MAX_PASSWORD_BYTES {
        return std::result::Result::Err(crypto_parameter_error());
    }
    if salt.len() < crate::KSPWALLET_V1_MIN_KDF_SALT_BYTES || salt.len() > crate::KSPWALLET_V1_MAX_KDF_SALT_BYTES {
        return std::result::Result::Err(crypto_parameter_error());
    }
    if memory_kib == 0
        || memory_kib > crate::KSPWALLET_V1_MAX_ARGON2_MEMORY_KIB
        || iterations == 0
        || iterations > crate::KSPWALLET_V1_MAX_ARGON2_ITERATIONS
        || parallelism == 0
        || parallelism > crate::KSPWALLET_V1_MAX_ARGON2_PARALLELISM
        || memory_kib < parallelism * 8
    {
        return std::result::Result::Err(crypto_parameter_error());
    }
    let params_result = argon2::Params::new(memory_kib, iterations, parallelism, std::option::Option::Some(SECRET_KEY_BYTES));
    let params = match params_result {
        std::result::Result::Ok(params) => params,
        std::result::Result::Err(_) => return std::result::Result::Err(crypto_parameter_error()),
    };
    let argon2 = argon2::Argon2::new(argon2::Algorithm::Argon2id, argon2::Version::V0x13, params);
    let mut output = [0_u8; SECRET_KEY_BYTES];
    let derive_result = argon2.hash_password_into(password, salt, output.as_mut_slice());
    if derive_result.is_err() {
        zeroize::Zeroize::zeroize(&mut output);
        return std::result::Result::Err(crypto_parameter_error());
    }
    return std::result::Result::Ok(crate::SecretKeyV1::from_bytes(output));
}

fn randomness_error() -> ksp_core_lib::Error {
    return ksp_core_lib::Error::new(crate::ERROR_CODE_RANDOMNESS_FAILED, "Wallet cryptographic randomness could not be obtained");
}

fn crypto_parameter_error() -> ksp_core_lib::Error {
    return ksp_core_lib::Error::new(crate::ERROR_CODE_CRYPTO_PARAMETERS_INVALID, "Wallet cryptographic parameters are invalid or unsupported");
}

fn authentication_error() -> ksp_core_lib::Error {
    return ksp_core_lib::Error::new(crate::ERROR_CODE_AUTHENTICATION_FAILED, "Wallet cryptographic authentication failed");
}

#[cfg(test)]
#[path = "../unit_tests/crypto.rs"]
mod tests;