// Copyright 2023 PingCAP, Inc.

mod metrics;

use std::{
    collections::BTreeMap,
    ops::Deref,
    sync::{Arc, Mutex, RwLock, Weak},
    time::Duration,
};

use bytes::Buf;
use cloud::{KmsProvider, error::KmsError};
use dashmap::{DashMap, DashSet, mapref::entry::Entry};
use derive_more::Deref;
use hmac::{Hmac, Mac, NewMac};
use kvenginepb::EncryptionMeta;
use openssl::{
    symm,
    symm::{Cipher, Crypter, Mode},
};
use rand::RngCore;
use serde_derive::{Deserialize, Serialize};
use strum::{Display, EnumString};
use thiserror::Error;
use tikv_util::{
    AsyncSafeRuntime, backoff::ExponentialBackoff, defer, metrics::CRITICAL_ERROR, time::Instant,
};
use tokio::sync::watch;
use zeroize::Zeroize;

use crate::metrics::*;

type Hmac256 = Hmac<sha2::Sha256>;

// Legacy encryption key format (used in serverless).
// - Exports a 64-byte encryption key.
// - At runtime, this is reduced to a 32-byte data key using HMAC-SHA256.
// - Encoded with data format = 0.
pub const KEY_TYPE_AES_256_CTR_LEGACY: u8 = 1;
const KEY_EXPORTED_SIZE_LEGACY: usize = 1 + 4 + 64;
const ENCRYPTION_DATA_FORMAT_LEGACY: u8 = 0;

// Next-gen key format.
// - Exports a 32-byte data encryption key directly, encrypted by the master
//   key.
// - Encoded with data format = 1.
pub const KEY_TYPE_AES_256_CTR: u8 = 2;
const KEY_EXPORTED_SIZE: usize = 1 + 32;
const ENCRYPTION_DATA_FORMAT: u8 = 1;
const MAX_DATA_KEY_ID: u32 = 0x00FF_FFFF;
const CMEK_HEALTH_CHECK_INTERVAL: Duration = Duration::from_secs(3 * 60 * 60);
const MASTER_KEY_RETRY_INITIAL_INTERVAL: Duration = Duration::from_secs(5);
const MASTER_KEY_RETRY_MAX_INTERVAL: Duration = Duration::from_secs(5 * 60);
const MASTER_KEY_DECRYPT_TIMEOUT: Duration = Duration::from_secs(15);

// Use zero IV for exported key wrapping compatibility because newer openssl
// versions do not allow None IV for AES-CTR.
// Ref: https://github.com/rust-openssl/rust-openssl/blob/db9c9e2f5db2ad7b45fd894e8d297ee15bfd0c7c/openssl/CHANGELOG.md?plain=1#L68
const KEY_WRAP_ZERO_IV: [u8; 16] = [0; 16];

// Next-gen key type that suggests no encryption is used.
pub const KEY_TYPE_UNENCRYPTED: u8 = 3;

pub type Result<T> = std::result::Result<T, EncryptionKeyError>;

#[derive(Debug, Error)]
pub enum EncryptionKeyError {
    #[error("Invalid empty master key")]
    InvalidMasterKey,

    #[error("Invalid exported key: {0}")]
    InvalidExportedKey(String),

    #[error("Invalid encryption meta: {0}")]
    InvalidEncryptionMeta(String),

    #[error("Requested data key {0} missing in keyspace {1}")]
    DataKeyMissing(u32, u32),

    #[error("Invalid encryption state: {0}")]
    InvalidEncryptionState(String),

    #[error("Encryption info unknown for keyspace {0}")]
    EncryptionInfoUnknownForKeyspace(u32),

    #[error("Master key unavailable: {0}")]
    MasterKeyUnavailable(String),

    #[error("Unsupported KMS vendor: {0}")]
    UnsupportedVendor(String),

    #[error(transparent)]
    Cloud(#[from] cloud::Error),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum EncryptionMethod {
    /// Legacy exported-key properties stored as `ENCRYPTION_KEY`.
    ///
    /// A process should use this mode consistently once legacy-encrypted
    /// shards are active.
    Legacy,
    /// Per-keyspace encryption metadata stored as `ENCRYPTION_META_KEY`.
    ///
    /// This mode is used for CMEK-aware keyspaces and should not be mixed with
    /// legacy mode within the same steady-state process.
    Metadata,
}

#[derive(Clone, Default, Debug, Serialize, Deserialize)]
#[serde(default)]
pub struct KeyspaceEncryptionConfig {
    pub enabled: bool,
    pub cmek_id: Option<String>,
    pub vendor: Option<String>,
    pub region: Option<String>,
    pub endpoint: Option<String>,
    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
    pub encryption_context: BTreeMap<String, String>,
    /// Test-only escape hatch for split-time mixed-mode admission.
    #[cfg(feature = "testexport")]
    #[serde(default)]
    pub bypass_mode_check: bool,
}

impl KeyspaceEncryptionConfig {
    /// KMS encryption context is provided to cloud KMS APIs as a map, but our
    /// control-plane wiring often represents it as a JSON object string.
    ///
    /// For operational safety, we cap the total JSON string length.
    pub const MAX_ENCRYPTION_CONTEXT_JSON_LEN: usize = 8192;

    pub fn validate_encryption_context(
        encryption_context: &BTreeMap<String, String>,
    ) -> std::result::Result<(), String> {
        if encryption_context.is_empty() {
            return Ok(());
        }

        let json = serde_json::to_string(encryption_context)
            .map_err(|err| format!("failed to serialize encryption_context as JSON: {}", err))?;
        if json.len() > Self::MAX_ENCRYPTION_CONTEXT_JSON_LEN {
            return Err(format!(
                "encryption_context JSON length must be <= {} chars, got {}",
                Self::MAX_ENCRYPTION_CONTEXT_JSON_LEN,
                json.len()
            ));
        }
        Ok(())
    }

    /// Builds an enabled keyspace encryption config with the provided KMS
    /// fields.
    pub fn enabled(
        cmek_id: Option<String>,
        vendor: Option<String>,
        region: Option<String>,
        endpoint: Option<String>,
        encryption_context: BTreeMap<String, String>,
    ) -> Self {
        Self {
            enabled: true,
            cmek_id,
            vendor,
            region,
            endpoint,
            encryption_context,
            #[cfg(feature = "testexport")]
            bypass_mode_check: false,
        }
    }

    /// Parses the JSON payload stored in PD keyspace config under the
    /// `encryption` key.
    pub fn from_pd_config_value(value: &str) -> serde_json::Result<Self> {
        use serde::de::Error as _;

        let config: Self = serde_json::from_str(value)?;
        Self::validate_encryption_context(&config.encryption_context)
            .map_err(serde_json::Error::custom)?;
        Ok(config)
    }

    /// Serializes the config into the JSON payload stored in PD keyspace config
    /// under the `encryption` key.
    pub fn to_pd_config_value(&self) -> serde_json::Result<String> {
        use serde::ser::Error as _;

        Self::validate_encryption_context(&self.encryption_context)
            .map_err(serde_json::Error::custom)?;
        serde_json::to_string(self)
    }

    /// Returns the active encryption intent carried by this keyspace config.
    ///
    /// `None` means no active encrypted mode is requested. This keeps legacy
    /// clusters without per-keyspace config backward compatible.
    pub fn intended_encryption_method(&self) -> Result<Option<EncryptionMethod>> {
        if !self.enabled {
            return Ok(None);
        }

        let has_vendor = self
            .vendor
            .as_ref()
            .is_some_and(|vendor| !vendor.is_empty());
        let has_cmek_id = self
            .cmek_id
            .as_ref()
            .is_some_and(|cmek_id| !cmek_id.is_empty());

        match (has_vendor, has_cmek_id) {
            (false, false) => Ok(Some(EncryptionMethod::Legacy)),
            (true, true) => Ok(Some(EncryptionMethod::Metadata)),
            _ => Err(EncryptionKeyError::InvalidEncryptionState(format!(
                "partial CMEK intent in keyspace encryption config: vendor_present={}, cmek_id_present={}",
                has_vendor, has_cmek_id
            ))),
        }
    }
}

impl EncryptionKeyError {
    pub fn is_retryable(&self) -> bool {
        matches!(
            self,
            EncryptionKeyError::Cloud(_) | EncryptionKeyError::MasterKeyUnavailable(_)
        )
    }

    fn reason_label(&self) -> &'static str {
        match self {
            EncryptionKeyError::InvalidMasterKey => "invalid_master_key",
            EncryptionKeyError::InvalidExportedKey(_) => "invalid_exported_key",
            EncryptionKeyError::InvalidEncryptionMeta(_) => "invalid_encryption_meta",
            EncryptionKeyError::DataKeyMissing(..) => "data_key_missing",
            EncryptionKeyError::InvalidEncryptionState(_) => "invalid_encryption_state",
            EncryptionKeyError::EncryptionInfoUnknownForKeyspace(_) => {
                "encryption_info_unknown_for_keyspace"
            }
            EncryptionKeyError::MasterKeyUnavailable(_) => "master_key_unavailable",
            EncryptionKeyError::UnsupportedVendor(_) => "unsupported_vendor",
            EncryptionKeyError::Cloud(_) => "cloud",
        }
    }
}

#[derive(Clone)]
pub struct EncryptionKey {
    pub core: Arc<EncryptionKeyCore>,
}

/// Runtime state for a shard/peer encryption key.
///
/// `Ready` carries a key that can be used immediately. `Pending` retains the
/// metadata that defines its key but does not resolve it until encryption or
/// decryption is required. `Invalid` records a metadata load error so callers
/// can delay the failure until the key is actually used. `Disabled` means
/// encryption is disabled.
#[derive(Clone, Default)]
pub enum EncryptionState {
    Ready(EncryptionKey),
    Pending {
        keyspace_id: u32,
        data_key_id: u32,
        encryption_meta: Arc<EncryptionMeta>,
        key_manager: Arc<EncryptionKeyManager>,
    },
    Invalid {
        keyspace_id: Option<u32>,
        err: String,
    },
    #[default]
    Disabled,
}

/// Converts an optional ready encryption key into an encryption state.
pub trait ToEncryptionState {
    fn to_encryption_state(&self) -> EncryptionState;
}

impl ToEncryptionState for Option<EncryptionKey> {
    fn to_encryption_state(&self) -> EncryptionState {
        self.clone()
            .map_or(EncryptionState::Disabled, EncryptionState::Ready)
    }
}

impl ToEncryptionState for Option<&EncryptionKey> {
    fn to_encryption_state(&self) -> EncryptionState {
        self.cloned()
            .map_or(EncryptionState::Disabled, EncryptionState::Ready)
    }
}

impl std::fmt::Debug for EncryptionState {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Ready(key) => f.debug_tuple("ready_encryption_key").field(key).finish(),
            Self::Pending {
                keyspace_id,
                data_key_id,
                ..
            } => f
                .debug_struct("pending_encryption_key")
                .field("keyspace_id", keyspace_id)
                .field("data_key_id", data_key_id)
                .finish(),
            Self::Invalid { keyspace_id, err } => f
                .debug_struct("invalid_encryption_state")
                .field("keyspace_id", keyspace_id)
                .field("err", err)
                .finish(),
            Self::Disabled => f.write_str("encryption_disabled"),
        }
    }
}

impl EncryptionState {
    fn pending_from_metadata(
        key_manager: Arc<EncryptionKeyManager>,
        keyspace_id: u32,
        data_key_id: u32,
        encryption_meta: Arc<EncryptionMeta>,
    ) -> Self {
        Self::Pending {
            keyspace_id,
            data_key_id,
            encryption_meta,
            key_manager,
        }
    }

    /// Records an encryption metadata error as a runtime state.
    pub fn from_error(keyspace_id: Option<u32>, err: EncryptionKeyError) -> Self {
        Self::Invalid {
            keyspace_id,
            err: err.to_string(),
        }
    }

    /// Returns the keyspace ID carried by this state, if any.
    pub fn keyspace_id(&self) -> Option<u32> {
        match self {
            Self::Ready(key) => key.keyspace_id,
            Self::Pending { keyspace_id, .. } => Some(*keyspace_id),
            Self::Invalid { keyspace_id, .. } => *keyspace_id,
            Self::Disabled => None,
        }
    }

    /// Returns the data key ID carried by this state, if any.
    pub fn data_key_id(&self) -> Option<u32> {
        match self {
            Self::Ready(key) => key.get_key_id().ok(),
            Self::Pending { data_key_id, .. } => Some(*data_key_id),
            Self::Invalid { .. } => None,
            Self::Disabled => None,
        }
    }

    /// Returns whether this state represents an encrypted key.
    pub fn is_encrypted(&self) -> bool {
        match self {
            Self::Ready(key) => !key.is_noop(),
            Self::Pending { .. } => true,
            Self::Invalid { .. } => true,
            Self::Disabled => false,
        }
    }

    /// Returns the raft-log encryption header carried by this state, if any.
    pub fn encryption_header(&self) -> Option<u32> {
        match self {
            Self::Ready(key) => Some(key.encryption_header()),
            Self::Pending { data_key_id, .. } => {
                Some(((ENCRYPTION_DATA_FORMAT as u32) << 24) | data_key_id)
            }
            Self::Invalid { .. } => None,
            Self::Disabled => None,
        }
    }

    /// Resolves this state into an immediately usable encryption key.
    pub fn resolve_ready_key(&self) -> Result<Option<EncryptionKey>> {
        let result = match self {
            Self::Ready(key) => return Ok(Some(key.clone())),
            Self::Pending {
                keyspace_id,
                encryption_meta,
                key_manager,
                ..
            } => key_manager
                .get_data_key_by_header(
                    *keyspace_id,
                    self.encryption_header().unwrap(),
                    Some(encryption_meta),
                )
                .map(|key| key.with_key_manager(key_manager.clone())),
            Self::Invalid { err, .. } => {
                Err(EncryptionKeyError::InvalidEncryptionState(err.clone()))
            }
            Self::Disabled => return Ok(None),
        };
        EncryptionKey::observe_key_result(&result);
        result.map(Some)
    }

    /// Resolves this state and caches the ready key in-place when possible.
    pub fn resolve_ready_key_and_cache(&mut self) -> Result<Option<EncryptionKey>> {
        match self {
            Self::Ready(key) => Ok(Some(key.clone())),
            Self::Pending { .. } => {
                let key = self.resolve_ready_key()?.unwrap();
                *self = Self::Ready(key.clone());
                Ok(Some(key))
            }
            Self::Invalid { err, .. } => {
                Err(EncryptionKeyError::InvalidEncryptionState(err.clone()))
            }
            Self::Disabled => Ok(None),
        }
    }

    /// Returns a ready key matching `encryption_header`.
    pub fn switch_if_header_mismatch(&self, encryption_header: u32) -> Result<EncryptionKey> {
        match self {
            Self::Ready(key) => key.switch_if_header_mismatch(encryption_header),
            Self::Pending {
                keyspace_id,
                encryption_meta,
                key_manager,
                ..
            } => {
                let result = key_manager
                    .get_data_key_by_header(*keyspace_id, encryption_header, Some(encryption_meta))
                    .map(|key| key.with_key_manager(key_manager.clone()));
                EncryptionKey::observe_key_result(&result);
                result
            }
            Self::Invalid { err, .. } => {
                Err(EncryptionKeyError::InvalidEncryptionState(err.clone()))
            }
            Self::Disabled => Err(EncryptionKeyError::InvalidEncryptionState(
                "encryption key is missing".into(),
            )),
        }
    }

    /// Returns a ready key matching `encryption_header` and caches it when it
    /// resolves this state's current pending data key.
    pub fn switch_if_header_mismatch_and_cache(
        &mut self,
        encryption_header: u32,
    ) -> Result<EncryptionKey> {
        let should_cache_current = match self {
            Self::Pending { data_key_id, .. } => {
                (((ENCRYPTION_DATA_FORMAT as u32) << 24) | *data_key_id) == encryption_header
            }
            _ => false,
        };
        let key = self.switch_if_header_mismatch(encryption_header)?;
        if should_cache_current {
            *self = Self::Ready(key.clone());
        }
        Ok(key)
    }
}

impl std::fmt::Debug for EncryptionKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("encryption_key")
            .field(
                "export",
                &self.export().iter().fold(String::new(), |mut acc, byte| {
                    acc.push_str(format!("{:02X}", byte).as_str());
                    acc
                }),
            )
            .finish()
    }
}

impl EncryptionKey {
    pub fn new(cipher_text: Vec<u8>, plain_text: Vec<u8>, key_type: u8) -> Self {
        let plain_text = match key_type {
            KEY_TYPE_AES_256_CTR_LEGACY => {
                EncryptionKeyCore::transform_legacy_encryption_key(&plain_text)
            }
            KEY_TYPE_AES_256_CTR => plain_text,
            _ => unimplemented!(),
        };

        Self {
            core: Arc::new(EncryptionKeyCore {
                keyspace_id: None,
                data_key_id: None,
                key_type,
                cipher_text,
                plain_text,
                key_manager: None,
            }),
        }
    }

