craton-hsm 0.10.5

// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Craton Software Company
use dashmap::DashMap;
use parking_lot::RwLock;
use std::sync::Arc;

use crate::error::{HsmError, HsmResult};
use crate::pkcs11_abi::constants::*;
use crate::pkcs11_abi::types::*;
use crate::session::handle::ObjectHandleAllocator;
use crate::store::encrypted_store::EncryptedStore;
use crate::store::key_material::RawKeyMaterial;
use crate::store::object::StoredObject;

/// Maximum number of objects that can be stored simultaneously.
/// Prevents resource exhaustion via unbounded object creation.
/// 64 attributes × 8 KiB each × 10,000 objects ≈ 5 GiB theoretical max.
const MAX_OBJECTS: usize = 10_000;

/// Maximum allowed size for CKA_VALUE key material in bytes.
/// Largest supported key: 4096-bit RSA private key DER ≈ 2.4 KiB.
/// We allow up to 8 KiB to accommodate future key types with headroom.
const MAX_KEY_MATERIAL_LEN: usize = 8192;

/// In-memory object store with optional persistence.
///
/// When `persist_store` is set, token objects (CKA_TOKEN=true) are
/// automatically persisted to the EncryptedStore on create/insert/destroy.
/// Session objects (CKA_TOKEN=false) are never persisted.
pub struct ObjectStore {
    objects: DashMap<CK_OBJECT_HANDLE, Arc<RwLock<StoredObject>>>,
    handle_alloc: ObjectHandleAllocator,
    /// Optional persistent backend. When present, token objects are
    /// serialized and stored encrypted on disk.
    persist_store: Option<EncryptedStore>,
    /// Encryption key for the persistent store (derived from user PIN).
    /// Set when persistence is enabled and a user logs in.
    /// Wrapped in `Zeroizing` so the key is automatically zeroed on drop,
    /// even if `clear_persist_key()` is not explicitly called.
    persist_key: parking_lot::Mutex<Option<zeroize::Zeroizing<[u8; 32]>>>,
    /// Maps object handles to their opaque store keys (random hex strings).
    /// Prevents sequential handle enumeration in the persistent store.
    handle_to_store_key: parking_lot::Mutex<std::collections::HashMap<CK_OBJECT_HANDLE, String>>,
}

impl Default for ObjectStore {
    fn default() -> Self {
        Self::new()
    }
}

impl ObjectStore {
    pub fn new() -> Self {
        Self {
            objects: DashMap::new(),
            handle_alloc: ObjectHandleAllocator::new(),
            persist_store: None,
            persist_key: parking_lot::Mutex::new(None),
            handle_to_store_key: parking_lot::Mutex::new(std::collections::HashMap::new()),
        }
    }

    /// Return the number of objects currently stored.
    pub fn objects_len(&self) -> usize {
        self.objects.len()
    }

    /// Create an ObjectStore with persistence enabled.
    pub fn with_persistence(store: EncryptedStore) -> Self {
        Self {
            objects: DashMap::new(),
            handle_alloc: ObjectHandleAllocator::new(),
            persist_store: Some(store),
            persist_key: parking_lot::Mutex::new(None),
            handle_to_store_key: parking_lot::Mutex::new(std::collections::HashMap::new()),
        }
    }

    /// Generate an opaque, random store key for a new object.
    /// Prevents sequential handle enumeration in the persistent store.
    ///
    /// Routes through the SP 800-90A HMAC_DRBG (per-call instantiation, matching
    /// the `crypto::keygen` pattern) so the bytes are subject to the approved
    /// DRBG's continuous health-test gating and prediction-resistance reseed.
    /// Falls back to `OsRng` only on DRBG instantiation failure (which would
    /// require OS entropy to be unavailable — already a catastrophic state).
    fn generate_store_key(&self, handle: CK_OBJECT_HANDLE) -> String {
        let mut random_bytes = [0u8; 16];
        match crate::crypto::drbg::HmacDrbg::new()
            .and_then(|mut drbg| drbg.generate(&mut random_bytes))
        {
            Ok(()) => {}
            Err(e) => {
                // DRBG unavailable — fall back to OsRng so we never emit a
                // sequential / predictable store key. This branch is logged
                // so operators can spot a degraded entropy environment.
                use rand::rngs::OsRng;
                use rand::RngCore;
                tracing::warn!(
                    "store-key DRBG instantiation failed ({:?}); \
                     falling back to OsRng for opaque store-key generation",
                    e
                );
                OsRng.fill_bytes(&mut random_bytes);
            }
        }
        let key = format!("obj_{}", hex::encode(random_bytes));
        self.handle_to_store_key.lock().insert(handle, key.clone());
        key
    }

    /// Look up the store key for a handle.
    /// If no mapping exists, generates a new random key to avoid falling
    /// back to a sequential handle-based key that would leak creation order.
    fn get_store_key(&self, handle: CK_OBJECT_HANDLE) -> String {
        let map = self.handle_to_store_key.lock();
        if let Some(key) = map.get(&handle) {
            return key.clone();
        }
        drop(map);
        // No mapping found — generate a random one rather than exposing
        // the sequential handle number in the persistent store.
        self.generate_store_key(handle)
    }

    /// Set the encryption key for persistence (called after user login).
    pub fn set_persist_key(&self, key: [u8; 32]) {
        *self.persist_key.lock() = Some(zeroize::Zeroizing::new(key));
    }