    pub fn new_unencrypted_key() -> Self {
        Self {
            core: Arc::new(EncryptionKeyCore {
                keyspace_id: None,
                data_key_id: None,
                key_type: KEY_TYPE_UNENCRYPTED,
                cipher_text: vec![],
                plain_text: vec![],
                key_manager: None,
            }),
        }
    }

    pub fn with_keyspace_id(&self, keyspace_id: u32) -> Self {
        EncryptionKey {
            core: Arc::new(self.core.with_keyspace_id(keyspace_id)),
        }
    }

    pub fn with_key_manager(&self, m: Arc<EncryptionKeyManager>) -> Self {
        EncryptionKey {
            core: Arc::new(self.core.with_key_manager(m)),
        }
    }

    /// Sets the logical data-key ID for V2 keys.
    ///
    /// `EncryptionMeta` owns this ID. It must stay stable across master-key
    /// rewraps even when the wrapped ciphertext changes.
    ///
    /// Returns an error if the key type cannot carry a logical ID or the ID
    /// does not fit in the 24-bit encryption header space.
    pub fn with_key_id(&self, key_id: u32) -> Result<Self> {
        self.core.with_key_id(key_id).map(|core| EncryptionKey {
            core: Arc::new(core),
        })
    }

    /// Returns a ready encryption key that matches the given encryption header.
    pub fn switch_if_header_mismatch(&self, encryption_header: u32) -> Result<Self> {
        if self.encryption_header() == encryption_header {
            return Ok(self.clone());
        }
        ENCRYPTION_KEY_SWITCH_COUNTER.inc();
        let result = (|| {
            let key_manager = self.key_manager.clone().ok_or_else(|| {
                EncryptionKeyError::InvalidEncryptionState(
                    "encryption key manager is not attached to the encryption key".into(),
                )
            })?;
            let keyspace_id = self.keyspace_id.ok_or_else(|| {
                EncryptionKeyError::InvalidEncryptionState(
                    "keyspace id is not attached to the encryption key".into(),
                )
            })?;
            key_manager
                .get_data_key_by_header(keyspace_id, encryption_header, None)
                .map(|key| key.with_key_manager(key_manager))
        })();
        Self::observe_key_result(&result);
        result
    }

    fn observe_key_result(result: &Result<Self>) {
        if let Err(err) = &result {
            ENCRYPTION_KEY_UNAVAILABLE_COUNT
                .with_label_values(&[err.reason_label()])
                .inc();
        }
    }
}

/// Generates a new 24-bit logical data-key ID for V2 `EncryptionMeta`.
///
/// Callers should persist the returned value in `EncryptionMeta` and treat it
/// as the source of truth for future lookups.
pub fn generate_data_key_id() -> u32 {
    rand::thread_rng().next_u32() & MAX_DATA_KEY_ID
}

impl Deref for EncryptionKey {
    type Target = EncryptionKeyCore;
    fn deref(&self) -> &Self::Target {
        &self.core
    }
}

pub struct EncryptionKeyCore {
    pub key_type: u8,
    // Legacy keys do not persist a logical data-key ID. V2 keys must bind one
    // from EncryptionMeta before writing an encryption header.
    pub data_key_id: Option<u32>,
    pub keyspace_id: Option<u32>,
    pub cipher_text: Vec<u8>,
    pub plain_text: Vec<u8>,
    pub key_manager: Option<Arc<EncryptionKeyManager>>,
}

impl EncryptionKeyCore {
    pub fn encrypt(&self, data: &[u8], iv_high: u64, iv_low: u32, buf: &mut Vec<u8>) {
        if self.key_type == KEY_TYPE_UNENCRYPTED {
            buf.extend_from_slice(data);
            return;
        }
        assert!(
            !self.plain_text.is_empty(),
            "encryption key is missing plaintext while encrypting"
        );
        let mut iv = [0u8; 16];
        iv[0..8].copy_from_slice(&iv_high.to_be_bytes());
        iv[8..12].copy_from_slice(&iv_low.to_be_bytes());
        let cipher = Cipher::aes_256_ctr();
        let mut crypter = Crypter::new(cipher, Mode::Encrypt, &self.plain_text, Some(&iv)).unwrap();
        let origin_len = buf.len();
        let additional_len = data.len() + cipher.block_size();
        buf.resize(origin_len + additional_len, 0);
        let dst = &mut buf[origin_len..];
        let count = crypter.update(data, dst).unwrap();
        let rest = crypter.finalize(&mut dst[count..]).unwrap();
        buf.truncate(origin_len + count + rest);
    }

    pub fn decrypt(&self, ciphertext: &[u8], iv_high: u64, iv_low: u32, buf: &mut Vec<u8>) {
        if self.key_type == KEY_TYPE_UNENCRYPTED {
            buf.extend_from_slice(ciphertext);
            return;
        }
        assert!(
            !self.plain_text.is_empty(),
            "encryption key is missing plaintext while decrypting"
        );
        let mut iv = [0u8; 16];
        iv[0..8].copy_from_slice(&iv_high.to_be_bytes());
        iv[8..12].copy_from_slice(&iv_low.to_be_bytes());
        let cipher = Cipher::aes_256_ctr();
        let mut crypter = Crypter::new(cipher, Mode::Decrypt, &self.plain_text, Some(&iv)).unwrap();
        let origin_len = buf.len();
        let additional_len = ciphertext.len() + cipher.block_size();
        buf.resize(origin_len + additional_len, 0);
        let dst = &mut buf[origin_len..];
        let count = crypter.update(ciphertext, dst).unwrap();
        let rest = crypter.finalize(&mut dst[count..]).unwrap();
        buf.truncate(origin_len + count + rest);
    }

    pub fn export(&self) -> Vec<u8> {
        match self.key_type {
            KEY_TYPE_AES_256_CTR_LEGACY => {
                let mut data = Vec::with_capacity(KEY_EXPORTED_SIZE_LEGACY);
                let current_ver: u32 = 0;
                data.push(KEY_TYPE_AES_256_CTR_LEGACY);
                data.extend_from_slice(&current_ver.to_be_bytes());
                data.extend_from_slice(&self.cipher_text);
                data
            }
            KEY_TYPE_AES_256_CTR => {
                let mut data = Vec::with_capacity(KEY_EXPORTED_SIZE);
                data.push(KEY_TYPE_AES_256_CTR);
                data.extend_from_slice(&self.cipher_text);
                data
            }
            KEY_TYPE_UNENCRYPTED => vec![KEY_TYPE_UNENCRYPTED],
            _ => unreachable!(),
        }
    }

    /// Returns true if this key performs no encryption and passes data
    /// through unchanged.
    pub fn is_noop(&self) -> bool {
        self.key_type == KEY_TYPE_UNENCRYPTED
    }

    // Transforms a legacy 64-byte encryption key into a 32-byte data key using
    // HMAC-SHA256 with a fixed version.
    //
    // V1 always uses version 0 since key rotation was not supported.
    fn transform_legacy_encryption_key(plain_text: &[u8]) -> Vec<u8> {
        let key_version: u32 = 0;
        let mut hmac = Hmac256::new_varkey(plain_text).unwrap();
        hmac.update(&key_version.to_be_bytes());
        hmac.finalize().into_bytes().to_vec()
    }

    pub fn get_key_id(&self) -> Result<u32> {
        match self.key_type {
            KEY_TYPE_AES_256_CTR_LEGACY => Ok(0),
            KEY_TYPE_AES_256_CTR => self.data_key_id.ok_or_else(|| {
                EncryptionKeyError::InvalidEncryptionState(
                    "data key id is not attached to the encryption key".into(),
                )
            }),
            KEY_TYPE_UNENCRYPTED => Ok(0),
            _ => unreachable!(),
        }
    }

    pub fn get_type(&self) -> u8 {
        self.key_type
    }

    pub fn with_keyspace_id(&self, keyspace_id: u32) -> Self {
        Self {
            key_type: self.key_type,
            data_key_id: self.data_key_id,
            keyspace_id: Some(keyspace_id),
            cipher_text: self.cipher_text.clone(),
            plain_text: self.plain_text.clone(),
            key_manager: self.key_manager.clone(),
        }
    }

    pub fn with_key_manager(&self, key_manager: Arc<EncryptionKeyManager>) -> Self {
        Self {
            key_type: self.key_type,
            data_key_id: self.data_key_id,
            keyspace_id: self.keyspace_id,
            cipher_text: self.cipher_text.clone(),
            plain_text: self.plain_text.clone(),
            key_manager: Some(key_manager),
        }
    }

    pub fn with_key_id(&self, key_id: u32) -> Result<Self> {
        if self.key_type != KEY_TYPE_AES_256_CTR {
            return Err(EncryptionKeyError::InvalidEncryptionState(format!(
                "encryption key type {} cannot carry a logical data key id",
                self.key_type
            )));
        }
        if key_id > MAX_DATA_KEY_ID {
            return Err(EncryptionKeyError::InvalidEncryptionState(format!(
                "data key id {} exceeds the 24-bit encryption header space",
                key_id
            )));
        }
        Ok(Self {
            key_type: self.key_type,
            data_key_id: Some(key_id),
            keyspace_id: self.keyspace_id,
            cipher_text: self.cipher_text.clone(),
            plain_text: self.plain_text.clone(),
            key_manager: self.key_manager.clone(),
        })
    }

    /// Returns the 4-byte encryption header that identifies the data key within
    /// the keyspace. The first byte is the data format version; the remaining 3
    /// bytes encode the data key ID.
    ///
    /// This header is written alongside the encrypted data and is required
    /// later to locate and use the correct data key during decryption.
    pub fn encryption_header(&self) -> u32 {
        match self.key_type {
            KEY_TYPE_AES_256_CTR_LEGACY => ENCRYPTION_DATA_FORMAT_LEGACY as u32,
            KEY_TYPE_AES_256_CTR => {
                let data_key_id = self.data_key_id.expect(
                    "v2 encryption key missing data_key_id when encoding encryption header",
                );
                ((ENCRYPTION_DATA_FORMAT as u32) << 24) | data_key_id
            }
            KEY_TYPE_UNENCRYPTED => 0,
            _ => unreachable!(),
        }
    }

    pub fn encryption_block_size(&self) -> usize {
        Cipher::aes_256_ctr().block_size()
    }
}

impl Drop for EncryptionKeyCore {
    fn drop(&mut self) {
        // Zeroize sensitive fields.
        self.plain_text.zeroize();
    }
}

#[derive(Clone, Copy, Debug, EnumString, Display, PartialEq, Eq)]
#[strum(serialize_all = "kebab_case", ascii_case_insensitive)]
pub enum KmsVendor {
    Aws,
    Aliyun,
    Tencentcloud,
    Volcengine,
    Test,
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Eq, Hash, Debug, Default)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct MasterKeyConfig {
    // Encrypted master key, base64 encoded.
    pub cipher_text: String,
    // KMS info
    pub key_id: String,
    pub region: String,
    pub vendor: String,
    pub endpoint: String,
    #[serde(skip_serializing_if = "BTreeMap::is_empty")]
    pub encryption_context: BTreeMap<String, String>,
}

impl MasterKeyConfig {
    pub fn override_from_env(&mut self) {
        Self::env_or_default("CSE_MASTER_KEY_ID", &mut self.key_id);
        Self::env_or_default("CSE_MASTER_KEY_CIPHER_TEXT", &mut self.cipher_text);
        Self::env_or_default("CSE_MASTER_KEY_VENDOR", &mut self.vendor);
        Self::env_or_default("CSE_MASTER_KEY_ENDPOINT", &mut self.endpoint);
        Self::env_or_default("AWS_REGION", &mut self.region);
    }

    fn env_or_default(name: &str, val: &mut String) {
        if let Ok(v) = std::env::var(name) {
            *val = v;
        }
    }

    fn parsed_vendor(&self) -> Option<KmsVendor> {
        self.vendor.parse::<KmsVendor>().ok()
    }

    fn cloud_kms_config(&self) -> cloud::Result<cloud::kms::Config> {
        cloud::kms::Config::from_parts(
            self.key_id.clone(),
            self.region.clone(),
            self.endpoint.clone(),
            self.vendor.clone(),
            self.encryption_context.clone(),
        )
    }

    /// Generates a new master key and its encrypted form using the configured
    /// KMS vendor.
    pub async fn generate_new_master_key(&self) -> Result<(MasterKey, cloud::EncryptedKey)> {
        let start = Instant::now_coarse();
        defer! {
            MASTER_KEY_OPS_DURATION
                .generate
                .observe(start.saturating_elapsed().as_secs_f64())
        };
        let result = async {
            let vendor = self.parsed_vendor();
            match vendor {
                Some(KmsVendor::Aws) => {
                    let kms = aws::AwsKms::new(self.cloud_kms_config()?)?;
                    let key_pair = kms.generate_data_key().await?;

                    let master_key = MasterKey::new_with_kms_vendor(
                        key_pair.plaintext.as_slice(),
                        KmsVendor::Aws,
                    );
                    let master_key_ciphertext =
                        cloud::EncryptedKey::new(key_pair.encrypted.into_inner())?;
                    Ok((master_key, master_key_ciphertext))
                }
                Some(KmsVendor::Aliyun) => {
                    let key_pair = aliyun::generate_data_key(
                        &self.key_id,
                        &self.region,
                        &self.endpoint,
                        &self.encryption_context,
                    )
                    .await
                    .map_err(|e| cloud::Error::KmsError(KmsError::Other(Box::new(e))))?;
                    let master_key = MasterKey::new_with_kms_vendor(
                        key_pair.plaintext.as_slice(),
                        KmsVendor::Aliyun,
                    );
                    let master_key_ciphertext =
                        cloud::EncryptedKey::new(key_pair.encrypted.into_inner())?;
                    Ok((master_key, master_key_ciphertext))
                }
                Some(KmsVendor::Tencentcloud) | Some(KmsVendor::Volcengine) => {
                    Err(EncryptionKeyError::UnsupportedVendor(self.vendor.clone()))
                }
                Some(KmsVendor::Test) => {
                    let mut master_key = self.key_id.as_bytes().to_vec();
                    master_key.resize(32, 11);
                    let master_key_ciphertext = cloud::EncryptedKey::new(master_key.clone())?; // no encryption
                    Ok((
                        MasterKey::new_with_kms_vendor(&master_key, KmsVendor::Test),
                        master_key_ciphertext,
                    ))
                }
                None => Err(EncryptionKeyError::UnsupportedVendor(self.vendor.clone())),
            }
        }
        .await;
        if let Err(err) = &result {
            observe_master_key_failure("generate", &self.key_id, &self.vendor, err);
        }
        result
    }

    /// Decrypts the encrypted master key using the configured KMS vendor.
    ///
    /// `is_legacy` is only for the process-level legacy master-key config,
    /// where an empty vendor means encryption is disabled. CMEK metadata must
    /// pass `false` so missing or unknown vendors fail closed.
    pub async fn decrypt(&self, is_legacy: bool) -> Result<MasterKey> {
        let start = Instant::now_coarse();

        let decrypt = async {
            let vendor = self.parsed_vendor();
            match vendor {
                Some(KmsVendor::Aws) => {
                    let encrypted_key_bytes = base64::decode(&self.cipher_text).map_err(|e| {
                        EncryptionKeyError::InvalidEncryptionMeta(format!(
                            "invalid master key ciphertext: {}",
                            e
                        ))
                    })?;
                    let encrypted_key =
                        cloud::EncryptedKey::new(encrypted_key_bytes).map_err(|e| {
                            EncryptionKeyError::InvalidEncryptionMeta(format!(
                                "invalid master key ciphertext: {}",
                                e
                            ))
                        })?;

                    let aws_kms = aws::AwsKms::new(self.cloud_kms_config()?)?;
                    let master_key_plain_text = aws_kms.decrypt_data_key(&encrypted_key).await?;
                    Ok(MasterKey::new_with_kms_vendor(
                        &master_key_plain_text,
                        KmsVendor::Aws,
                    ))
                }
                Some(KmsVendor::Aliyun) => {
                    let master_key_plain_text = aliyun::decrypt_master_key(
                        &self.cipher_text,
                        &self.region,
                        &self.endpoint,
                        &self.encryption_context,
                    )
                    .await
                    .map_err(|e| cloud::Error::KmsError(KmsError::Other(Box::new(e))))?;
                    Ok(MasterKey::new_with_kms_vendor(
                        &master_key_plain_text,
                        KmsVendor::Aliyun,
                    ))
                }
                Some(KmsVendor::Tencentcloud) => {
                    let master_key_plain_text = tencentcloud::decrypt_master_key(
                        &self.cipher_text,
                        &self.region,
                        &self.endpoint,
                    )
                    .await
                    .map_err(|e| cloud::Error::KmsError(KmsError::Other(Box::new(e))))?;
                    Ok(MasterKey::new_with_kms_vendor(
                        &master_key_plain_text,
                        KmsVendor::Tencentcloud,
                    ))
                }
                Some(KmsVendor::Volcengine) => {
                    let master_key_plain_text = volcengine::decrypt_master_key(
                        &self.cipher_text,
                        &self.region,
                        &self.endpoint,
                        &self.encryption_context,
                    )
                    .await
                    .map_err(|e| cloud::Error::KmsError(KmsError::Other(Box::new(e))))?;
                    Ok(MasterKey::new_with_kms_vendor(
                        &master_key_plain_text,
                        KmsVendor::Volcengine,
                    ))
                }
                Some(KmsVendor::Test) => {
                    let mut master_key = self.key_id.as_bytes().to_vec();
                    master_key.resize(32, 11);
                    Ok(MasterKey::new_with_kms_vendor(&master_key, KmsVendor::Test))
                }
                None if is_legacy && self.vendor.is_empty() => Ok(MasterKey::new(&[])),
                None => Err(EncryptionKeyError::UnsupportedVendor(self.vendor.clone())),
            }
        };
        let result = match tokio::time::timeout(MASTER_KEY_DECRYPT_TIMEOUT, decrypt).await {
            Ok(result) => result,
            Err(_) => Err(EncryptionKeyError::MasterKeyUnavailable(format!(
                "master key decrypt timed out after {}s",
                MASTER_KEY_DECRYPT_TIMEOUT.as_secs()
            ))),
        };
        match &result {
            Ok(_) => {
                let elapsed = start.saturating_elapsed().as_secs_f64();
                log::info!(
                    "[CMEK] decrypt master key took {}s, cmek_id={}, vendor={}",
                    elapsed,
                    self.key_id,
                    self.vendor
                );
                MASTER_KEY_OPS_DURATION.decrypt.observe(elapsed);
            }
            Err(err) => observe_master_key_failure("decrypt", &self.key_id, &self.vendor, err),
        }
        result
    }
}