    /// Clear the persistence key (called on logout).
    /// The `Zeroizing` wrapper also guarantees zeroization on drop even if
    /// this method is never called (e.g., process crash).
    pub fn clear_persist_key(&self) {
        // Setting to None drops the Zeroizing<[u8; 32]>, which auto-zeroizes.
        *self.persist_key.lock() = None;
    }

    /// Load persisted token objects from the encrypted store.
    /// Called during initialization after the user logs in and the
    /// persist key is set.
    ///
    /// The persist key is held under lock for the entire load operation
    /// to avoid copying it to the stack where a panic could leak it.
    pub fn load_from_store(&self) -> HsmResult<usize> {
        let store = match &self.persist_store {
            Some(s) => s,
            None => return Ok(0),
        };

        // Hold the lock for the duration of all decryption operations
        // instead of copying the key out, minimizing key exposure.
        let guard = self.persist_key.lock();
        let key = match guard.as_ref() {
            Some(k) => k,
            None => return Ok(0), // No key set yet, can't decrypt
        };

        let keys = store.list_keys()?;
        let mut loaded = 0;

        for store_key in keys {
            if let Some(data) = store.load_encrypted(&store_key, key)? {
                // The codec reads key bytes directly into RawKeyMaterial; it does
                // not use a generic deserializer that could leave heap copies of
                // secret material outside a zeroizing allocation.
                let result = crate::store::object_codec::decode(&data);
                match result {
                    Ok(obj) => {
                        // Defense-in-depth: verify critical security invariants
                        // on deserialized objects. AES-GCM provides authenticity,
                        // but this catches logic bugs or schema migration issues.
                        if let Err(reason) = validate_deserialized_object(&obj) {
                            tracing::warn!(
                                "Rejected deserialized object '{}': {}",
                                store_key,
                                reason
                            );
                            continue;
                        }

                        let handle = obj.handle;
                        // Record the store_key mapping for this handle
                        self.handle_to_store_key.lock().insert(handle, store_key);
                        // Ensure handle allocator is past this handle
                        self.handle_alloc.ensure_past(handle);
                        self.objects.insert(handle, Arc::new(RwLock::new(obj)));
                        loaded += 1;
                    }
                    Err(e) => {
                        tracing::warn!("Failed to deserialize object '{}': {}", store_key, e);
                    }
                }
            }
        }

        Ok(loaded)
    }

    /// Persist a token object to the encrypted store (if enabled).
    ///
    /// Holds the persist key under lock for the entire encryption operation
    /// to avoid copying the key to a stack temporary.
    fn persist_object(&self, obj: &StoredObject) {
        if !obj.token_object {
            return; // Only persist token objects
        }

        let store = match &self.persist_store {
            Some(s) => s,
            None => return,
        };

        // Hold the lock for the entire operation instead of copying key out
        let guard = self.persist_key.lock();
        let key = match guard.as_ref() {
            Some(k) => k,
            None => return, // No key, can't persist
        };

        // Use opaque store key — generate one if this handle is new
        let store_key = {
            let map = self.handle_to_store_key.lock();
            map.get(&obj.handle).cloned()
        }
        .unwrap_or_else(|| self.generate_store_key(obj.handle));

        match crate::store::object_codec::encode(obj) {
            Ok(data) => {
                let result = store.store_encrypted(&store_key, &data, key);
                if let Err(e) = result {
                    tracing::error!("Failed to persist object {}: {:?}", obj.handle, e);
                }
            }
            Err(e) => {
                tracing::error!("Failed to serialize object {}: {}", obj.handle, e);
            }
        }
    }

    /// Persist several token objects under a single store transaction.
    ///
    /// Equivalent to calling [`ObjectStore::persist_object`] on each object,
    /// but pays one `fsync` for the whole set instead of one per object. Used
    /// by [`ObjectStore::insert_objects`].
    ///
    /// Non-token objects in `objs` are skipped, exactly as in the single-object
    /// path. If no token objects remain, this is a no-op and never touches disk.
    fn persist_objects(&self, objs: &[StoredObject]) {
        let store = match &self.persist_store {
            Some(s) => s,
            None => return,
        };

        let guard = self.persist_key.lock();
        let key = match guard.as_ref() {
            Some(k) => k,
            None => return, // No key, can't persist
        };

        let mut items: Vec<(String, zeroize::Zeroizing<Vec<u8>>)> = Vec::with_capacity(objs.len());
        for obj in objs {
            if !obj.token_object {
                continue;
            }
            let store_key = {
                let map = self.handle_to_store_key.lock();
                map.get(&obj.handle).cloned()
            }
            .unwrap_or_else(|| self.generate_store_key(obj.handle));

            match crate::store::object_codec::encode(obj) {
                Ok(data) => items.push((store_key, data)),
                Err(e) => {
                    tracing::error!("Failed to serialize object {}: {}", obj.handle, e);
                }
            }
        }

        if items.is_empty() {
            return;
        }

        if let Err(e) = store.store_encrypted_batch(&items, key) {
            tracing::error!("Failed to persist {} object(s): {:?}", items.len(), e);
        }
    }

    /// Remove a persisted object from the encrypted store.
    fn unpersist_object(&self, handle: CK_OBJECT_HANDLE) {
        let store = match &self.persist_store {
            Some(s) => s,
            None => return,
        };

        let store_key = self.get_store_key(handle);
        if let Err(e) = store.delete(&store_key) {
            tracing::error!("Failed to unpersist object {}: {:?}", handle, e);
        }
        // Remove the mapping
        self.handle_to_store_key.lock().remove(&handle);
    }

    /// Create a new object from a template of attributes
    pub fn create_object(
        &self,
        template: &[(CK_ATTRIBUTE_TYPE, Vec<u8>)],
    ) -> HsmResult<CK_OBJECT_HANDLE> {
        // Enforce maximum object count to prevent resource exhaustion
        if self.objects.len() >= MAX_OBJECTS {
            return Err(HsmError::DeviceMemory);
        }

        // Find class in template (required)
        let class = template
            .iter()
            .find(|(t, _)| *t == CKA_CLASS)
            .and_then(|(_, v)| read_ck_ulong(v))
            .ok_or(HsmError::TemplateIncomplete)?;

        let handle = self.handle_alloc.next()?;
        let mut obj = StoredObject::new(handle, class);

        // Apply all attributes from template
        for (attr_type, value) in template {
            apply_attribute(&mut obj, *attr_type, value)?;
        }

        // If an RSA private key object was created with discrete components instead
        // of CKA_VALUE, synthesize key_material (PKCS#8 DER) so it can be used for
        // cryptographic operations (RFC 8017 / PKCS#11 v3.0 §2.1.3 Table 39).
        if obj.class == CKO_PRIVATE_KEY
            && obj.key_type == Some(CKK_RSA)
            && obj.key_material.is_none()
            && obj.extra_attributes.contains_key(&CKA_PRIVATE_EXPONENT)
        {
            let modulus = obj.modulus.as_ref().ok_or(HsmError::TemplateInconsistent)?;
            let pub_exp = obj
                .public_exponent
                .as_ref()
                .ok_or(HsmError::TemplateInconsistent)?;
            let priv_exp = obj
                .extra_attributes
                .get(&CKA_PRIVATE_EXPONENT)
                .ok_or(HsmError::TemplateInconsistent)?;
            let prime1 = obj
                .extra_attributes
                .get(&CKA_PRIME_1)
                .ok_or(HsmError::TemplateInconsistent)?;
            let prime2 = obj
                .extra_attributes
                .get(&CKA_PRIME_2)
                .ok_or(HsmError::TemplateInconsistent)?;

            use rsa::pkcs8::EncodePrivateKey;
            use rsa::{BigUint, RsaPrivateKey};

            let n = BigUint::from_bytes_be(modulus);
            let e = BigUint::from_bytes_be(pub_exp);
            let d = BigUint::from_bytes_be(priv_exp);
            let p = BigUint::from_bytes_be(prime1);
            let q = BigUint::from_bytes_be(prime2);

            let private_key = RsaPrivateKey::from_components(n, e, d, vec![p, q])
                .map_err(|_| HsmError::AttributeValueInvalid)?;
            private_key
                .validate()
                .map_err(|_| HsmError::AttributeValueInvalid)?;

            let der = private_key
                .to_pkcs8_der()
                .map_err(|_| HsmError::GeneralError)?;

            obj.key_material = Some(RawKeyMaterial::new(der.as_bytes().to_vec()));
            if obj.modulus_bits.is_none() {
                obj.modulus_bits = Some((modulus.len() * 8) as CK_ULONG);
            }
        }

        // Persist if it's a token object
        self.persist_object(&obj);

        self.objects.insert(handle, Arc::new(RwLock::new(obj)));
        Ok(handle)
    }

    /// Insert a pre-built object (used by keygen).
    ///
    /// Returns the handle on success, or `Err(DeviceMemory)` if the store
    /// has reached its maximum object capacity.
    pub fn insert_object(&self, obj: StoredObject) -> HsmResult<CK_OBJECT_HANDLE> {
        // Enforce maximum object count to prevent resource exhaustion
        if self.objects.len() >= MAX_OBJECTS {
            return Err(HsmError::DeviceMemory);
        }

        let handle = obj.handle;

        // Persist if it's a token object
        self.persist_object(&obj);

        self.objects.insert(handle, Arc::new(RwLock::new(obj)));
        Ok(handle)
    }

    /// Insert several pre-built objects, persisting them together.
    ///
    /// This exists because the objects produced by one logical operation must
    /// be persisted as a unit. `C_GenerateKeyPair` is the motivating case: the
    /// public and private objects were previously written with two separate
    /// store commits, so key-pair generation paid two `fsync`s, and a crash
    /// between them could persist the public key without its private half.
    /// Committing them together is both faster and more correct.
    ///
    /// Objects become visible in the in-memory map only after persistence has
    /// been attempted, matching the ordering of [`ObjectStore::insert_object`].
    ///
    /// Returns the handles in the same order as `objs`, or `Err(DeviceMemory)`
    /// if the store cannot accommodate the whole set.
    pub fn insert_objects(&self, objs: Vec<StoredObject>) -> HsmResult<Vec<CK_OBJECT_HANDLE>> {
        if objs.is_empty() {
            return Ok(Vec::new());
        }
        // Check capacity for the whole set up front so a multi-object insert
        // cannot half-succeed against the cap.
        if self.objects.len().saturating_add(objs.len()) > MAX_OBJECTS {
            return Err(HsmError::DeviceMemory);
        }

        let handles: Vec<CK_OBJECT_HANDLE> = objs.iter().map(|o| o.handle).collect();

        self.persist_objects(&objs);

        for obj in objs {
            self.objects.insert(obj.handle, Arc::new(RwLock::new(obj)));
        }
        Ok(handles)
    }

    /// Allocate a new handle
    pub fn next_handle(&self) -> HsmResult<CK_OBJECT_HANDLE> {
        self.handle_alloc.next()
    }