pub fn master_key_encryption_context(
    master_key: &kvenginepb::MasterKey,
) -> BTreeMap<String, String> {
    master_key
        .get_encryption_context()
        .iter()
        .map(|(key, value)| (key.clone(), value.clone()))
        .collect()
}

pub fn set_master_key_encryption_context(
    master_key: &mut kvenginepb::MasterKey,
    encryption_context: &BTreeMap<String, String>,
) {
    master_key.set_encryption_context(
        encryption_context
            .iter()
            .map(|(key, value)| (key.clone(), value.clone()))
            .collect(),
    );
}

fn observe_master_key_failure(
    operation: &'static str,
    cmek_id: &str,
    vendor: &str,
    err: &EncryptionKeyError,
) {
    let critical_type = match operation {
        "generate" => "cmek_master_key_generate",
        "decrypt" => "cmek_master_key_decrypt",
        _ => "cmek_master_key",
    };
    CRITICAL_ERROR.with_label_values(&[critical_type]).inc();
    log::warn!(
        "[CMEK] master key operation failed: type={}, cmek_id={}, vendor={}, retryable={}",
        operation,
        cmek_id,
        vendor,
        err.is_retryable(),
    );
}

#[derive(Clone)]
pub struct MasterKey {
    core: Arc<MasterKeyCore>,
}

impl MasterKey {
    pub fn new(plain_text: &[u8]) -> Self {
        Self::new_with_real_kms_vendor(plain_text, false)
    }

    pub fn new_with_kms_vendor(plain_text: &[u8], vendor: KmsVendor) -> Self {
        Self::new_with_real_kms_vendor(plain_text, vendor != KmsVendor::Test)
    }

    fn new_with_real_kms_vendor(plain_text: &[u8], uses_real_kms_vendor: bool) -> Self {
        Self {
            core: Arc::new(MasterKeyCore {
                master_key: plain_text.to_vec(),
                uses_real_kms_vendor,
            }),
        }
    }
}

impl Deref for MasterKey {
    type Target = MasterKeyCore;

    fn deref(&self) -> &Self::Target {
        &self.core
    }
}

// ExportedMasterKey is a dedicated type for exported master key,
// which is used to avoid accidentally printing the master key.
#[derive(Deref, PartialEq)]
pub struct ExportedMasterKey(Vec<u8>);

impl ExportedMasterKey {
    pub fn new(master_key: Vec<u8>) -> Self {
        Self(master_key)
    }
}

impl std::fmt::Debug for ExportedMasterKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_tuple("ExportedMasterKey")
            .field(&"REDACTED".to_string())
            .finish()
    }
}

pub struct MasterKeyCore {
    master_key: Vec<u8>,
    uses_real_kms_vendor: bool,
}

impl MasterKeyCore {
    pub fn is_valid(&self) -> bool {
        !self.master_key.is_empty()
    }

    pub fn uses_real_kms_vendor(&self) -> bool {
        self.uses_real_kms_vendor
    }

    pub fn generate_encryption_key(&self) -> EncryptionKey {
        self.generate_encryption_key_of_type(64, KEY_TYPE_AES_256_CTR_LEGACY)
    }

    pub fn generate_encryption_key_v2(&self) -> EncryptionKey {
        self.generate_encryption_key_of_type(32, KEY_TYPE_AES_256_CTR)
    }

    fn generate_encryption_key_of_type(
        &self,
        plain_text_len: usize,
        key_type: u8,
    ) -> EncryptionKey {
        assert!(self.is_valid());
        let mut rng = rand::thread_rng();
        let mut plain_text = vec![0u8; plain_text_len];
        rng.fill_bytes(&mut plain_text[..]);
        let cipher_text = symm::encrypt(
            Cipher::aes_256_ctr(),
            &self.master_key,
            Some(&KEY_WRAP_ZERO_IV),
            &plain_text,
        )
        .unwrap();
        EncryptionKey::new(cipher_text, plain_text, key_type)
    }

    pub fn decrypt_encryption_key(&self, mut exported: &[u8]) -> Result<EncryptionKey> {
        if !self.is_valid() {
            return Err(EncryptionKeyError::InvalidMasterKey);
        }
        let key_len = exported.len();
        let key_type = exported.first().copied().ok_or_else(|| {
            EncryptionKeyError::InvalidExportedKey("missing key type byte".into())
        })?;
        exported.advance(1);
        let expected_len = match key_type {
            KEY_TYPE_AES_256_CTR_LEGACY => KEY_EXPORTED_SIZE_LEGACY,
            KEY_TYPE_AES_256_CTR => KEY_EXPORTED_SIZE,
            KEY_TYPE_UNENCRYPTED => 1,
            _ => {
                return Err(EncryptionKeyError::InvalidExportedKey(format!(
                    "unsupported key type {}",
                    key_type
                )));
            }
        };
        if expected_len != key_len {
            return Err(EncryptionKeyError::InvalidExportedKey(format!(
                "invalid key length {} for key type {}",
                key_len, key_type
            )));
        }

        if key_type == KEY_TYPE_UNENCRYPTED {
            return Ok(EncryptionKey::new_unencrypted_key());
        }

        if key_type == KEY_TYPE_AES_256_CTR_LEGACY {
            let key_ver = exported.get_u32();
            if key_ver != 0 {
                return Err(EncryptionKeyError::InvalidExportedKey(format!(
                    "unsupported legacy key version {}",
                    key_ver
                )));
            }
        }
        let cipher_text = exported.to_vec();
        let plain_text = symm::decrypt(
            Cipher::aes_256_ctr(),
            &self.master_key,
            Some(&KEY_WRAP_ZERO_IV),
            &cipher_text,
        )
        .map_err(|e| EncryptionKeyError::InvalidExportedKey(format!("decryption failed: {}", e)))?;
        Ok(EncryptionKey::new(cipher_text, plain_text, key_type))
    }

    pub fn export(&self) -> ExportedMasterKey {
        ExportedMasterKey::new(self.master_key.clone())
    }
}

impl Drop for MasterKeyCore {
    fn drop(&mut self) {
        // Zeroize sensitive fields.
        self.master_key.zeroize();
    }
}

use protobuf::Message;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct EncryptionKeyRef {
    keyspace_id: u32,
    data_key_id: u32,
    unencrypted: bool,
}

/// Extracts a ready current data key from serialized `EncryptionMeta` bytes.
///
/// Internally registers the `EncryptionMeta` if needed. KMS must be available
/// because the returned key is immediately usable for encryption/decryption.
pub fn current_data_key_from_encryption_meta_bytes(
    m: Arc<EncryptionKeyManager>,
    encryption_meta_bytes: &[u8],
    keyspace_id: Option<u32>,
) -> Result<EncryptionKey> {
    let (entry, incoming, data_key_id) =
        m.register_encryption_meta_if_needed(encryption_meta_bytes, keyspace_id)?;
    let result = entry.get_data_key_from_meta(&incoming, data_key_id);
    EncryptionKey::observe_key_result(&result);
    result.map(|key| key.with_key_manager(m))
}

/// Registers serialized `EncryptionMeta` and returns state for its current data
/// key without contacting KMS.
///
/// Invalid metadata is recorded as `EncryptionState::Invalid` so metadata-load
/// paths can finish and report the failure when the key is actually used.
pub fn encryption_state_from_meta_bytes(
    m: Arc<EncryptionKeyManager>,
    encryption_meta_bytes: &[u8],
    keyspace_id: Option<u32>,
) -> EncryptionState {
    let result = m.register_encryption_meta_if_needed(encryption_meta_bytes, keyspace_id);
    let (_, incoming, data_key_id) = match result {
        Ok(result) => result,
        Err(err) => {
            let keyspace_id = encryption_meta_keyspace_id(encryption_meta_bytes, keyspace_id);
            return EncryptionState::from_error(keyspace_id, err);
        }
    };
    let target_keyspace_id = keyspace_id
        .filter(|keyspace_id| *keyspace_id > 0)
        .unwrap_or_else(|| incoming.get_keyspace_id());
    EncryptionState::pending_from_metadata(m, target_keyspace_id, data_key_id, incoming)
}

fn encryption_meta_keyspace_id(
    encryption_meta_bytes: &[u8],
    keyspace_id: Option<u32>,
) -> Option<u32> {
    if keyspace_id.is_some_and(|keyspace_id| keyspace_id > 0) {
        return keyspace_id;
    }

    let mut incoming = kvenginepb::EncryptionMeta::default();
    incoming.merge_from_bytes(encryption_meta_bytes).ok()?;
    let incoming_keyspace_id = incoming.get_keyspace_id();
    (incoming_keyspace_id > 0).then_some(incoming_keyspace_id)
}

/// A singleton that manages encryption keys across all keyspaces.
pub struct EncryptionKeyManager {
    // Tokio runtime for handling KMS operations (master key decryption).
    runtime: Arc<AsyncSafeRuntime>,

    // Shutdown signal for the periodic CMEK health check task.
    cmek_health_check_shutdown_tx: watch::Sender<bool>,

    // Best-effort process signal used only to gate enabling encryption for new
    // keyspaces. Its intent comes from loaded shard properties or, at startup,
    // from a real-vendor global master key. Existing shard reads still follow
    // shard properties, not this latch.
    mode: Mutex<Option<EncryptionMethod>>,

    // shard_id -> (keyspace_id, data_key_id)
    shard_to_current_key: DashMap<u64, EncryptionKeyRef>,

    // keyspace_id -> number of active shards on this store
    keyspace_active_shard_counts: Arc<DashMap<u32, usize>>,

    // keyspace_id -> keyspace encryption entry
    keyspace_to_encryption_entry: Arc<DashMap<u32, Arc<KeyspaceEncryptionEntry>>>,

    // Keyspaces that are known to have no encryption.
    unencrypted_keyspaces: DashSet<u32>,
}

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

impl Drop for EncryptionKeyManager {
    fn drop(&mut self) {
        let _ = self.cmek_health_check_shutdown_tx.send(true);
        // Release cached keyspace entries before the owning runtime field is
        // dropped, so the runtime can usually become uniquely owned and enter
        // its shutdown path immediately when the manager is torn down.
        self.keyspace_to_encryption_entry.clear();
        self.keyspace_active_shard_counts.clear();
        self.shard_to_current_key.clear();
    }
}

impl EncryptionKeyManager {
    pub fn new() -> Self {
        Self::new_with_cmek_health_check(true)
    }

    fn new_with_cmek_health_check(enable_cmek_health_check: bool) -> Self {
        let runtime = Arc::new(AsyncSafeRuntime::new(
            tokio::runtime::Builder::new_multi_thread()
                .worker_threads(1)
                .thread_name("encryption-manager")
                .enable_all()
                .build()
                .unwrap(),
        ));
        let (cmek_health_check_shutdown_tx, cmek_health_check_shutdown_rx) = watch::channel(false);

        let manager = Self {
            runtime,
            cmek_health_check_shutdown_tx,
            mode: Mutex::new(None),
            shard_to_current_key: DashMap::new(),
            keyspace_active_shard_counts: Arc::new(DashMap::new()),
            keyspace_to_encryption_entry: Arc::new(DashMap::new()),
            unencrypted_keyspaces: DashSet::new(),
        };
        if enable_cmek_health_check {
            manager.start_cmek_health_check_task(cmek_health_check_shutdown_rx);
        }
        manager
    }

    #[cfg(test)]
    fn new_for_test() -> Self {
        Self::new_with_cmek_health_check(false)
    }

    // Records the best-effort signal for future new-keyspace admission.
    // Metadata wins if both modes are observed so the final latch does not
    // depend on shard load order.
    pub fn set_mode_if_unset(&self, next: EncryptionMethod) {
        let mut mode = self.mode.lock().unwrap();
        match *mode {
            None => {
                *mode = Some(next);
            }
            Some(EncryptionMethod::Legacy) if next == EncryptionMethod::Metadata => {
                log::warn!(
                    "incoming metadata encryption mode overrides current legacy process mode, current_mode={:?}, incoming_mode={:?}",
                    EncryptionMethod::Legacy,
                    next,
                );
                *mode = Some(next);
            }
            Some(current_mode) if current_mode != next => {
                log::warn!(
                    "incoming encryption mode does not match current process mode, continuing: current_mode={:?}, incoming_mode={:?}",
                    current_mode,
                    next
                );
            }
            Some(_) => {}
        }
    }

    pub fn current_mode(&self) -> Option<EncryptionMethod> {
        *self.mode.lock().unwrap()
    }

    // Seeds the best-effort signal from a real-vendor global master key.
    pub fn try_set_legacy_mode_from_global_master_key(&self, master_key: &MasterKey) {
        if self.current_mode().is_none() && master_key.uses_real_kms_vendor() {
            self.set_mode_if_unset(EncryptionMethod::Legacy);
        }
    }

    /// Registers the given `EncryptionMeta` if needed and returns the parsed
    /// metadata and its current data key ID.
    fn register_encryption_meta_if_needed(
        &self,
        encryption_meta_bytes: &[u8],
        keyspace_id: Option<u32>,
    ) -> Result<(Arc<KeyspaceEncryptionEntry>, Arc<EncryptionMeta>, u32)> {
        let mut incoming = kvenginepb::EncryptionMeta::default();
        incoming
            .merge_from_bytes(encryption_meta_bytes)
            .map_err(|e| {
                EncryptionKeyError::InvalidEncryptionMeta(format!(
                    "failed to parse encryption meta: {}",
                    e
                ))
            })?;
        validate_encryption_meta(&incoming)?;
        let incoming_keyspace_id = incoming.get_keyspace_id();

        let target_keyspace_id = keyspace_id
            .filter(|keyspace_id| *keyspace_id > 0)
            .unwrap_or_else(|| incoming.get_keyspace_id());

        if target_keyspace_id != incoming_keyspace_id {
            log::warn!(
                "[CMEK] incoming encryption meta has a mismatched keyspace id; use target keyspace id as source of truth, target_keyspace_id={}, incoming_keyspace_id={}",
                target_keyspace_id,
                incoming_keyspace_id,
            );
        }
        self.unencrypted_keyspaces.remove(&target_keyspace_id);

        let entry = self
            .keyspace_to_encryption_entry
            .entry(target_keyspace_id)
            .or_insert_with(|| {
                Arc::new(KeyspaceEncryptionEntry::new(
                    target_keyspace_id,
                    &incoming,
                    self.runtime.clone(),
                ))
            })
            .clone();
        entry.update_cached_meta_if_preferred(&incoming);
        let cached = entry.encryption_meta();
        let incoming = if cached.as_ref() == &incoming {
            cached
        } else {
            Arc::new(incoming)
        };

        let data_key_id = incoming.get_current().data_key_id;
        if !incoming.data_keys.contains_key(&data_key_id) {
            return Err(EncryptionKeyError::DataKeyMissing(
                data_key_id,
                target_keyspace_id,
            ));
        }
        Ok((entry, incoming, data_key_id))
    }