    /// Destroy an object.
    ///
    /// Checks the `destroyable` flag **before** removing from the map to avoid
    /// a TOCTOU race where a non-destroyable object could be permanently lost.
    pub fn destroy_object(&self, handle: CK_OBJECT_HANDLE) -> HsmResult<()> {
        // Atomic check-and-remove: the destroyable check and removal happen
        // while the DashMap shard lock is held, eliminating the TOCTOU window.
        let removed = self
            .objects
            .remove_if(&handle, |_k, v| v.read().destroyable);

        match removed {
            Some(_) => {
                // Successfully removed a destroyable object
                self.unpersist_object(handle);
                Ok(())
            }
            None => {
                // Either the object doesn't exist, or it's not destroyable
                if self.objects.contains_key(&handle) {
                    // Object exists but is not destroyable
                    Err(HsmError::GeneralError)
                } else {
                    Err(HsmError::ObjectHandleInvalid)
                }
            }
        }
    }

    /// Get an object by handle
    pub fn get_object(&self, handle: CK_OBJECT_HANDLE) -> HsmResult<Arc<RwLock<StoredObject>>> {
        self.objects
            .get(&handle)
            .map(|o| o.value().clone())
            .ok_or(HsmError::ObjectHandleInvalid)
    }

    /// Get object size
    pub fn get_object_size(&self, handle: CK_OBJECT_HANDLE) -> HsmResult<CK_ULONG> {
        let obj = self.get_object(handle)?;
        let size = obj.read().approximate_size();
        Ok(size)
    }

    /// Clear all objects (used by C_InitToken per PKCS#11 spec)
    pub fn clear(&self) {
        self.objects.clear();
        self.handle_to_store_key.lock().clear();

        // Also clear the persistent store
        if let Some(ref store) = self.persist_store {
            if let Err(e) = store.clear() {
                tracing::error!("Failed to clear persistent store: {:?}", e);
            }
        }
    }

    /// Clear all objects AND rotate the in-memory persistence key.
    ///
    /// Used by `C_InitToken` to satisfy FIPS 140-3 §7.7 (CSP zeroization
    /// on token reinitialization). The plain [`clear`] only calls
    /// `delete_table` on redb; the encrypted ciphertexts remain in the
    /// underlying data file and the wrapping key is still loaded in
    /// memory. Rotating the persist key to a fresh DRBG-generated value
    /// makes any residual ciphertext on disk cryptographically
    /// inaccessible — the old key is dropped (the `Zeroizing` wrapper
    /// auto-zeroizes it on drop), so neither memory nor disk retains
    /// usable plaintext after this call.
    ///
    /// This is "logical zeroization" as accepted by FIPS 140-3 §7.7:
    /// destroying the wrapping key is sufficient when full physical
    /// erasure of the storage medium is impractical.
    ///
    /// Returns `Err(HsmError::GeneralError)` if DRBG instantiation or
    /// generation fails — in that case the persist key is NOT rotated
    /// (the previous key remains active) so callers can detect the
    /// failure and avoid leaving the store in an inconsistent state.
    /// In-memory and on-disk wipes are still performed before the
    /// rotation attempt, so partial progress is the worst case.
    pub fn clear_and_rotate(&self) -> HsmResult<()> {
        // Step 1: standard wipe (in-memory + redb table delete).
        self.clear();

        // Step 2: only rotate the persist key if one was actually set —
        // a store that was never logged into has nothing to rotate, and
        // generating a new key would create state that confuses later
        // re-derivation from the PIN.
        let had_key = self.persist_key.lock().is_some();
        if !had_key {
            return Ok(());
        }

        // Generate a fresh AES-256-GCM master key from the DRBG.
        // Per project policy, randomness MUST flow through HmacDrbg
        // (health-tested, prediction-resistant) — never OsRng directly.
        let mut new_key = zeroize::Zeroizing::new([0u8; 32]);
        let mut drbg = crate::crypto::drbg::HmacDrbg::new().map_err(|e| {
            tracing::error!("clear_and_rotate: failed to instantiate DRBG: {:?}", e);
            HsmError::GeneralError
        })?;
        drbg.generate(new_key.as_mut()).map_err(|e| {
            tracing::error!("clear_and_rotate: DRBG generate failed: {:?}", e);
            HsmError::GeneralError
        })?;

        // Atomic swap: assigning into the Mutex<Option<Zeroizing<...>>>
        // drops the previous Zeroizing wrapper, which zeroizes the old
        // key bytes in memory before deallocation.
        *self.persist_key.lock() = Some(zeroize::Zeroizing::new(*new_key));
        Ok(())
    }

    /// Check if persistence is enabled
    pub fn has_persistence(&self) -> bool {
        self.persist_store.is_some()
    }

    /// Find objects matching a template, scoped to a specific slot.
    ///
    /// Always iterates *all* objects and performs a full template match on
    /// each one (even those that will be filtered out) to avoid leaking
    /// information about the number or distribution of private objects via
    /// timing side-channels.
    pub fn find_objects(
        &self,
        template: &[(CK_ATTRIBUTE_TYPE, Vec<u8>)],
        is_logged_in: bool,
    ) -> Vec<CK_OBJECT_HANDLE> {
        self.find_objects_for_slot(template, is_logged_in, None)
    }

    /// Find objects matching a template, optionally scoped to a specific slot.
    ///
    /// When `slot_id` is `Some`, only objects belonging to that slot are
    /// returned. This prevents cross-slot object access in multi-slot
    /// deployments.
    ///
    /// Always iterates *all* objects to maintain constant-time behavior.
    pub fn find_objects_for_slot(
        &self,
        template: &[(CK_ATTRIBUTE_TYPE, Vec<u8>)],
        is_logged_in: bool,
        slot_id: Option<CK_ULONG>,
    ) -> Vec<CK_OBJECT_HANDLE> {
        // Pre-allocate to maximum possible size so push() never reallocates.
        // This prevents timing differences between logged-in and logged-out
        // states from leaking information about private object counts.
        let total = self.objects.len();
        let mut results = Vec::with_capacity(total);
        for entry in self.objects.iter() {
            let obj = entry.value().read();
            let matches = obj.matches_template(template);
            let visible = !obj.private || is_logged_in;
            let slot_ok = slot_id.map_or(true, |sid| obj.slot_id == sid);
            if matches && visible && slot_ok {
                results.push(*entry.key());
            }
        }
        results
    }

    /// Sweep all objects, applying date-based lifecycle transitions.
    ///
    /// Write-locks each object individually to check and persist any
    /// SP 800-57 state transitions triggered by `start_date` / `end_date`.
    /// Returns a list of transitions that occurred during the sweep.
    pub fn sweep_lifecycle(&self) -> Vec<crate::store::lifecycle::LifecycleTransition> {
        let mut transitions = Vec::new();
        // Collect objects that need persistence outside the iterator to avoid
        // holding DashMap shard locks during disk I/O.
        let mut to_persist: Vec<crate::store::object::StoredObject> = Vec::new();

        for entry in self.objects.iter() {
            let mut obj = entry.value().write();
            if let Some((old_state, new_state)) = obj.transition_lifecycle_if_needed() {
                let label = String::from_utf8_lossy(&obj.label).to_string();
                transitions.push(crate::store::lifecycle::LifecycleTransition {
                    handle: obj.handle,
                    label,
                    old_state,
                    new_state,
                });
                // Clone while still holding the write lock to prevent TOCTOU:
                // no other thread can modify the object between transition and clone.
                if obj.token_object {
                    to_persist.push(obj.clone());
                }
            }
            // Write lock is dropped here (end of loop iteration)
        }

        // Persist outside the DashMap iterator to avoid lock ordering issues
        for obj_clone in to_persist {
            self.persist_object(&obj_clone);
        }

        transitions
    }

    /// Export all objects as cloned StoredObjects (for backup).
    ///
    /// Returns a `ZeroizingObjects` wrapper that zeroizes all key material
    /// when dropped, preventing callers from accidentally leaking plaintext
    /// key copies.
    pub fn export_all_objects(&self) -> ZeroizingObjects {
        let objects: Vec<StoredObject> = self
            .objects
            .iter()
            .map(|entry| entry.value().read().clone())
            .collect();
        ZeroizingObjects(objects)
    }
}

/// Apply an attribute value to an object during creation.
///
/// **Security note**: this function must only be called during object
/// creation (`create_object`).  Post-creation attribute modification is
/// handled by `C_SetAttributeValue` which enforces `CKA_MODIFIABLE` and
/// restricts the set of mutable attributes to `CKA_LABEL` / `CKA_ID`.
pub fn apply_attribute(
    obj: &mut StoredObject,
    attr_type: CK_ATTRIBUTE_TYPE,
    value: &[u8],
) -> HsmResult<()> {
    match attr_type {
        CKA_CLASS => {
            obj.class = read_ck_ulong(value).ok_or(HsmError::AttributeValueInvalid)?;
        }
        CKA_KEY_TYPE => {
            let key_type = read_ck_ulong(value).ok_or(HsmError::AttributeValueInvalid)?;
            obj.key_type = Some(normalize_key_type(key_type));
        }
        CKA_LABEL => {
            obj.label = value.to_vec();
        }
        CKA_ID => {
            obj.id = value.to_vec();
        }
        CKA_TOKEN => {
            obj.token_object = !value.is_empty() && value[0] != 0;
        }
        CKA_PRIVATE => {
            obj.private = !value.is_empty() && value[0] != 0;
        }
        CKA_SENSITIVE => {
            obj.sensitive = !value.is_empty() && value[0] != 0;
        }
        CKA_EXTRACTABLE => {
            obj.extractable = !value.is_empty() && value[0] != 0;
        }
        CKA_MODIFIABLE => {
            obj.modifiable = !value.is_empty() && value[0] != 0;
        }
        CKA_DESTROYABLE => {
            obj.destroyable = !value.is_empty() && value[0] != 0;
        }
        CKA_ENCRYPT => {
            obj.can_encrypt = !value.is_empty() && value[0] != 0;
        }
        CKA_DECRYPT => {
            obj.can_decrypt = !value.is_empty() && value[0] != 0;
        }
        CKA_SIGN => {
            obj.can_sign = !value.is_empty() && value[0] != 0;
        }
        CKA_VERIFY => {
            obj.can_verify = !value.is_empty() && value[0] != 0;
        }
        CKA_WRAP => {
            obj.can_wrap = !value.is_empty() && value[0] != 0;
        }
        CKA_UNWRAP => {
            obj.can_unwrap = !value.is_empty() && value[0] != 0;
        }
        CKA_DERIVE => {
            obj.can_derive = !value.is_empty() && value[0] != 0;
        }
        CKA_VALUE => {
            if value.len() > MAX_KEY_MATERIAL_LEN {
                return Err(HsmError::AttributeValueInvalid);
            }
            if obj.class == CKO_SECRET_KEY && obj.value_len.is_none() {
                obj.value_len = Some(value.len() as CK_ULONG);
            }
            obj.key_material = Some(RawKeyMaterial::new(value.to_vec()));
        }
        CKA_MODULUS => {
            obj.modulus = Some(value.to_vec());
        }
        CKA_MODULUS_BITS => {
            obj.modulus_bits = Some(read_ck_ulong(value).ok_or(HsmError::AttributeValueInvalid)?);
        }
        CKA_PUBLIC_EXPONENT => {
            obj.public_exponent = Some(value.to_vec());
        }
        CKA_EC_PARAMS => {
            obj.ec_params = Some(value.to_vec());
        }
        CKA_EC_POINT => {
            // Callers supply the DER OCTET STRING form (or, for
            // compatibility, a bare point); store the raw point.
            obj.ec_point = Some(crate::store::ec_point::decode(value).to_vec());
        }
        CKA_VALUE_LEN => {
            obj.value_len = Some(read_ck_ulong(value).ok_or(HsmError::AttributeValueInvalid)?);
        }
        CKA_START_DATE => {
            if value.len() == 8 {
                let mut date = [0u8; 8];
                date.copy_from_slice(value);
                obj.start_date = Some(date);
            } else if value.is_empty() {
                obj.start_date = None;
            } else {
                return Err(HsmError::GeneralError);
            }
        }
        CKA_END_DATE => {
            if value.len() == 8 {
                let mut date = [0u8; 8];
                date.copy_from_slice(value);
                obj.end_date = Some(date);
            } else if value.is_empty() {
                obj.end_date = None;
            } else {
                return Err(HsmError::GeneralError);
            }
        }
        _ => {
            // Limit extra attribute value size and count to prevent resource exhaustion
            if value.len() > 8192 {
                return Err(HsmError::ArgumentsBad);
            }
            if obj.extra_attributes.len() >= 64 && !obj.extra_attributes.contains_key(&attr_type) {
                return Err(HsmError::ArgumentsBad);
            }
            obj.extra_attributes.insert(attr_type, value.to_vec());
        }
    }
    Ok(())
}