    /// Registers the current encryption key used by a shard (region).
    pub fn register_current_key_for_shard(&self, shard_id: u64, k: EncryptionKey) -> Result<()> {
        let Some(keyspace_id) = k.keyspace_id else {
            return Err(EncryptionKeyError::InvalidEncryptionState(
                "keyspace id is not attached to the encryption key".into(),
            ));
        };
        if k.get_type() == KEY_TYPE_AES_256_CTR_LEGACY {
            // We don't expect the EncryptionKeyManager to deal with legacy keys.
            return Ok(());
        }
        self.register_current_key_ref_for_shard_inner(
            shard_id,
            EncryptionKeyRef {
                keyspace_id,
                data_key_id: k.get_key_id()?,
                unencrypted: k.get_type() == KEY_TYPE_UNENCRYPTED,
            },
        )
    }

    /// Registers the current encryption key state used by a shard (region).
    pub fn register_current_encryption_state_for_shard(
        &self,
        shard_id: u64,
        key_state: &EncryptionState,
    ) -> Result<()> {
        if matches!(
            key_state,
            EncryptionState::Invalid { .. } | EncryptionState::Disabled
        ) {
            self.deregister_current_key_for_shard(shard_id);
            return Ok(());
        }
        let Some(keyspace_id) = key_state.keyspace_id() else {
            return Ok(());
        };
        let registered_ref = EncryptionKeyRef {
            keyspace_id,
            data_key_id: key_state.data_key_id().ok_or_else(|| {
                EncryptionKeyError::InvalidEncryptionState(
                    "data key id is not attached to the encryption key".into(),
                )
            })?,
            unencrypted: false,
        };
        self.register_current_key_ref_for_shard_inner(shard_id, registered_ref)
    }

    fn register_current_key_ref_for_shard_inner(
        &self,
        shard_id: u64,
        key_state: EncryptionKeyRef,
    ) -> Result<()> {
        let keyspace_id = key_state.keyspace_id;
        match self
            .shard_to_current_key
            .insert(shard_id, key_state.clone())
        {
            Some(previous) if previous.keyspace_id == keyspace_id => {}
            Some(previous) => {
                if self.decrement_keyspace_active_shard_count(previous.keyspace_id) {
                    self.clear_master_key_cache_for_inactive_keyspace(previous.keyspace_id);
                }
                self.increment_keyspace_active_shard_count(keyspace_id);
            }
            None => self.increment_keyspace_active_shard_count(keyspace_id),
        }
        self.update_encryption_stats();
        Ok(())
    }

    /// Returns the current encryption key for the given shard, if available.
    ///
    /// Might be used by TiFlash FFI.
    pub fn current_key_by_shard(&self, shard_id: u64) -> Result<EncryptionKey> {
        let key_ref = self.shard_to_current_key.get(&shard_id).ok_or_else(|| {
            EncryptionKeyError::InvalidEncryptionState(format!(
                "no data key found for shard {}",
                shard_id
            ))
        })?;
        if key_ref.unencrypted {
            return Ok(EncryptionKey::new_unencrypted_key().with_keyspace_id(key_ref.keyspace_id));
        }
        let entry = self
            .keyspace_to_encryption_entry
            .get(&key_ref.keyspace_id)
            .ok_or_else(|| {
                EncryptionKeyError::InvalidEncryptionState(format!(
                    "encryption info unknown for keyspace {}",
                    key_ref.keyspace_id
                ))
            })?;
        entry.get_data_key(key_ref.data_key_id)
    }

    pub fn encryption_meta_by_keyspace(&self, keyspace_id: u32) -> Option<Arc<EncryptionMeta>> {
        self.keyspace_to_encryption_entry
            .get(&keyspace_id)
            .map(|entry| entry.value().encryption_meta())
    }

    /// Deregisters the encryption key associated with a shard.
    pub fn deregister_current_key_for_shard(&self, shard_id: u64) {
        if let Some((_, key_ref)) = self.shard_to_current_key.remove(&shard_id) {
            if self.decrement_keyspace_active_shard_count(key_ref.keyspace_id) {
                self.clear_master_key_cache_for_inactive_keyspace(key_ref.keyspace_id);
            }
            self.update_encryption_stats();
        }
    }

    fn increment_keyspace_active_shard_count(&self, keyspace_id: u32) {
        match self.keyspace_active_shard_counts.entry(keyspace_id) {
            Entry::Occupied(mut entry) => *entry.get_mut() += 1,
            Entry::Vacant(entry) => {
                entry.insert(1);
            }
        }
    }

    fn decrement_keyspace_active_shard_count(&self, keyspace_id: u32) -> bool {
        match self.keyspace_active_shard_counts.entry(keyspace_id) {
            Entry::Occupied(mut entry) => {
                if *entry.get() == 1 {
                    entry.remove();
                    true
                } else {
                    *entry.get_mut() -= 1;
                    false
                }
            }
            Entry::Vacant(_) => false,
        }
    }

    fn clear_master_key_cache_for_inactive_keyspace(&self, keyspace_id: u32) {
        if let Some(entry) = self.keyspace_to_encryption_entry.get(&keyspace_id) {
            entry.clear_master_key_cache();
        }
    }

    fn update_encryption_stats(&self) {
        ENCRYPTION_REGION_COUNT_GAUGE.set(self.shard_to_current_key.len() as i64);
        ENCRYPTION_KEYSPACE_ACTIVE_COUNT_GAUGE.set(self.keyspace_active_shard_counts.len() as i64);
        ENCRYPTION_KEYSPACE_REGISTERED_COUNT_GAUGE
            .set(self.keyspace_to_encryption_entry.len() as i64);
        ENCRYPTION_KEYSPACE_CMEK_HEALTH_CHECK_FAILURE_GAUGE
            .set(self.count_cmek_health_check_failures() as i64);
    }

    /// Returns a reference to the internal runtime used for KMS operations.
    pub fn get_runtime(&self) -> &tokio::runtime::Runtime {
        self.runtime.as_ref()
    }

    /// Returns the current encryption key for the given keyspace, if available.
    pub fn current_key_by_keyspace(&self, keyspace_id: u32) -> Result<Option<EncryptionKey>> {
        if self.unencrypted_keyspaces.contains(&keyspace_id) {
            return Ok(None);
        }

        if let Some(entry) = self.keyspace_to_encryption_entry.get(&keyspace_id) {
            let data_key_id = entry.current_data_key_id();
            let data_key = entry.get_data_key(data_key_id)?;
            return Ok(Some(data_key));
        }

        Err(EncryptionKeyError::EncryptionInfoUnknownForKeyspace(
            keyspace_id,
        ))
    }

    /// Marks a keyspace as unencrypted.
    pub fn mark_keyspace_as_unencrypted(&self, keyspace_id: u32) {
        self.unencrypted_keyspaces.insert(keyspace_id);
        self.update_encryption_stats();
    }

    /// Returns the master key from the given `EncryptionMeta`.
    /// If the keyspace has not been registered yet, it will be registered
    /// first.
    pub fn get_or_decrypt_master_key(&self, m: &EncryptionMeta) -> Result<Arc<MasterKey>> {
        let entry = self
            .keyspace_to_encryption_entry
            .entry(m.keyspace_id)
            .or_insert_with(|| {
                Arc::new(KeyspaceEncryptionEntry::new(
                    m.keyspace_id,
                    m,
                    self.runtime.clone(),
                ))
            })
            .clone();
        entry.decrypt_master_key(m)
    }

    /// Retrieves the data key for a given keyspace and encryption header.
    fn get_data_key_by_header(
        &self,
        keyspace_id: u32,
        encryption_header: u32,
        encryption_meta: Option<&Arc<EncryptionMeta>>,
    ) -> Result<EncryptionKey> {
        if encryption_header == 0 {
            return Ok(EncryptionKey::new_unencrypted_key().with_keyspace_id(keyspace_id));
        }
        let encryption_format = (encryption_header >> 24) as u8;
        match encryption_format {
            ENCRYPTION_DATA_FORMAT_LEGACY => {
                Err(EncryptionKeyError::InvalidEncryptionMeta(format!(
                    "legacy encryption header {:#010x} is unsupported by EncryptionKeyManager",
                    encryption_header
                )))
            }
            ENCRYPTION_DATA_FORMAT => {
                let data_key_id = encryption_header & MAX_DATA_KEY_ID;
                let entry = self
                    .keyspace_to_encryption_entry
                    .get(&keyspace_id)
                    .ok_or_else(|| {
                        EncryptionKeyError::InvalidEncryptionState(format!(
                            "encryption info unknown for keyspace {}",
                            keyspace_id
                        ))
                    })?;
                if let Some(encryption_meta) = encryption_meta {
                    entry.get_data_key_from_meta(encryption_meta, data_key_id)
                } else {
                    entry.get_data_key(data_key_id)
                }
            }
            _ => Err(EncryptionKeyError::InvalidEncryptionMeta(format!(
                "unsupported encryption header {:#010x}",
                encryption_header
            ))),
        }
    }

    fn start_cmek_health_check_task(&self, mut shutdown_rx: watch::Receiver<bool>) {
        let runtime = self.runtime.clone();
        let active_keyspace_counts = self.keyspace_active_shard_counts.clone();
        let keyspace_to_encryption_entry = self.keyspace_to_encryption_entry.clone();
        runtime.spawn(async move {
            let mut ticker = tokio::time::interval(CMEK_HEALTH_CHECK_INTERVAL);
            ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

            loop {
                tokio::select! {
                    changed = shutdown_rx.changed() => {
                        if changed.is_err() || *shutdown_rx.borrow() {
                            break;
                        }
                    }
                    _ = ticker.tick() => {
                        Self::run_cmek_health_check_once(
                            active_keyspace_counts.as_ref(),
                            keyspace_to_encryption_entry.as_ref(),
                        )
                        .await;
                    }
                }
            }
        });
    }

    fn count_cmek_health_check_failures(&self) -> usize {
        self.keyspace_active_shard_counts
            .iter()
            .filter_map(|entry| self.keyspace_to_encryption_entry.get(entry.key()))
            .filter(|entry| entry.health_check_status() == KeyspaceCmekHealthStatus::Failed)
            .count()
    }