/// Apply a post-creation attribute modification (`C_SetAttributeValue`),
/// enforcing PKCS#11 v3.0 §4.9 (Table 39, notes 11–12) one-way rules for
/// `CKA_SENSITIVE` (false -> true only) and `CKA_EXTRACTABLE` (true -> false only).
pub fn modify_attribute(
    obj: &mut StoredObject,
    attr_type: CK_ATTRIBUTE_TYPE,
    value: &[u8],
) -> HsmResult<()> {
    if !obj.modifiable {
        return Err(HsmError::AttributeReadOnly);
    }
    match attr_type {
        CKA_LABEL => {
            obj.label = value.to_vec();
            Ok(())
        }
        CKA_ID => {
            obj.id = value.to_vec();
            Ok(())
        }
        CKA_SENSITIVE => {
            let new_val = !value.is_empty() && value[0] != 0;
            if obj.sensitive && !new_val {
                return Err(HsmError::AttributeReadOnly);
            }
            obj.sensitive = new_val;
            Ok(())
        }
        CKA_EXTRACTABLE => {
            let new_val = !value.is_empty() && value[0] != 0;
            if !obj.extractable && new_val {
                return Err(HsmError::AttributeReadOnly);
            }
            obj.extractable = new_val;
            Ok(())
        }
        _ => Err(HsmError::AttributeReadOnly),
    }
}

/// Read attribute value from a stored object.
///
/// `is_logged_in` is the login state of the calling session. PKCS#11 mandates
/// that objects with `CKA_PRIVATE=true` are invisible at the object level to
/// non-logged-in sessions: any attempt to access such an object — including
/// reading its attributes through a known handle — must behave as if the
/// object does not exist. This layered check guards the `C_GetAttributeValue`
/// path, which calls `read_attribute` directly on a caller-supplied handle
/// and would otherwise bypass the find-objects visibility filter.
///
/// Returns:
/// - `Err(HsmError::ObjectHandleInvalid)` if the object is private and the
///   caller is not logged in (the spec-mandated response — the object should
///   appear not to exist).
/// - `Err(HsmError::AttributeSensitive)` if reading `CKA_VALUE` on a
///   sensitive, non-extractable object.
/// - `Ok(None)` if the attribute is unknown / not present on the object.
pub fn read_attribute(
    obj: &StoredObject,
    attr_type: CK_ATTRIBUTE_TYPE,
    is_logged_in: bool,
) -> Result<Option<Vec<u8>>, HsmError> {
    // PKCS#11 §4.4 / §5.7: a private object must be inaccessible to
    // non-logged-in sessions at the OBJECT level — not just at the
    // find-objects level. Without this guard, a caller holding a handle
    // (e.g. one obtained while logged in and reused after logout) could read
    // CKA_VALUE of a private, non-sensitive, extractable object via
    // C_GetAttributeValue. Return CKR_OBJECT_HANDLE_INVALID so the object
    // appears not to exist, as the spec requires.
    if obj.private && !is_logged_in {
        return Err(HsmError::ObjectHandleInvalid);
    }
    match attr_type {
        CKA_CLASS => Ok(Some(ck_ulong_to_bytes(obj.class))),
        CKA_KEY_TYPE => Ok(obj.key_type.map(ck_ulong_to_bytes)),
        CKA_LABEL => Ok(Some(obj.label.clone())),
        CKA_ID => Ok(Some(obj.id.clone())),
        CKA_TOKEN => Ok(Some(vec![if obj.token_object { 1 } else { 0 }])),
        CKA_PRIVATE => Ok(Some(vec![if obj.private { 1 } else { 0 }])),
        CKA_SENSITIVE => Ok(Some(vec![if obj.sensitive { 1 } else { 0 }])),
        CKA_EXTRACTABLE => Ok(Some(vec![if obj.extractable { 1 } else { 0 }])),
        CKA_MODIFIABLE => Ok(Some(vec![if obj.modifiable { 1 } else { 0 }])),
        CKA_DESTROYABLE => Ok(Some(vec![if obj.destroyable { 1 } else { 0 }])),
        CKA_COPYABLE => Ok(Some(vec![if obj.copyable { 1 } else { 0 }])),
        CKA_ENCRYPT => Ok(Some(vec![if obj.can_encrypt { 1 } else { 0 }])),
        CKA_DECRYPT => Ok(Some(vec![if obj.can_decrypt { 1 } else { 0 }])),
        CKA_SIGN => Ok(Some(vec![if obj.can_sign { 1 } else { 0 }])),
        CKA_VERIFY => Ok(Some(vec![if obj.can_verify { 1 } else { 0 }])),
        CKA_WRAP => Ok(Some(vec![if obj.can_wrap { 1 } else { 0 }])),
        CKA_UNWRAP => Ok(Some(vec![if obj.can_unwrap { 1 } else { 0 }])),
        CKA_DERIVE => Ok(Some(vec![if obj.can_derive { 1 } else { 0 }])),
        CKA_VALUE => {
            // CKA_VALUE access policy (applies to secret/private keys; data,
            // certificate, and public key objects have non-sensitive CKA_VALUE
            // per PKCS#11 v3.0 §4.5–4.7):
            // - sensitive=true AND extractable=false → block (key is fully protected)
            // - sensitive=true AND extractable=true  → allow (key is exportable)
            // - sensitive=false → always allow (regardless of extractable)
            if crate::store::object::is_secret_class(obj.class) && obj.sensitive && !obj.extractable
            {
                return Err(HsmError::AttributeSensitive);
            }
            // PKCS#11 v3.2: CKA_VALUE of an ML-KEM / ML-DSA / SLH-DSA public
            // key is the encoded public key. Generated key pairs keep it in
            // `public_key_data`, not `key_material`.
            if obj.class == CKO_PUBLIC_KEY && obj.key_material.is_none() && obj.is_pqc_key() {
                return Ok(obj.public_key_data.clone());
            }
            Ok(obj.key_material.as_ref().map(|km| km.as_bytes().to_vec()))
        }
        CKA_PARAMETER_SET => Ok(obj.parameter_set().map(ck_ulong_to_bytes)),
        CKA_MODULUS => Ok(obj.modulus.clone()),
        CKA_MODULUS_BITS => Ok(obj.modulus_bits.map(ck_ulong_to_bytes)),
        CKA_PUBLIC_EXPONENT => Ok(obj.public_exponent.clone()),
        CKA_EC_PARAMS => Ok(obj.ec_params.clone()),
        // PKCS#11: CKA_EC_POINT is the DER encoding of the ECPoint
        // (an OCTET STRING around the SEC1 / RFC 8032 public key).
        CKA_EC_POINT => Ok(obj
            .ec_point
            .as_deref()
            .map(crate::store::ec_point::encode_der)),
        CKA_VALUE_LEN => Ok(obj.value_len.map(ck_ulong_to_bytes)),
        CKA_START_DATE => Ok(obj.start_date.map(|d| d.to_vec())),
        CKA_END_DATE => Ok(obj.end_date.map(|d| d.to_vec())),
        t if crate::store::object::SENSITIVE_COMPONENT_ATTRIBUTES.contains(&t) => {
            if obj.class == CKO_PRIVATE_KEY {
                if obj.sensitive || !obj.extractable {
                    return Err(HsmError::AttributeSensitive);
                }
                Ok(obj.extra_attributes.get(&t).cloned())
            } else {
                Ok(None)
            }
        }
        _ => Ok(obj.extra_attributes.get(&attr_type).cloned()),
    }
}

/// Parse a `CK_ULONG` from raw bytes. Returns `None` if the input is
/// too short, instead of silently returning 0 (which could cause false
/// matches or incorrect attribute values).
pub fn read_ck_ulong(bytes: &[u8]) -> Option<CK_ULONG> {
    let size = std::mem::size_of::<CK_ULONG>();
    if bytes.len() < size {
        return None;
    }
    let mut buf = [0u8; std::mem::size_of::<CK_ULONG>()];
    buf.copy_from_slice(&bytes[..size]);
    Some(CK_ULONG::from_ne_bytes(buf))
}

fn ck_ulong_to_bytes(val: CK_ULONG) -> Vec<u8> {
    val.to_ne_bytes().to_vec()
}