    async fn run_cmek_health_check_once(
        active_keyspace_counts: &DashMap<u32, usize>,
        keyspace_to_encryption_entry: &DashMap<u32, Arc<KeyspaceEncryptionEntry>>,
    ) {
        let active_keyspace_ids: Vec<u32> = active_keyspace_counts
            .iter()
            .map(|entry| *entry.key())
            .collect();
        let mut failed_keyspaces = 0;

        for keyspace_id in active_keyspace_ids {
            let Some(entry) = keyspace_to_encryption_entry
                .get(&keyspace_id)
                .map(|entry| entry.clone())
            else {
                continue;
            };
            match entry.check_cmek_health().await {
                CheckCmekHealthResult::Healthy { recovered } => {
                    if recovered {
                        log::info!(
                            "CMEK encryption metadata health check recovered for keyspace {}",
                            keyspace_id
                        );
                    }
                }
                CheckCmekHealthResult::Failed { newly_failed, err } => {
                    failed_keyspaces += 1;
                    log::warn!(
                        "CMEK encryption metadata health check failed for keyspace {}: {:?} (newly_failed={})",
                        keyspace_id,
                        err,
                        newly_failed
                    );
                }
            }
        }

        ENCRYPTION_KEYSPACE_CMEK_HEALTH_CHECK_FAILURE_GAUGE.set(failed_keyspaces as i64);
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum KeyspaceCmekHealthStatus {
    Unknown,
    Healthy,
    Failed,
}

enum CheckCmekHealthResult {
    Healthy {
        recovered: bool,
    },
    Failed {
        newly_failed: bool,
        err: EncryptionKeyError,
    },
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
enum MasterKeyRetryState {
    #[default]
    Idle,
    Retrying,
    PermanentlyFailed,
}

#[derive(Default)]
struct MasterKeyCacheState {
    master_key: Option<Arc<MasterKey>>,
    last_err: Option<String>,
    retry_state: MasterKeyRetryState,
}

fn validate_encryption_meta(meta: &EncryptionMeta) -> Result<()> {
    if !meta.has_current() {
        return Err(EncryptionKeyError::InvalidEncryptionMeta(
            "missing current encryption epoch".to_owned(),
        ));
    }
    if !meta.has_master_key() {
        return Err(EncryptionKeyError::InvalidEncryptionMeta(
            "missing master key".to_owned(),
        ));
    }
    let current = meta.get_current();
    if !meta.data_keys.contains_key(&current.data_key_id) {
        return Err(EncryptionKeyError::InvalidEncryptionMeta(format!(
            "current data key {} missing in keyspace {}",
            current.data_key_id, meta.keyspace_id
        )));
    }
    Ok(())
}

/// Holds all encryption-related info for a single keyspace.
///
/// It handles the caching and lookup of encryption metadata, data keys,
/// and master keys.
struct KeyspaceEncryptionEntry {
    // The keyspace this entry is registered for. It can differ from the
    // keyspace_id inside the cached `EncryptionMeta` in restore scenarios where
    // the serialized meta is reused across clusters.
    keyspace_id: u32,

    // Tokio runtime for handling KMS operations (master key decryption).
    runtime: Arc<AsyncSafeRuntime>,

    // meta is read-often, updated on rotation
    meta: RwLock<Arc<EncryptionMeta>>,

    // data_key_id -> key resolved from the current cached metadata
    data_key_cache: DashMap<u32, EncryptionKey>,

    // master key config -> decrypted master key, or a decrypt failure.
    // Retryable failures are retried in the background so foreground
    // read/apply paths do not block on KMS after the first failure.
    master_key_cache: DashMap<MasterKeyConfig, Arc<Mutex<MasterKeyCacheState>>>,

    // Tracks the last periodic health-check result for this keyspace.
    health_check_status: Mutex<KeyspaceCmekHealthStatus>,
}

impl KeyspaceEncryptionEntry {
    pub fn new(
        keyspace_id: u32,
        encryption_meta: &EncryptionMeta,
        rt: Arc<AsyncSafeRuntime>,
    ) -> Self {
        Self {
            keyspace_id,
            meta: RwLock::new(Arc::new(encryption_meta.clone())),
            data_key_cache: DashMap::new(),
            master_key_cache: DashMap::new(),
            runtime: rt,
            health_check_status: Mutex::new(KeyspaceCmekHealthStatus::Unknown),
        }
    }

    pub fn update_cached_meta_if_preferred(&self, incoming: &kvenginepb::EncryptionMeta) {
        let mut g = self
            .meta
            .write()
            .expect("RwLock on encryption meta is poisoned");
        let cached = g.as_ref();
        let cached_file_id = cached.get_current().get_file_id();
        let incoming_file_id = incoming.get_current().get_file_id();
        let master_key_changed =
            Self::build_master_key_config(cached) != Self::build_master_key_config(incoming);
        let prefer_due_to_rotation =
            self.is_rotated_config(incoming) && incoming_file_id > cached_file_id;
        let prefer_due_to_restore =
            self.is_restored_config(incoming) && self.is_initial_config(cached);
        // Handle the special case where a keyspace is restored from another
        // keyspace. The target keyspace is expected to be empty, but it may
        // already have its own initial encryption config from keyspace
        // creation. That initial config must be overridden by the incoming
        // config from the source keyspace so the target keyspace can read SST
        // files encrypted by the source keyspace.
        //
        // That restored-config override is the only special case. After it
        // happens, later metas are expected to belong to the target keyspace
        // again, so matching metas should advance this cache by file id.
        if prefer_due_to_rotation || prefer_due_to_restore {
            log::info!(
                "[CMEK] update cached encryption meta: keyspace_id={}, cached_meta_file_id={}, incoming_meta_file_id={}, cached_meta_keyspace_id={}, incoming_keyspace_id={}, cached_data_key_count={}, incoming_data_key_count={}, prefer_due_to_rotation={}, prefer_due_to_restore={}, master_key_changed={}",
                self.keyspace_id,
                cached_file_id,
                incoming_file_id,
                cached.get_keyspace_id(),
                incoming.get_keyspace_id(),
                cached.get_data_keys().len(),
                incoming.get_data_keys().len(),
                prefer_due_to_rotation,
                prefer_due_to_restore,
                master_key_changed,
            );
            *g = Arc::new(incoming.clone());
            if prefer_due_to_restore || master_key_changed {
                self.data_key_cache.clear();
            }
            if master_key_changed {
                self.clear_master_key_cache();
            }
            *self.health_check_status.lock().unwrap() = KeyspaceCmekHealthStatus::Unknown;
        }
    }

    fn clear_master_key_cache(&self) {
        self.master_key_cache.clear();
    }

    /// Returns true when the meta still looks like the first config created
    /// for a brand-new encrypted keyspace.
    fn is_initial_config(&self, meta: &kvenginepb::EncryptionMeta) -> bool {
        meta.get_keyspace_id() == self.keyspace_id && meta.get_data_keys().len() <= 1
    }

    /// Returns true when the meta originates from a different keyspace.
    fn is_restored_config(&self, meta: &kvenginepb::EncryptionMeta) -> bool {
        meta.get_keyspace_id() != self.keyspace_id
    }

    fn is_rotated_config(&self, meta: &kvenginepb::EncryptionMeta) -> bool {
        meta.get_keyspace_id() == self.keyspace_id && meta.get_data_keys().len() > 1
    }

    fn encryption_meta(&self) -> Arc<EncryptionMeta> {
        self.meta
            .read()
            .expect("RwLock on encryption meta is poisoned")
            .clone()
    }

    fn current_data_key_id(&self) -> u32 {
        self.meta
            .read()
            .expect("RwLock on encryption meta is poisoned")
            .get_current()
            .data_key_id
    }

    /// Retrieves or decrypts the data key from the cached keyspace metadata and
    /// caches it.
    fn get_data_key(&self, data_key_id: u32) -> Result<EncryptionKey> {
        let cached_meta = self
            .meta
            .read()
            .expect("RwLock on encryption meta is poisoned");
        if let Some(key) = self.data_key_cache.get(&data_key_id) {
            return Ok(key.value().clone());
        }
        let encryption_meta = cached_meta.clone();
        drop(cached_meta);
        self.resolve_data_key_from_meta(&encryption_meta, data_key_id)
    }

    /// Retrieves or decrypts the data key from the given metadata.
    fn get_data_key_from_meta(
        &self,
        encryption_meta: &Arc<EncryptionMeta>,
        data_key_id: u32,
    ) -> Result<EncryptionKey> {
        let cached_meta = self
            .meta
            .read()
            .expect("RwLock on encryption meta is poisoned");
        if Arc::ptr_eq(encryption_meta, &cached_meta) {
            if let Some(key) = self.data_key_cache.get(&data_key_id) {
                return Ok(key.value().clone());
            }
        }
        drop(cached_meta);
        self.resolve_data_key_from_meta(encryption_meta, data_key_id)
    }

    /// Resolves the requested data key from the provided metadata.
    fn resolve_data_key_from_meta(
        &self,
        encryption_meta: &Arc<EncryptionMeta>,
        data_key_id: u32,
    ) -> Result<EncryptionKey> {
        let start = Instant::now_coarse();
        defer!({
            ENCRYPTION_KEY_GET_DURATION.observe(start.saturating_elapsed().as_secs_f64());
        });
        let data_key = encryption_meta
            .data_keys
            .get(&data_key_id)
            .ok_or_else(|| EncryptionKeyError::DataKeyMissing(data_key_id, self.keyspace_id))?;

        let master_key = self.decrypt_master_key(encryption_meta.as_ref())?;
        let key = master_key
            .decrypt_encryption_key(&data_key.ciphertext)
            .and_then(|key| {
                if key.get_type() != KEY_TYPE_AES_256_CTR {
                    return Err(EncryptionKeyError::InvalidEncryptionMeta(format!(
                        "unexpected data key type {} in keyspace {}",
                        key.get_type(),
                        self.keyspace_id
                    )));
                }
                key.with_key_id(data_key_id)
            })?
            .with_keyspace_id(self.keyspace_id);
        self.cache_data_key_if_current(encryption_meta, data_key_id, key.clone());
        Ok(key)
    }

    fn cache_data_key_if_current(
        &self,
        encryption_meta: &Arc<EncryptionMeta>,
        data_key_id: u32,
        key: EncryptionKey,
    ) {
        let cached_meta = self
            .meta
            .read()
            .expect("RwLock on encryption meta is poisoned");
        if !Arc::ptr_eq(encryption_meta, &cached_meta) {
            return;
        }
        if let Entry::Vacant(slot) = self.data_key_cache.entry(data_key_id) {
            slot.insert(key);
        }
    }

    /// Decrypts the master key from the given `EncryptionMeta`.
    ///
    /// If the decrypted master key is already cached, it is returned directly.
    /// Otherwise, it is decrypted via KMS and stored for future reuse.
    pub fn decrypt_master_key(&self, m: &EncryptionMeta) -> Result<Arc<MasterKey>> {
        if !m.has_master_key() {
            return Err(EncryptionKeyError::InvalidEncryptionMeta(format!(
                "missing master key in keyspace {}",
                m.keyspace_id
            )));
        }
        let master_key_config = Self::build_master_key_config(m);

        // The mutex single-flights the first foreground KMS attempt. Retry
        // state is set only after that attempt fails.
        let entry = self
            .master_key_cache
            .entry(master_key_config.clone())
            .or_insert_with(|| Arc::new(Mutex::new(MasterKeyCacheState::default())))
            .clone();

        let mut guard = entry.lock().unwrap();
        if let Some(master_key) = &guard.master_key {
            return Ok(master_key.clone());
        }
        match guard.retry_state {
            MasterKeyRetryState::Idle => {}
            MasterKeyRetryState::Retrying => {
                return Err(EncryptionKeyError::MasterKeyUnavailable(format!(
                    "master key decrypt is retrying: {}",
                    guard.last_err.as_deref().unwrap_or("unknown error")
                )));
            }
            MasterKeyRetryState::PermanentlyFailed => {
                return Err(EncryptionKeyError::MasterKeyUnavailable(format!(
                    "master key decrypt permanently failed: {}",
                    guard.last_err.as_deref().unwrap_or("unknown error")
                )));
            }
        }

        let fut = master_key_config.decrypt(false);
        let master_key = if tokio::runtime::Handle::try_current().is_ok() {
            // Inside a Tokio runtime (e.g. in
            // `TxnChunkHandler::acquire_keyspace_info`), calling
            // `block_on` directly would stall the scheduler. Use
            // `block_in_place` instead.
            tokio::task::block_in_place(|| self.runtime.block_on(fut))
        } else {
            self.runtime.block_on(fut)
        };
        let master_key = match master_key {
            Ok(master_key) => master_key,
            Err(err) => {
                let retryable = err.is_retryable();
                guard.last_err = Some(err.to_string());
                if retryable {
                    guard.retry_state = MasterKeyRetryState::Retrying;
                } else {
                    guard.retry_state = MasterKeyRetryState::PermanentlyFailed;
                }

                if retryable {
                    self.spawn_master_key_retry_loop(
                        Arc::downgrade(&entry),
                        master_key_config.clone(),
                        m.keyspace_id,
                        m.get_current().data_key_id,
                    );
                }
                log::warn!(
                    "[CMEK] master key decrypt failed for keyspace: keyspace_id={}, data_key_id={}, cmek_id={}, vendor={}, retryable={}, err={}",
                    m.keyspace_id,
                    m.get_current().data_key_id,
                    master_key_config.key_id,
                    master_key_config.vendor,
                    retryable,
                    err,
                );
                return Err(err);
            }
        };
        let arc_key = Arc::new(master_key);
        guard.master_key = Some(arc_key.clone());
        guard.last_err = None;
        guard.retry_state = MasterKeyRetryState::Idle;
        Ok(arc_key)
    }

    fn spawn_master_key_retry_loop(
        &self,
        entry: Weak<Mutex<MasterKeyCacheState>>,
        master_key_config: MasterKeyConfig,
        keyspace_id: u32,
        data_key_id: u32,
    ) {
        self.runtime.spawn(async move {
            let mut backoff = ExponentialBackoff::new(
                MASTER_KEY_RETRY_INITIAL_INTERVAL,
                MASTER_KEY_RETRY_MAX_INTERVAL,
                usize::MAX,
            );
            loop {
                let retry_interval = match backoff.next_delay() {
                    Ok(delay) => delay,
                    Err(err) => {
                        log::warn!(
                            "[CMEK] master key decrypt retry stopped: keyspace_id={}, data_key_id={}, cmek_id={}, vendor={}, err={}",
                            keyspace_id,
                            data_key_id,
                            master_key_config.key_id,
                            master_key_config.vendor,
                            err,
                        );
                        return;
                    }
                };
                tokio::time::sleep(retry_interval).await;
                let Some(entry) = entry.upgrade() else {
                    log::info!(
                        "[CMEK] master key decrypt retry stopped because cache entry is no longer referenced: keyspace_id={}, data_key_id={}, cmek_id={}, vendor={}",
                        keyspace_id,
                        data_key_id,
                        master_key_config.key_id,
                        master_key_config.vendor,
                    );
                    return;
                };
                // `decrypt` observes master key failure metrics on every
                // failed attempt.
                let result = master_key_config.decrypt(false).await;
                match result {
                    Ok(master_key) => {
                        let mut guard = entry.lock().unwrap();
                        guard.master_key = Some(Arc::new(master_key));
                        guard.last_err = None;
                        guard.retry_state = MasterKeyRetryState::Idle;
                        log::info!(
                            "[CMEK] master key decrypt recovered in background: keyspace_id={}, data_key_id={}, cmek_id={}, vendor={}",
                            keyspace_id,
                            data_key_id,
                            master_key_config.key_id,
                            master_key_config.vendor,
                        );
                        return;
                    }
                    Err(err) => {
                        let retryable = err.is_retryable();
                        {
                            let mut guard = entry.lock().unwrap();
                            guard.last_err = Some(err.to_string());
                            if !retryable {
                                guard.retry_state = MasterKeyRetryState::PermanentlyFailed;
                            }
                        }

                        log::warn!(
                            "[CMEK] background master key decrypt failed: keyspace_id={}, data_key_id={}, cmek_id={}, vendor={}, retryable={}, attempts={}, retry_interval_secs={}, err={}",
                            keyspace_id,
                            data_key_id,
                            master_key_config.key_id,
                            master_key_config.vendor,
                            retryable,
                            backoff.current_attempts(),
                            retry_interval.as_secs(),
                            err,
                        );
                        if !retryable {
                            return;
                        }
                    }
                }
            }
        });
    }

    fn build_master_key_config(m: &EncryptionMeta) -> MasterKeyConfig {
        let master_key = m.get_master_key();
        MasterKeyConfig {
            key_id: master_key.cmek_id.clone(),
            vendor: master_key.vendor.clone(),
            region: master_key.region.clone(),
            endpoint: master_key.endpoint.clone(),
            cipher_text: base64::encode(master_key.ciphertext.clone()),
            encryption_context: master_key_encryption_context(master_key),
        }
    }

    async fn check_cmek_health(&self) -> CheckCmekHealthResult {
        let meta = self.encryption_meta();
        let result = Self::build_master_key_config(&meta)
            .decrypt(false)
            .await
            .and_then(|master_key| {
                if master_key.is_valid() {
                    Ok(Arc::new(master_key))
                } else {
                    Err(EncryptionKeyError::InvalidMasterKey)
                }
            });

        let mut status = self.health_check_status.lock().unwrap();
        let previous = *status;
        match result {
            Ok(_) => {
                *status = KeyspaceCmekHealthStatus::Healthy;
                CheckCmekHealthResult::Healthy {
                    recovered: previous == KeyspaceCmekHealthStatus::Failed,
                }
            }
            Err(err) => {
                *status = KeyspaceCmekHealthStatus::Failed;
                CheckCmekHealthResult::Failed {
                    newly_failed: previous != KeyspaceCmekHealthStatus::Failed,
                    err,
                }
            }
        }
    }

    fn health_check_status(&self) -> KeyspaceCmekHealthStatus {
        *self.health_check_status.lock().unwrap()
    }
}

#[cfg(test)]
mod tests {
    use std::sync::{Barrier, Mutex as StdMutex};

    use lazy_static::lazy_static;
    use rand::Rng;

    use super::*;

    lazy_static! {
        static ref METRICS_TEST_LOCK: StdMutex<()> = StdMutex::new(());
    }

    fn reset_cmek_health_check_metric() {
        ENCRYPTION_KEYSPACE_CMEK_HEALTH_CHECK_FAILURE_GAUGE.set(0);
    }

    #[test]
    fn test_kms_vendor_parse() {
        let cases = [
            ("aws", KmsVendor::Aws),
            ("Aws", KmsVendor::Aws),
            ("AWS", KmsVendor::Aws),
            ("aliyun", KmsVendor::Aliyun),
            ("Aliyun", KmsVendor::Aliyun),
            ("tencentcloud", KmsVendor::Tencentcloud),
            ("volcengine", KmsVendor::Volcengine),
            ("test", KmsVendor::Test),
        ];

        for (input, expected) in cases {
            assert_eq!(expected, input.parse::<KmsVendor>().unwrap());
        }
    }

    #[test]
    fn test_master_key_encryption_context_round_trip() {
        let mut master_key = kvenginepb::MasterKey::default();
        let encryption_context = BTreeMap::from([
            ("ExternalId".to_owned(), "customer-42".to_owned()),
            ("tenant".to_owned(), "acme".to_owned()),
        ]);
        set_master_key_encryption_context(&mut master_key, &encryption_context);
        assert_eq!(
            master_key_encryption_context(&master_key),
            encryption_context
        );
    }

    #[test]
    fn test_encryption() {
        let master_key_plain_text = rand::thread_rng().gen::<[u8; 32]>().to_vec();
        let master_key = MasterKey::new(&master_key_plain_text);
        for encryption_key in [
            master_key.generate_encryption_key(),
            master_key.generate_encryption_key_v2(),
        ] {
            let mut encrypted = vec![];
            encryption_key.encrypt(b"hello", 123, 12, &mut encrypted);

            // Test the key export code path.
            let exported = encryption_key.export();
            let encryption_key = master_key
                .decrypt_encryption_key(&exported)
                .expect("failed to decrypt encryption key");

            let mut decrypted = vec![];
            encryption_key.decrypt(&encrypted, 123, 12, &mut decrypted);
            assert_eq!(decrypted, b"hello");
        }
    }

    fn new_test_runtime() -> Arc<tokio::runtime::Runtime> {
        Arc::new(
            tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap(),
        )
    }

    fn build_test_encryption_meta(keyspace_id: u32) -> (kvenginepb::EncryptionMeta, EncryptionKey) {
        let mut config = MasterKeyConfig::default();
        config.vendor = "test".to_owned();
        let (master_key, master_key_ciphertext) = new_test_runtime()
            .block_on(config.generate_new_master_key())
            .unwrap();

        let data_key_id = generate_data_key_id();
        let data_key = master_key
            .generate_encryption_key_v2()
            .with_key_id(data_key_id)
            .unwrap();
        let mut meta = kvenginepb::EncryptionMeta::default();
        meta.set_keyspace_id(keyspace_id);

        let mut mk = kvenginepb::MasterKey::default();
        mk.vendor = config.vendor;
        mk.cmek_id = config.key_id;
        mk.region = config.region;
        mk.endpoint = config.endpoint;
        mk.ciphertext = master_key_ciphertext.to_vec();
        set_master_key_encryption_context(&mut mk, &config.encryption_context);
        meta.set_master_key(mk);

        meta.mut_current().set_data_key_id(data_key_id);
        let mut dk = kvenginepb::DataKey::new();
        dk.set_ciphertext(data_key.export());
        meta.mut_data_keys().insert(data_key_id, dk);
        (meta, data_key)
    }

    fn add_test_data_key(meta: &mut kvenginepb::EncryptionMeta) -> EncryptionKey {
        let master_key = MasterKeyConfig {
            key_id: meta.get_master_key().get_cmek_id().to_owned(),
            vendor: meta.get_master_key().get_vendor().to_owned(),
            region: meta.get_master_key().get_region().to_owned(),
            endpoint: meta.get_master_key().get_endpoint().to_owned(),
            cipher_text: base64::encode(meta.get_master_key().get_ciphertext()),
            encryption_context: master_key_encryption_context(meta.get_master_key()),
        };
        let master_key = new_test_runtime()
            .block_on(master_key.decrypt(false))
            .unwrap();
        let data_key = loop {
            let data_key_id = generate_data_key_id();
            if !meta.get_data_keys().contains_key(&data_key_id) {
                break master_key
                    .generate_encryption_key_v2()
                    .with_key_id(data_key_id)
                    .unwrap();
            }
        };
        let mut dk = kvenginepb::DataKey::new();
        dk.set_ciphertext(data_key.export());
        meta.mut_data_keys()
            .insert(data_key.get_key_id().unwrap(), dk);
        data_key
    }

    #[test]
    fn test_master_key_gen() {
        let mut config = MasterKeyConfig::default();
        config.vendor = "test".to_owned();
        let runtime = new_test_runtime();
        let (master_key, encrypted_key) =
            runtime.block_on(config.generate_new_master_key()).unwrap();
        config.cipher_text = base64::encode(encrypted_key.to_vec());
        let decoded_master_key = runtime.block_on(config.decrypt(false)).unwrap();
        assert_eq!(master_key.export(), decoded_master_key.export());
    }

    #[test]
    fn test_master_key_config_decrypt_rejects_invalid_aws_ciphertext() {
        let config = MasterKeyConfig {
            vendor: "aws".to_owned(),
            key_id: "test-key".to_owned(),
            cipher_text: "%%%not-base64%%%".to_owned(),
            ..Default::default()
        };

        let err = match new_test_runtime().block_on(config.decrypt(false)) {
            Ok(_) => panic!("invalid aws ciphertext should return error"),
            Err(err) => err,
        };
        assert!(matches!(err, EncryptionKeyError::InvalidEncryptionMeta(_)));
    }

    #[test]
    fn test_master_key_config_decrypt_rejects_unsupported_vendor() {
        let _guard = METRICS_TEST_LOCK
            .lock()
            .unwrap_or_else(|err| err.into_inner());

        let critical_error = CRITICAL_ERROR.with_label_values(&["cmek_master_key_decrypt"]);
        let before = critical_error.get();

        let config = MasterKeyConfig {
            vendor: "unknown-vendor".to_owned(),
            key_id: "test-key".to_owned(),
            ..Default::default()
        };
        let err = match new_test_runtime().block_on(config.decrypt(false)) {
            Ok(_) => panic!("unsupported vendor should return error"),
            Err(err) => err,
        };
        assert!(matches!(
            err,
            EncryptionKeyError::UnsupportedVendor(v) if v == "unknown-vendor"
        ));
        assert_eq!(critical_error.get(), before + 1);
    }

    #[test]
    fn test_master_key_config_decrypt_empty_vendor_only_allowed_for_legacy() {
        let config = MasterKeyConfig::default();
        let master_key = new_test_runtime().block_on(config.decrypt(true)).unwrap();
        assert!(!master_key.is_valid());

        let config = MasterKeyConfig::default();
        let err = match new_test_runtime().block_on(config.decrypt(false)) {
            Ok(_) => panic!("empty CMEK vendor should return error"),
            Err(err) => err,
        };
        assert!(matches!(
            err,
            EncryptionKeyError::UnsupportedVendor(v) if v.is_empty()
        ));
    }

    #[test]
    fn test_failed_master_key_decrypt_observes_critical_error_on_early_return() {
        let _guard = METRICS_TEST_LOCK
            .lock()
            .unwrap_or_else(|err| err.into_inner());

        let critical_error = CRITICAL_ERROR.with_label_values(&["cmek_master_key_decrypt"]);
        let before = critical_error.get();

        let config = MasterKeyConfig {
            vendor: "aws".to_owned(),
            key_id: "test-key".to_owned(),
            cipher_text: "%%%not-base64%%%".to_owned(),
            ..Default::default()
        };
        let result = new_test_runtime().block_on(config.decrypt(false));
        assert!(
            result.is_err(),
            "invalid AWS master key ciphertext should fail before calling KMS"
        );

        assert_eq!(critical_error.get(), before + 1);
    }

    #[test]
    fn test_generate_encryption_key_default_and_v2_types() {
        let master_key_plain_text = rand::thread_rng().gen::<[u8; 32]>().to_vec();
        let master_key = MasterKey::new(&master_key_plain_text);

        let legacy_key = master_key.generate_encryption_key();
        assert_eq!(legacy_key.get_type(), KEY_TYPE_AES_256_CTR_LEGACY);
        assert_eq!(legacy_key.export().len(), KEY_EXPORTED_SIZE_LEGACY);

        let v2_key = master_key.generate_encryption_key_v2();
        assert_eq!(v2_key.get_type(), KEY_TYPE_AES_256_CTR);
        assert_eq!(v2_key.export().len(), KEY_EXPORTED_SIZE);
    }

    #[test]
    fn test_aes_256_ctr_none_iv_is_compatible_with_zero_iv() {
        let key = rand::thread_rng().gen::<[u8; 32]>();
        let data = b"compatibility-check-for-none-vs-zero-iv";
        let zero_iv = [0u8; 16];

        let cipher_text_zero = symm::encrypt(Cipher::aes_256_ctr(), &key, Some(&zero_iv), data)
            .expect("encrypt with zero iv");

        // If the current openssl version supports None IV, verify it is compatible with
        // zero IV.
        if let Ok(cipher_text_none) = symm::encrypt(Cipher::aes_256_ctr(), &key, None, data) {
            assert_eq!(cipher_text_none, cipher_text_zero);

            if let Ok(plain_text_none) =
                symm::decrypt(Cipher::aes_256_ctr(), &key, None, &cipher_text_none)
            {
                assert_eq!(plain_text_none, data);
            }
            let plain_text_zero_from_none = symm::decrypt(
                Cipher::aes_256_ctr(),
                &key,
                Some(&zero_iv),
                &cipher_text_none,
            )
            .expect("decrypt none iv data with zero iv");
            assert_eq!(plain_text_zero_from_none, data);
        }

        let plain_text_zero = symm::decrypt(
            Cipher::aes_256_ctr(),
            &key,
            Some(&zero_iv),
            &cipher_text_zero,
        )
        .expect("decrypt with zero iv");

        assert_eq!(plain_text_zero, data);
    }

    #[test]
    fn test_generate_data_key_id_fits_header_width() {
        let data_key_id = generate_data_key_id();
        assert!(data_key_id <= MAX_DATA_KEY_ID);
    }

    #[test]
    fn test_decrypt_encryption_key_rejects_empty_exported_key() {
        let master_key_plain_text = rand::thread_rng().gen::<[u8; 32]>().to_vec();
        let master_key = MasterKey::new(&master_key_plain_text);

        let err = master_key
            .decrypt_encryption_key(&[])
            .expect_err("empty exported key should return error");
        assert!(matches!(
            err,
            EncryptionKeyError::InvalidExportedKey(msg) if msg == "missing key type byte"
        ));
    }

    #[test]
    fn test_decrypt_encryption_key_rejects_invalid_legacy_version() {
        let master_key_plain_text = rand::thread_rng().gen::<[u8; 32]>().to_vec();
        let master_key = MasterKey::new(&master_key_plain_text);
        let legacy_key = master_key.generate_encryption_key();
        let mut exported = legacy_key.export();
        exported[1..5].copy_from_slice(&1u32.to_be_bytes());

        let err = master_key
            .decrypt_encryption_key(&exported)
            .expect_err("legacy key with invalid version should return error");
        assert!(matches!(
            err,
            EncryptionKeyError::InvalidExportedKey(msg)
                if msg == "unsupported legacy key version 1"
        ));
    }

    #[test]
    fn test_v2_encryption_header_requires_bound_data_key_id() {
        let master_key_plain_text = rand::thread_rng().gen::<[u8; 32]>().to_vec();
        let master_key = MasterKey::new(&master_key_plain_text);
        let data_key = master_key.generate_encryption_key_v2();

        let err = data_key
            .get_key_id()
            .expect_err("unbound v2 key should reject get_key_id");
        assert!(matches!(err, EncryptionKeyError::InvalidEncryptionState(_)));

        let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            data_key.encryption_header()
        }));
        panic.expect_err("unbound v2 key should panic when encoding header");
    }