/// Defense-in-depth: validate that a deserialized `StoredObject` has sane
/// security-critical fields. AES-GCM ensures authenticity of the ciphertext,
/// but this guards against schema migration bugs, serialization quirks, or
/// accidental corruption that produces a valid-but-dangerous object.
fn validate_deserialized_object(obj: &StoredObject) -> Result<(), &'static str> {
    // Handle 0 is CK_INVALID_HANDLE in PKCS#11
    if obj.handle == 0 {
        return Err("handle is CK_INVALID_HANDLE (0)");
    }

    // Note: a key that is both non-sensitive and non-extractable is valid
    // PKCS#11 (CKA_VALUE is readable; only wrapping is forbidden), and the
    // creation paths accept it. Rejecting it here made such keys vanish from
    // the token on the next load, so it is deliberately not checked.

    // Lifecycle state "Destroyed" objects should never be persisted
    if obj.lifecycle_state == crate::store::object::KeyLifecycleState::Destroyed {
        return Err("persisted object has Destroyed lifecycle state");
    }

    Ok(())
}

/// Wrapper around `Vec<StoredObject>` that ensures all cloned key material
/// and sensitive metadata is zeroized when the collection is dropped.
pub struct ZeroizingObjects(Vec<StoredObject>);

impl std::ops::Deref for ZeroizingObjects {
    type Target = [StoredObject];
    fn deref(&self) -> &[StoredObject] {
        &self.0
    }
}

impl Drop for ZeroizingObjects {
    fn drop(&mut self) {
        use zeroize::Zeroize;
        for obj in self.0.iter_mut() {
            // RawKeyMaterial is zeroized by its own Drop impl, but we must
            // also zeroize other fields that may contain sensitive data.
            obj.label.zeroize();
            obj.id.zeroize();
            if let Some(ref mut pk) = obj.public_key_data {
                pk.zeroize();
            }
            if let Some(ref mut m) = obj.modulus {
                m.zeroize();
            }
            if let Some(ref mut e) = obj.public_exponent {
                e.zeroize();
            }
            if let Some(ref mut p) = obj.ec_params {
                p.zeroize();
            }
            if let Some(ref mut p) = obj.ec_point {
                p.zeroize();
            }
            for val in obj.extra_attributes.values_mut() {
                val.zeroize();
            }
        }
        // Now drop all elements (which triggers RawKeyMaterial::drop → zeroize key bytes)
        self.0.clear();
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn rsa_private_components_sensitivity() {
        let mut obj = StoredObject::new(1, CKO_PRIVATE_KEY);
        obj.key_type = Some(CKK_RSA);
        obj.sensitive = true;
        obj.extractable = false;
        obj.extra_attributes
            .insert(CKA_PRIVATE_EXPONENT, vec![1, 2, 3]);
        obj.extra_attributes.insert(CKA_PRIME_1, vec![4, 5, 6]);

        // Sensitive and non-extractable: returns AttributeSensitive
        assert!(matches!(
            read_attribute(&obj, CKA_PRIVATE_EXPONENT, true),
            Err(HsmError::AttributeSensitive)
        ));
        assert!(matches!(
            read_attribute(&obj, CKA_PRIME_1, true),
            Err(HsmError::AttributeSensitive)
        ));

        // Sensitive but extractable: also sensitive per Table 39 note 7
        obj.extractable = true;
        assert!(matches!(
            read_attribute(&obj, CKA_PRIVATE_EXPONENT, true),
            Err(HsmError::AttributeSensitive)
        ));

        // Non-sensitive and extractable: permitted
        obj.sensitive = false;
        obj.extractable = true;
        assert_eq!(
            read_attribute(&obj, CKA_PRIVATE_EXPONENT, true).unwrap(),
            Some(vec![1, 2, 3])
        );
        assert_eq!(
            read_attribute(&obj, CKA_PRIME_1, true).unwrap(),
            Some(vec![4, 5, 6])
        );
    }

    #[test]
    fn rsa_private_key_created_from_components_has_key_material() {
        use crate::crypto::drbg::DrbgRng;
        use rsa::traits::{PrivateKeyParts, PublicKeyParts};
        use rsa::RsaPrivateKey;

        let mut rng = DrbgRng::new().unwrap();
        let priv_key = RsaPrivateKey::new(&mut rng, 2048).unwrap();
        let n = priv_key.n().to_bytes_be();
        let e = priv_key.e().to_bytes_be();
        let d = priv_key.d().to_bytes_be();
        let p = priv_key.primes()[0].to_bytes_be();
        let q = priv_key.primes()[1].to_bytes_be();

        let store = ObjectStore::new();
        let template = vec![
            (CKA_CLASS, ck_ulong_to_bytes(CKO_PRIVATE_KEY)),
            (CKA_KEY_TYPE, ck_ulong_to_bytes(CKK_RSA)),
            (CKA_MODULUS, n.clone()),
            (CKA_PUBLIC_EXPONENT, e.clone()),
            (CKA_PRIVATE_EXPONENT, d),
            (CKA_PRIME_1, p),
            (CKA_PRIME_2, q),
            (CKA_DECRYPT, vec![CK_TRUE]),
        ];

        let handle = store
            .create_object(&template)
            .expect("create_object should succeed");
        let obj = store.get_object(handle).unwrap();
        let obj = obj.read();
        assert!(
            obj.key_material.is_some(),
            "key_material must be populated from components"
        );

        let priv_der = obj.key_material.as_ref().unwrap().as_bytes();
        let pt = b"RSA discrete component decryption test";
        let ct = crate::crypto::sign::rsa_oaep_encrypt(
            &n,
            &e,
            pt,
            crate::crypto::sign::OaepHash::Sha256,
            &[],
        )
        .unwrap();
        let recovered = crate::crypto::sign::rsa_oaep_decrypt(
            priv_der,
            &ct,
            crate::crypto::sign::OaepHash::Sha256,
            &[],
        )
        .unwrap();
        assert_eq!(&recovered, pt);
    }
}