    #[test]
    fn test_with_key_id_rejects_out_of_range_ids() {
        let master_key_plain_text = rand::thread_rng().gen::<[u8; 32]>().to_vec();
        let master_key = MasterKey::new(&master_key_plain_text);
        let data_key = master_key.generate_encryption_key_v2();

        let err = data_key
            .with_key_id(MAX_DATA_KEY_ID + 1)
            .expect_err("out-of-range v2 data_key_id should return error");
        assert!(matches!(err, EncryptionKeyError::InvalidEncryptionState(_)));
    }

    #[test]
    fn test_encryption_key_manager() {
        let mut config = MasterKeyConfig::default();
        config.vendor = "test".to_owned();
        let (master_key, master_key_ciphertext) = new_test_runtime()
            .block_on(config.generate_new_master_key())
            .unwrap();

        let data_key1 = master_key
            .generate_encryption_key_v2()
            .with_key_id(1)
            .unwrap();
        let data_key2 = master_key
            .generate_encryption_key_v2()
            .with_key_id(2)
            .unwrap();
        let keyspace_id = 100;

        let mut m = kvenginepb::EncryptionMeta::default();
        m.set_keyspace_id(keyspace_id);

        let mut mk = kvenginepb::MasterKey::default();
        mk.vendor = config.vendor.clone();
        mk.ciphertext = master_key_ciphertext.to_vec();
        m.set_master_key(mk);

        m.mut_current().data_key_id = data_key2.get_key_id().unwrap();
        for k in [&data_key1, &data_key2] {
            let mut dk = kvenginepb::DataKey::new();
            dk.set_ciphertext(k.export());
            m.mut_data_keys().insert(k.get_key_id().unwrap(), dk);
        }

        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        let k = current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &m.write_to_bytes().unwrap(),
            None,
        )
        .expect("failed to get current key");
        let meta = mgr.encryption_meta_by_keyspace(keyspace_id).unwrap();
        assert_eq!(meta.get_keyspace_id(), keyspace_id);
        assert_eq!(meta.get_master_key().get_vendor(), config.vendor);
        assert_eq!(k.core.plain_text, data_key2.core.plain_text);

        // Test fetching data key by keyspace.
        let kk = mgr
            .current_key_by_keyspace(keyspace_id)
            .expect("get key should succeed")
            .expect("keyspace should have encryption key");
        assert_eq!(kk.core.plain_text, data_key2.core.plain_text);

        let err = mgr.current_key_by_keyspace(keyspace_id + 1).unwrap_err();
        assert!(matches!(
            err,
            EncryptionKeyError::EncryptionInfoUnknownForKeyspace(_)
        ));
        mgr.mark_keyspace_as_unencrypted(keyspace_id + 1);
        assert!(
            mgr.current_key_by_keyspace(keyspace_id + 1)
                .unwrap()
                .is_none()
        );

        // Test switching keys by encryption header.
        for dk in [&data_key1, &data_key2] {
            assert_eq!(
                k.switch_if_header_mismatch(dk.encryption_header())
                    .unwrap()
                    .core
                    .plain_text,
                dk.core.plain_text
            )
        }
        let mut rotated_meta = m.clone();
        rotated_meta.mut_current().set_file_id(1);
        let data_key3 = master_key
            .generate_encryption_key_v2()
            .with_key_id(3)
            .unwrap();
        rotated_meta
            .mut_current()
            .set_data_key_id(data_key3.get_key_id().unwrap());
        let mut dk = kvenginepb::DataKey::new();
        dk.set_ciphertext(data_key3.export());
        rotated_meta
            .mut_data_keys()
            .insert(data_key3.get_key_id().unwrap(), dk);
        let rotated_current_key = current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &rotated_meta.write_to_bytes().unwrap(),
            None,
        )
        .expect("failed to register rotated meta");
        assert_eq!(
            rotated_current_key.core.plain_text,
            data_key3.core.plain_text
        );
        let entry = mgr.keyspace_to_encryption_entry.get(&keyspace_id).unwrap();
        assert!(
            entry
                .data_key_cache
                .contains_key(&data_key1.get_key_id().unwrap())
        );
        assert!(
            entry
                .data_key_cache
                .contains_key(&data_key3.get_key_id().unwrap())
        );
        drop(entry);

        let missing_header = master_key
            .generate_encryption_key_v2()
            .with_key_id(4)
            .unwrap();
        let err = k
            .switch_if_header_mismatch(missing_header.encryption_header())
            .unwrap_err();
        assert!(matches!(err, EncryptionKeyError::DataKeyMissing(_, _)));

        // Ensure that `keyspace_id` is set on `data_key1`.
        let data_key1 = k
            .switch_if_header_mismatch(data_key1.encryption_header())
            .unwrap();
        assert!(matches!(
            mgr.current_key_by_shard(5),
            Err(EncryptionKeyError::InvalidEncryptionState(_))
        ));
        mgr.register_current_key_for_shard(5, data_key1.clone())
            .unwrap();
        mgr.current_key_by_shard(5).unwrap();
        assert_eq!(
            mgr.current_key_by_shard(5).unwrap().core.plain_text,
            data_key1.core.plain_text
        );
        mgr.register_current_encryption_state_for_shard(5, &EncryptionState::Disabled)
            .unwrap();
        assert!(matches!(
            mgr.current_key_by_shard(5),
            Err(EncryptionKeyError::InvalidEncryptionState(_))
        ));

        mgr.register_current_key_for_shard(5, data_key1.clone())
            .unwrap();
        mgr.register_current_encryption_state_for_shard(
            5,
            &EncryptionState::from_error(
                Some(keyspace_id),
                EncryptionKeyError::InvalidEncryptionState("bad encryption state".to_owned()),
            ),
        )
        .unwrap();
        assert!(matches!(
            mgr.current_key_by_shard(5),
            Err(EncryptionKeyError::InvalidEncryptionState(_))
        ));

        mgr.register_current_key_for_shard(5, data_key1.clone())
            .unwrap();
        mgr.deregister_current_key_for_shard(5);
        assert!(matches!(
            mgr.current_key_by_shard(5),
            Err(EncryptionKeyError::InvalidEncryptionState(_))
        ));
    }

    #[test]
    fn test_restored_config_replaces_initial_meta_without_losing_pending_key() {
        let target_keyspace_id = 42;
        let source_keyspace_id = 7;

        let (mut target_meta, target_current_key) = build_test_encryption_meta(target_keyspace_id);
        target_meta.mut_current().set_file_id(200);
        let shared_data_key_id = target_current_key.get_key_id().unwrap();

        let (mut incoming_meta, incoming_current_key) =
            build_test_encryption_meta(source_keyspace_id);
        incoming_meta.mut_current().set_file_id(100);
        let incoming_current_key = incoming_current_key
            .with_key_id(shared_data_key_id)
            .unwrap();
        incoming_meta
            .mut_current()
            .set_data_key_id(shared_data_key_id);
        incoming_meta.mut_data_keys().clear();
        let mut incoming_data_key = kvenginepb::DataKey::new();
        incoming_data_key.set_ciphertext(incoming_current_key.export());
        incoming_meta
            .mut_data_keys()
            .insert(shared_data_key_id, incoming_data_key);
        assert_ne!(
            target_current_key.core.plain_text,
            incoming_current_key.core.plain_text
        );

        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        let mut target_pending_key = encryption_state_from_meta_bytes(
            mgr.clone(),
            &target_meta.write_to_bytes().unwrap(),
            None,
        );
        assert!(
            mgr.keyspace_to_encryption_entry
                .get(&target_keyspace_id)
                .unwrap()
                .data_key_cache
                .is_empty(),
            "registering initial metadata should not eagerly resolve its data key"
        );

        let initial_cached_key = current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &target_meta.write_to_bytes().unwrap(),
            None,
        )
        .expect("failed to resolve initial target key");
        assert_eq!(
            initial_cached_key.core.plain_text,
            target_current_key.core.plain_text
        );
        assert_eq!(
            mgr.keyspace_to_encryption_entry
                .get(&target_keyspace_id)
                .unwrap()
                .data_key_cache
                .len(),
            1,
            "initial current metadata should populate the manager data-key cache"
        );

        let restored_pending_key = encryption_state_from_meta_bytes(
            mgr.clone(),
            &incoming_meta.write_to_bytes().unwrap(),
            Some(target_keyspace_id),
        );
        assert!(
            mgr.keyspace_to_encryption_entry
                .get(&target_keyspace_id)
                .unwrap()
                .data_key_cache
                .is_empty(),
            "replacing cached metadata must drop data keys from the old metadata"
        );

        let restored_key = restored_pending_key
            .resolve_ready_key()
            .expect("failed to resolve current key from incoming restore meta")
            .expect("restore meta should contain an encryption key");
        assert_eq!(
            mgr.keyspace_to_encryption_entry
                .get(&target_keyspace_id)
                .unwrap()
                .data_key_cache
                .len(),
            1,
            "preferred restore metadata should still use the manager data-key cache"
        );

        assert_eq!(restored_key.keyspace_id, Some(target_keyspace_id));
        assert_eq!(
            restored_key.core.plain_text,
            incoming_current_key.core.plain_text
        );

        let cached_meta = mgr.encryption_meta_by_keyspace(target_keyspace_id).unwrap();
        assert_eq!(cached_meta.get_keyspace_id(), source_keyspace_id);
        assert_eq!(cached_meta.get_current().get_file_id(), 100);
        assert_eq!(
            cached_meta.get_current().get_data_key_id(),
            incoming_current_key.get_key_id().unwrap(),
        );

        let target_key = target_pending_key
            .switch_if_header_mismatch_and_cache(target_current_key.encryption_header())
            .expect("initial key should remain resolvable after restore meta replacement");
        assert_eq!(
            mgr.keyspace_to_encryption_entry
                .get(&target_keyspace_id)
                .unwrap()
                .data_key_cache
                .len(),
            1,
            "historical retained metadata should not be inserted into the manager cache"
        );
        assert_eq!(target_key.keyspace_id, Some(target_keyspace_id));
        assert_eq!(
            target_key.core.plain_text,
            target_current_key.core.plain_text
        );

        let current_key = mgr
            .current_key_by_keyspace(target_keyspace_id)
            .expect("get key should succeed")
            .expect("keyspace should have encryption key");
        assert_eq!(
            current_key.core.plain_text,
            incoming_current_key.core.plain_text
        );
    }

    #[test]
    fn test_rotated_config_replaces_restored_cached_meta() {
        let keyspace_id = 55;
        let (mut initial_meta, _initial_key) = build_test_encryption_meta(keyspace_id);
        initial_meta.mut_current().set_file_id(100);

        let source_keyspace_id = 7;
        let (mut restore_meta, _restore_current_key) =
            build_test_encryption_meta(source_keyspace_id);
        restore_meta.mut_current().set_file_id(50);

        let (mut rotated_meta, rotated_key) = build_test_encryption_meta(keyspace_id);
        rotated_meta.mut_current().set_file_id(200);
        add_test_data_key(&mut rotated_meta);

        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &initial_meta.write_to_bytes().unwrap(),
            None,
        )
        .expect("failed to register initial meta");

        current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &restore_meta.write_to_bytes().unwrap(),
            Some(keyspace_id),
        )
        .expect("failed to register restore meta");

        let rotated_current_key = current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &rotated_meta.write_to_bytes().unwrap(),
            None,
        )
        .expect("failed to register rotated meta");
        assert_eq!(
            rotated_current_key.core.plain_text,
            rotated_key.core.plain_text
        );

        let cached_meta = mgr.encryption_meta_by_keyspace(keyspace_id).unwrap();
        assert_eq!(cached_meta.get_current().get_file_id(), 200);
        assert_eq!(
            cached_meta.get_current().get_data_key_id(),
            rotated_key.get_key_id().unwrap()
        );
    }

    #[test]
    fn test_cmek_health_check_detects_failed_master_key_decrypt() {
        let _guard = METRICS_TEST_LOCK
            .lock()
            .unwrap_or_else(|err| err.into_inner());
        reset_cmek_health_check_metric();

        let keyspace_id = 321;
        let source_keyspace_id = 654;
        let shard_id = 11;
        let mgr = Arc::new(EncryptionKeyManager::new_for_test());

        let (mut healthy_meta, _healthy_key) = build_test_encryption_meta(keyspace_id);
        healthy_meta.mut_current().set_file_id(100);
        let current_key = current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &healthy_meta.write_to_bytes().unwrap(),
            None,
        )
        .expect("failed to register healthy meta");
        mgr.register_current_key_for_shard(shard_id, current_key)
            .expect("failed to register active shard");

        new_test_runtime().block_on(EncryptionKeyManager::run_cmek_health_check_once(
            mgr.keyspace_active_shard_counts.as_ref(),
            mgr.keyspace_to_encryption_entry.as_ref(),
        ));
        assert_eq!(ENCRYPTION_KEYSPACE_CMEK_HEALTH_CHECK_FAILURE_GAUGE.get(), 0);

        let mut broken_restore_meta = healthy_meta.clone();
        broken_restore_meta.set_keyspace_id(source_keyspace_id);
        broken_restore_meta.mut_current().set_file_id(200);
        broken_restore_meta
            .mut_master_key()
            .set_vendor("aws".to_owned());
        encryption_state_from_meta_bytes(
            mgr.clone(),
            &broken_restore_meta.write_to_bytes().unwrap(),
            Some(keyspace_id),
        );

        new_test_runtime().block_on(EncryptionKeyManager::run_cmek_health_check_once(
            mgr.keyspace_active_shard_counts.as_ref(),
            mgr.keyspace_to_encryption_entry.as_ref(),
        ));
        assert_eq!(ENCRYPTION_KEYSPACE_CMEK_HEALTH_CHECK_FAILURE_GAUGE.get(), 1);
    }

    #[test]
    fn test_cmek_health_check_rejects_unsupported_vendor() {
        let _guard = METRICS_TEST_LOCK
            .lock()
            .unwrap_or_else(|err| err.into_inner());
        reset_cmek_health_check_metric();

        let keyspace_id = 324;
        let source_keyspace_id = 657;
        let shard_id = 13;
        let mgr = Arc::new(EncryptionKeyManager::new_for_test());

        let (mut healthy_meta, _healthy_key) = build_test_encryption_meta(keyspace_id);
        healthy_meta.mut_current().set_file_id(100);
        let current_key = current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &healthy_meta.write_to_bytes().unwrap(),
            None,
        )
        .expect("failed to register healthy meta");
        mgr.register_current_key_for_shard(shard_id, current_key)
            .expect("failed to register active shard");

        let mut broken_restore_meta = healthy_meta.clone();
        broken_restore_meta.set_keyspace_id(source_keyspace_id);
        broken_restore_meta.mut_current().set_file_id(200);
        broken_restore_meta
            .mut_master_key()
            .set_vendor("unknown-vendor".to_owned());
        encryption_state_from_meta_bytes(
            mgr.clone(),
            &broken_restore_meta.write_to_bytes().unwrap(),
            Some(keyspace_id),
        );

        new_test_runtime().block_on(EncryptionKeyManager::run_cmek_health_check_once(
            mgr.keyspace_active_shard_counts.as_ref(),
            mgr.keyspace_to_encryption_entry.as_ref(),
        ));
        assert_eq!(ENCRYPTION_KEYSPACE_CMEK_HEALTH_CHECK_FAILURE_GAUGE.get(), 1);
    }

    #[test]
    fn test_cmek_health_check_recovers_after_metadata_fix() {
        let _guard = METRICS_TEST_LOCK
            .lock()
            .unwrap_or_else(|err| err.into_inner());
        reset_cmek_health_check_metric();

        let keyspace_id = 322;
        let source_keyspace_id = 655;
        let shard_id = 12;
        let mgr = Arc::new(EncryptionKeyManager::new_for_test());

        let (mut healthy_meta, _healthy_key) = build_test_encryption_meta(keyspace_id);
        healthy_meta.mut_current().set_file_id(100);
        let current_key = current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &healthy_meta.write_to_bytes().unwrap(),
            None,
        )
        .expect("failed to register healthy meta");
        mgr.register_current_key_for_shard(shard_id, current_key)
            .expect("failed to register active shard");

        let mut broken_restore_meta = healthy_meta.clone();
        broken_restore_meta.set_keyspace_id(source_keyspace_id);
        broken_restore_meta.mut_current().set_file_id(200);
        broken_restore_meta
            .mut_master_key()
            .set_vendor("aws".to_owned());
        encryption_state_from_meta_bytes(
            mgr.clone(),
            &broken_restore_meta.write_to_bytes().unwrap(),
            Some(keyspace_id),
        );

        new_test_runtime().block_on(EncryptionKeyManager::run_cmek_health_check_once(
            mgr.keyspace_active_shard_counts.as_ref(),
            mgr.keyspace_to_encryption_entry.as_ref(),
        ));
        assert_eq!(ENCRYPTION_KEYSPACE_CMEK_HEALTH_CHECK_FAILURE_GAUGE.get(), 1);

        let (mut fixed_meta, fixed_key) = build_test_encryption_meta(keyspace_id);
        fixed_meta.mut_current().set_file_id(300);
        add_test_data_key(&mut fixed_meta);
        fixed_meta
            .mut_current()
            .set_data_key_id(fixed_key.get_key_id().unwrap());
        mgr.keyspace_to_encryption_entry
            .get(&keyspace_id)
            .unwrap()
            .update_cached_meta_if_preferred(&fixed_meta);

        new_test_runtime().block_on(EncryptionKeyManager::run_cmek_health_check_once(
            mgr.keyspace_active_shard_counts.as_ref(),
            mgr.keyspace_to_encryption_entry.as_ref(),
        ));
        assert_eq!(ENCRYPTION_KEYSPACE_CMEK_HEALTH_CHECK_FAILURE_GAUGE.get(), 0);
    }

    #[test]
    fn test_cmek_health_check_skips_inactive_registered_keyspace() {
        let _guard = METRICS_TEST_LOCK
            .lock()
            .unwrap_or_else(|err| err.into_inner());
        reset_cmek_health_check_metric();

        let keyspace_id = 323;
        let source_keyspace_id = 656;
        let mgr = Arc::new(EncryptionKeyManager::new_for_test());

        let (mut healthy_meta, _healthy_key) = build_test_encryption_meta(keyspace_id);
        healthy_meta.mut_current().set_file_id(100);
        current_data_key_from_encryption_meta_bytes(
            mgr.clone(),
            &healthy_meta.write_to_bytes().unwrap(),
            None,
        )
        .expect("failed to register healthy meta");

        let mut broken_restore_meta = healthy_meta.clone();
        broken_restore_meta.set_keyspace_id(source_keyspace_id);
        broken_restore_meta.mut_current().set_file_id(200);
        broken_restore_meta
            .mut_master_key()
            .set_vendor("aws".to_owned());
        encryption_state_from_meta_bytes(
            mgr.clone(),
            &broken_restore_meta.write_to_bytes().unwrap(),
            Some(keyspace_id),
        );

        new_test_runtime().block_on(EncryptionKeyManager::run_cmek_health_check_once(
            mgr.keyspace_active_shard_counts.as_ref(),
            mgr.keyspace_to_encryption_entry.as_ref(),
        ));
        assert_eq!(ENCRYPTION_KEYSPACE_CMEK_HEALTH_CHECK_FAILURE_GAUGE.get(), 0);
    }

    #[test]
    fn test_invalid_encryption_meta_bytes_return_error() {
        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        let err =
            current_data_key_from_encryption_meta_bytes(mgr.clone(), b"not-a-valid-meta", None)
                .expect_err("invalid bytes should return error");
        assert!(matches!(err, EncryptionKeyError::InvalidEncryptionMeta(_)));

        let key_state = encryption_state_from_meta_bytes(mgr, b"not-a-valid-meta", Some(42));
        assert_eq!(key_state.keyspace_id(), Some(42));
        assert!(matches!(
            key_state.resolve_ready_key(),
            Err(EncryptionKeyError::InvalidEncryptionState(_))
        ));
    }

    #[test]
    fn test_keyspace_encryption_config_pd_value_round_trip() {
        let encryption_context = BTreeMap::from([
            ("ExternalId".to_owned(), "customer-42".to_owned()),
            ("tenant".to_owned(), "acme".to_owned()),
        ]);
        let cfg = KeyspaceEncryptionConfig::enabled(
            Some("arn:aws:kms:us-west-2:123456:key/abc-def".to_owned()),
            Some("aws".to_owned()),
            Some("us-west-2".to_owned()),
            Some("https://kms.example".to_owned()),
            encryption_context.clone(),
        );

        let pd_value = cfg.to_pd_config_value().unwrap();
        let parsed = KeyspaceEncryptionConfig::from_pd_config_value(&pd_value).unwrap();

        assert!(parsed.enabled);
        assert_eq!(parsed.cmek_id, cfg.cmek_id);
        assert_eq!(parsed.vendor, cfg.vendor);
        assert_eq!(parsed.region, cfg.region);
        assert_eq!(parsed.endpoint, cfg.endpoint);
        assert_eq!(parsed.encryption_context, encryption_context);
    }

    #[test]
    fn test_drop_unused_encryption_key_manager_in_async_context() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        runtime.block_on(async {
            let mgr = EncryptionKeyManager::new_for_test();
            drop(mgr);
        });
    }

    #[test]
    fn test_drop_encryption_key_manager_with_cached_entry_in_async_context() {
        let mut config = MasterKeyConfig::default();
        config.vendor = "test".to_owned();
        let (master_key, master_key_ciphertext) = new_test_runtime()
            .block_on(config.generate_new_master_key())
            .unwrap();
        let data_key = master_key
            .generate_encryption_key_v2()
            .with_key_id(1)
            .unwrap();
        let mut meta = kvenginepb::EncryptionMeta::default();
        meta.set_keyspace_id(888);

        let mut mk = kvenginepb::MasterKey::default();
        mk.vendor = config.vendor.clone();
        mk.ciphertext = master_key_ciphertext.to_vec();
        meta.set_master_key(mk);

        meta.mut_current().data_key_id = data_key.get_key_id().unwrap();
        let mut dk = kvenginepb::DataKey::new();
        dk.set_ciphertext(data_key.export());
        meta.mut_data_keys()
            .insert(data_key.get_key_id().unwrap(), dk);

        let meta_bytes = meta.write_to_bytes().unwrap();
        let runtime = tokio::runtime::Runtime::new().unwrap();
        runtime.block_on(async {
            let mgr = Arc::new(EncryptionKeyManager::new_for_test());
            let key = current_data_key_from_encryption_meta_bytes(mgr.clone(), &meta_bytes, None)
                .expect("should register encryption meta");
            drop(key);
            drop(mgr);
        });
    }

    #[test]
    fn test_encryption_key_manager_rejects_missing_master_key() {
        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        let mut meta = kvenginepb::EncryptionMeta::default();
        meta.set_keyspace_id(42);
        meta.mut_current().data_key_id = 1;
        meta.mut_data_keys()
            .insert(1, kvenginepb::DataKey::default());
        let meta_bytes = meta.write_to_bytes().unwrap();

        let err = current_data_key_from_encryption_meta_bytes(mgr.clone(), &meta_bytes, None)
            .expect_err("missing master key should return error");
        assert!(matches!(err, EncryptionKeyError::InvalidEncryptionMeta(_)));

        let key_state = encryption_state_from_meta_bytes(mgr, &meta_bytes, None);
        assert_eq!(key_state.keyspace_id(), Some(42));
        assert!(matches!(
            key_state.resolve_ready_key(),
            Err(EncryptionKeyError::InvalidEncryptionState(_))
        ));
    }

    fn start_retrying_master_key_cache(
        keyspace_id: u32,
        data_key_id: u32,
    ) -> (
        Arc<EncryptionKeyManager>,
        MasterKeyConfig,
        EncryptionKey,
        kvenginepb::EncryptionMeta,
    ) {
        let mut source_config = MasterKeyConfig::default();
        source_config.vendor = "test".to_owned();
        let (source_master_key, source_master_key_ciphertext) = new_test_runtime()
            .block_on(source_config.generate_new_master_key())
            .unwrap();
        let source_master_key_ciphertext = source_master_key_ciphertext.to_vec();
        let data_key = source_master_key
            .generate_encryption_key_v2()
            .with_key_id(data_key_id)
            .unwrap();

        let mut meta = kvenginepb::EncryptionMeta::default();
        meta.set_keyspace_id(keyspace_id);
        let mut mk = kvenginepb::MasterKey::default();
        mk.vendor = "aws".to_owned();
        mk.ciphertext = source_master_key_ciphertext.clone();
        meta.set_master_key(mk);
        meta.mut_current().data_key_id = data_key_id;
        let mut dk = kvenginepb::DataKey::default();
        dk.set_ciphertext(data_key.export());
        meta.mut_data_keys().insert(data_key_id, dk);
        let cache_key = KeyspaceEncryptionEntry::build_master_key_config(&meta);
        let meta_bytes = meta.write_to_bytes().unwrap();

        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        let err = encryption_state_from_meta_bytes(mgr.clone(), &meta_bytes, None)
            .resolve_ready_key()
            .expect_err("ready key resolution should fail when KMS is unavailable");
        assert!(err.is_retryable());
        assert!(
            mgr.keyspace_to_encryption_entry
                .get(&keyspace_id)
                .unwrap()
                .master_key_cache
                .get(&cache_key)
                .is_some()
        );

        (mgr, cache_key, data_key, meta)
    }

    #[test]
    fn test_encryption_key_manager_returns_pending_state_on_kms_failure() {
        let mut source_config = MasterKeyConfig::default();
        source_config.vendor = "test".to_owned();
        let (source_master_key, source_master_key_ciphertext) = new_test_runtime()
            .block_on(source_config.generate_new_master_key())
            .unwrap();
        let source_master_key_ciphertext = source_master_key_ciphertext.to_vec();
        let data_key_id = 7;
        let data_key = source_master_key
            .generate_encryption_key_v2()
            .with_key_id(data_key_id)
            .unwrap();

        let keyspace_id = 42;
        let mut meta = kvenginepb::EncryptionMeta::default();
        meta.set_keyspace_id(keyspace_id);
        let mut mk = kvenginepb::MasterKey::default();
        mk.vendor = "aws".to_owned();
        mk.ciphertext = source_master_key_ciphertext.clone();
        meta.set_master_key(mk);
        meta.mut_current().data_key_id = data_key_id;
        let mut dk = kvenginepb::DataKey::default();
        dk.set_ciphertext(data_key.export());
        meta.mut_data_keys().insert(data_key_id, dk);
        let meta_bytes = meta.write_to_bytes().unwrap();
        let source_master_key_config = KeyspaceEncryptionEntry::build_master_key_config(&meta);

        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        let key_state = encryption_state_from_meta_bytes(mgr.clone(), &meta_bytes, None);
        assert_eq!(key_state.data_key_id(), Some(data_key_id));
        assert_eq!(key_state.keyspace_id(), Some(keyspace_id));
        assert_eq!(
            key_state.encryption_header(),
            Some(data_key.encryption_header())
        );
        let err = key_state
            .resolve_ready_key()
            .expect_err("ready key resolution should fail when KMS is unavailable");
        assert!(err.is_retryable());

        let entry = mgr.keyspace_to_encryption_entry.get(&keyspace_id).unwrap();
        let master_key_cache = entry
            .master_key_cache
            .get(&source_master_key_config)
            .unwrap()
            .clone();
        {
            let cached = master_key_cache.lock().unwrap();
            assert!(cached.master_key.is_none());
            assert!(cached.last_err.is_some());
            assert_eq!(cached.retry_state, MasterKeyRetryState::Retrying);
        }

        let err = key_state
            .resolve_ready_key()
            .expect_err("foreground path should not retry while background retry is running");
        assert!(err.is_retryable());
    }

    #[test]
    fn test_master_key_retry_cache_is_removed_on_new_master_key_config() {
        let _guard = METRICS_TEST_LOCK
            .lock()
            .unwrap_or_else(|err| err.into_inner());
        let keyspace_id = 43;
        let data_key_id = 8;
        let shard_id = 11;
        let (mgr, retry_cache_key, data_key, mut new_meta) =
            start_retrying_master_key_cache(keyspace_id, data_key_id);
        mgr.register_current_key_for_shard(shard_id, data_key.with_keyspace_id(keyspace_id))
            .expect("failed to register active shard");

        let entry = mgr.keyspace_to_encryption_entry.get(&keyspace_id).unwrap();
        new_meta.mut_current().set_file_id(1);
        set_master_key_encryption_context(
            new_meta.mut_master_key(),
            &BTreeMap::from([("ExternalId".to_owned(), "new-context".to_owned())]),
        );
        new_meta
            .mut_data_keys()
            .insert(data_key_id + 1, kvenginepb::DataKey::default());
        entry.update_cached_meta_if_preferred(&new_meta);

        assert!(mgr.keyspace_active_shard_counts.contains_key(&keyspace_id));
        assert!(entry.master_key_cache.get(&retry_cache_key).is_none());
    }

    #[test]
    fn test_master_key_retry_cache_is_removed_when_keyspace_becomes_inactive() {
        let _guard = METRICS_TEST_LOCK
            .lock()
            .unwrap_or_else(|err| err.into_inner());
        let keyspace_id = 44;
        let data_key_id = 9;
        let shard_id = 12;
        let (mgr, retry_cache_key, data_key, _) =
            start_retrying_master_key_cache(keyspace_id, data_key_id);

        mgr.register_current_key_for_shard(shard_id, data_key.with_keyspace_id(keyspace_id))
            .expect("failed to register active shard");
        mgr.deregister_current_key_for_shard(shard_id);

        let entry = mgr.keyspace_to_encryption_entry.get(&keyspace_id).unwrap();
        assert!(entry.master_key_cache.get(&retry_cache_key).is_none());
    }

    #[test]
    fn test_encryption_key_manager_uses_meta_data_key_id_for_rewrapped_key() {
        let original_master_key_plain_text = rand::thread_rng().gen::<[u8; 32]>().to_vec();
        let original_master_key = MasterKey::new(&original_master_key_plain_text);

        let mut rewrap_config = MasterKeyConfig::default();
        rewrap_config.vendor = "test".to_owned();
        rewrap_config.key_id = "rewrapped-master-key".to_owned();
        let (rewrap_master_key, rewrap_master_key_ciphertext) = new_test_runtime()
            .block_on(rewrap_config.generate_new_master_key())
            .unwrap();

        let stable_data_key_id = 0x123456;
        let data_key = original_master_key
            .generate_encryption_key_v2()
            .with_key_id(stable_data_key_id)
            .unwrap();
        let rewrapped_cipher_text = symm::encrypt(
            Cipher::aes_256_ctr(),
            &rewrap_master_key.export(),
            None,
            &data_key.plain_text,
        )
        .unwrap();
        let rewrapped_data_key_without_id = EncryptionKey::new(
            rewrapped_cipher_text,
            data_key.plain_text.clone(),
            KEY_TYPE_AES_256_CTR,
        );
        let panic = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            rewrapped_data_key_without_id.encryption_header()
        }));
        panic.expect_err("rewrapped v2 key should still need an externally bound data_key_id");

        let mut meta = kvenginepb::EncryptionMeta::default();
        meta.set_keyspace_id(7);
        let mut mk = kvenginepb::MasterKey::default();
        mk.vendor = rewrap_config.vendor.clone();
        mk.cmek_id = rewrap_config.key_id.clone();
        mk.ciphertext = rewrap_master_key_ciphertext.to_vec();
        meta.set_master_key(mk);

        let rewrapped_data_key = rewrapped_data_key_without_id
            .with_key_id(stable_data_key_id)
            .unwrap();
        meta.mut_current().data_key_id = stable_data_key_id;
        let mut dk = kvenginepb::DataKey::default();
        dk.set_ciphertext(rewrapped_data_key.export());
        meta.mut_data_keys().insert(stable_data_key_id, dk);

        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        let key =
            current_data_key_from_encryption_meta_bytes(mgr, &meta.write_to_bytes().unwrap(), None)
                .expect("metadata data_key_id should be the source of truth");
        assert_eq!(key.get_key_id().unwrap(), stable_data_key_id);
        assert_eq!(key.core.plain_text, data_key.core.plain_text);
        assert_eq!(
            key.encryption_header(),
            rewrapped_data_key.encryption_header()
        );
    }

    #[test]
    fn test_switch_if_header_mismatch_requires_key_manager() {
        let master_key_plain_text = rand::thread_rng().gen::<[u8; 32]>().to_vec();
        let master_key = MasterKey::new(&master_key_plain_text);
        let encryption_key = master_key
            .generate_encryption_key_v2()
            .with_key_id(1)
            .unwrap()
            .with_keyspace_id(1);
        let err = encryption_key
            .switch_if_header_mismatch(0x01000002)
            .expect_err("missing key manager should return error");
        assert!(matches!(err, EncryptionKeyError::InvalidEncryptionState(_)));
    }

    #[test]
    fn test_register_current_key_for_shard_requires_keyspace_id() {
        let master_key_plain_text = rand::thread_rng().gen::<[u8; 32]>().to_vec();
        let master_key = MasterKey::new(&master_key_plain_text);
        let encryption_key = master_key.generate_encryption_key_v2();
        let mgr = EncryptionKeyManager::new_for_test();

        let err = mgr
            .register_current_key_for_shard(9, encryption_key)
            .expect_err("missing keyspace id should return error");
        assert!(matches!(err, EncryptionKeyError::InvalidEncryptionState(_)));
    }

    #[test]
    fn test_encryption_key_manager_prefers_metadata_mode() {
        let mgr = EncryptionKeyManager::new_for_test();
        mgr.set_mode_if_unset(EncryptionMethod::Legacy);
        mgr.set_mode_if_unset(EncryptionMethod::Metadata);
        assert_eq!(mgr.current_mode(), Some(EncryptionMethod::Metadata));

        let mgr = EncryptionKeyManager::new_for_test();
        mgr.set_mode_if_unset(EncryptionMethod::Metadata);
        mgr.set_mode_if_unset(EncryptionMethod::Legacy);
        assert_eq!(mgr.current_mode(), Some(EncryptionMethod::Metadata));
    }

    #[test]
    fn test_keyspace_encryption_config_intended_method() {
        assert_eq!(
            KeyspaceEncryptionConfig::default()
                .intended_encryption_method()
                .unwrap(),
            None
        );

        assert_eq!(
            KeyspaceEncryptionConfig {
                enabled: true,
                ..Default::default()
            }
            .intended_encryption_method()
            .unwrap(),
            Some(EncryptionMethod::Legacy)
        );

        assert_eq!(
            KeyspaceEncryptionConfig {
                enabled: true,
                cmek_id: Some("cmek-1".into()),
                vendor: Some("aws".into()),
                ..Default::default()
            }
            .intended_encryption_method()
            .unwrap(),
            Some(EncryptionMethod::Metadata)
        );

        let err = KeyspaceEncryptionConfig {
            enabled: true,
            vendor: Some("aws".into()),
            ..Default::default()
        }
        .intended_encryption_method()
        .expect_err("partial CMEK intent should be rejected");
        assert!(matches!(err, EncryptionKeyError::InvalidEncryptionState(_)));
    }

    #[test]
    fn test_master_key_non_test_vendor_state() {
        let aws_master_key = MasterKey::new_with_kms_vendor(&[1; 32], KmsVendor::Aws);
        assert!(aws_master_key.uses_real_kms_vendor());

        let test_master_key = MasterKey::new_with_kms_vendor(&[2; 32], KmsVendor::Test);
        assert!(!test_master_key.uses_real_kms_vendor());

        let plain_master_key = MasterKey::new(&[3; 32]);
        assert!(!plain_master_key.uses_real_kms_vendor());
    }

    #[test]
    fn test_encryption_key_manager_try_set_legacy_mode_from_global_master_key() {
        let mgr = EncryptionKeyManager::new_for_test();
        let aws_master_key = MasterKey::new_with_kms_vendor(&[1; 32], KmsVendor::Aws);
        mgr.try_set_legacy_mode_from_global_master_key(&aws_master_key);
        assert_eq!(mgr.current_mode(), Some(EncryptionMethod::Legacy));

        let mgr = EncryptionKeyManager::new_for_test();
        let test_master_key = MasterKey::new_with_kms_vendor(&[2; 32], KmsVendor::Test);
        mgr.try_set_legacy_mode_from_global_master_key(&test_master_key);
        assert_eq!(mgr.current_mode(), None);

        let mgr = EncryptionKeyManager::new_for_test();
        mgr.set_mode_if_unset(EncryptionMethod::Metadata);
        mgr.try_set_legacy_mode_from_global_master_key(&aws_master_key);
        assert_eq!(mgr.current_mode(), Some(EncryptionMethod::Metadata));
    }

    #[test]
    fn test_unencrypted_key() {
        let encryption_key = EncryptionKey::new_unencrypted_key();
        let mut encrypted = vec![];
        encryption_key.encrypt(b"hello", 123, 12, &mut encrypted);
        assert_eq!(encrypted, b"hello");

        let mut decrypted = vec![];
        encryption_key.decrypt(&encrypted, 123, 12, &mut decrypted);
        assert_eq!(decrypted, b"hello");
    }

    #[test]
    fn test_get_or_decrypt_master_key() {
        // Setup: generate a test master key and encryption meta
        let mut config = MasterKeyConfig::default();
        config.vendor = "test".to_owned();
        let rt = new_test_runtime();
        let (master_key, master_key_ciphertext) =
            rt.block_on(config.generate_new_master_key()).unwrap();

        let keyspace_id = 200;
        let mut m = kvenginepb::EncryptionMeta::default();
        m.set_keyspace_id(keyspace_id);

        let mut mk = kvenginepb::MasterKey::default();
        mk.vendor = config.vendor.clone();
        mk.ciphertext = master_key_ciphertext.to_vec();
        m.set_master_key(mk);

        let data_key_id = generate_data_key_id();
        let data_key = master_key
            .generate_encryption_key_v2()
            .with_key_id(data_key_id)
            .unwrap();
        m.mut_current().data_key_id = data_key_id;
        let mut dk = kvenginepb::DataKey::new();
        dk.set_ciphertext(data_key.export());
        m.mut_data_keys().insert(data_key_id, dk);

        // Test: get_or_decrypt_master_key should return the correct master key
        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        let decrypted_master_key = mgr.get_or_decrypt_master_key(&m).unwrap();

        // Verify: the decrypted master key should be able to generate valid data keys
        let new_data_key = decrypted_master_key.generate_encryption_key();
        let mut test_data = vec![];
        new_data_key.encrypt(b"test data", 456, 78, &mut test_data);

        // Re-decrypt the new data key using the master key export and verify
        let reloaded_key = decrypted_master_key
            .decrypt_encryption_key(&new_data_key.export())
            .unwrap();
        let mut decrypted_data = vec![];
        reloaded_key.decrypt(&test_data, 456, 78, &mut decrypted_data);
        assert_eq!(decrypted_data, b"test data");

        // Verify: calling get_or_decrypt_master_key again returns cached entry
        let decrypted_master_key2 = mgr.get_or_decrypt_master_key(&m).unwrap();
        assert_eq!(
            decrypted_master_key.export(),
            decrypted_master_key2.export()
        );
    }

    #[test]
    fn test_first_master_key_decrypt_is_single_flight() {
        let key_id = "single-flight-master-key";
        let keyspace_id = 201;
        let mut meta = kvenginepb::EncryptionMeta::default();
        meta.set_keyspace_id(keyspace_id);
        meta.mut_master_key().vendor = "test".to_owned();
        meta.mut_master_key().cmek_id = key_id.to_owned();

        let mgr = EncryptionKeyManager::new_for_test();
        let barrier = Barrier::new(17);
        let keys = std::thread::scope(|scope| {
            let handles = (0..16)
                .map(|_| {
                    scope.spawn(|| {
                        barrier.wait();
                        mgr.get_or_decrypt_master_key(&meta)
                            .map(|master_key| master_key.export())
                            .unwrap()
                    })
                })
                .collect::<Vec<_>>();
            barrier.wait();
            handles
                .into_iter()
                .map(|handle| handle.join().unwrap())
                .collect::<Vec<_>>()
        });

        let mut expected = key_id.as_bytes().to_vec();
        expected.resize(32, 11);
        for key in keys {
            assert_eq!(&key[..], expected.as_slice());
        }
    }

    #[test]
    fn test_encryption_meta_roundtrip() {
        // Setup: create a complete encryption meta with master key and data keys
        let mut config = MasterKeyConfig::default();
        config.vendor = "test".to_owned();
        let rt = new_test_runtime();
        let (master_key, master_key_ciphertext) =
            rt.block_on(config.generate_new_master_key()).unwrap();

        let keyspace_id = 300;
        let mut original_meta = kvenginepb::EncryptionMeta::default();
        original_meta.set_keyspace_id(keyspace_id);

        let mut mk = kvenginepb::MasterKey::default();
        mk.vendor = config.vendor.clone();
        mk.region = "us-west-2".to_owned();
        mk.ciphertext = master_key_ciphertext.to_vec();
        original_meta.set_master_key(mk);

        // Add multiple data keys.
        let data_key1_id = 1;
        let data_key2_id = 2;
        let data_key1 = master_key
            .generate_encryption_key_v2()
            .with_key_id(data_key1_id)
            .unwrap();
        let data_key2 = master_key
            .generate_encryption_key_v2()
            .with_key_id(data_key2_id)
            .unwrap();

        for (key_id, key) in [(data_key1_id, &data_key1), (data_key2_id, &data_key2)] {
            let mut dk = kvenginepb::DataKey::new();
            dk.set_ciphertext(key.export());
            original_meta.mut_data_keys().insert(key_id, dk);
        }

        // Set current epoch
        let mut current_epoch = kvenginepb::EncryptionEpoch::default();
        current_epoch.set_file_id(12345);
        current_epoch.set_data_key_id(data_key2_id);
        current_epoch.set_created_at(1000);
        original_meta.set_current(current_epoch);

        // Serialize to bytes
        let meta_bytes = original_meta.write_to_bytes().unwrap();

        // Deserialize and verify all fields match
        let mut loaded_meta = kvenginepb::EncryptionMeta::default();
        loaded_meta.merge_from_bytes(&meta_bytes).unwrap();

        assert_eq!(loaded_meta.keyspace_id, original_meta.keyspace_id);
        assert_eq!(
            loaded_meta.get_master_key().vendor,
            original_meta.get_master_key().vendor
        );
        assert_eq!(
            loaded_meta.get_master_key().region,
            original_meta.get_master_key().region
        );
        assert_eq!(
            loaded_meta.get_master_key().ciphertext,
            original_meta.get_master_key().ciphertext
        );
        assert_eq!(
            loaded_meta.get_current().file_id,
            original_meta.get_current().file_id
        );
        assert_eq!(
            loaded_meta.get_current().data_key_id,
            original_meta.get_current().data_key_id
        );
        assert_eq!(loaded_meta.data_keys.len(), original_meta.data_keys.len());

        // Verify data keys can be decrypted correctly after roundtrip
        let mgr = Arc::new(EncryptionKeyManager::new_for_test());
        let decrypted_master_key = mgr.get_or_decrypt_master_key(&loaded_meta).unwrap();

        for (key_id, dk) in loaded_meta.get_data_keys() {
            let loaded_data_key = decrypted_master_key
                .decrypt_encryption_key(&dk.ciphertext)
                .and_then(|key| key.with_key_id(*key_id))
                .unwrap();
            assert_eq!(loaded_data_key.get_key_id().unwrap(), *key_id);

            // Verify encryption/decryption works
            let mut encrypted = vec![];
            loaded_data_key.encrypt(b"roundtrip test", 999, 1, &mut encrypted);
            let mut decrypted = vec![];
            loaded_data_key.decrypt(&encrypted, 999, 1, &mut decrypted);
            assert_eq!(decrypted, b"roundtrip test");
        }
    }
}