use rand::RngCore;
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use zeroize::Zeroizing;
use crate::keys::{decrypt_bytes, encrypt_bytes, CryptoError, EncryptedKey};
use koi_common::encoding::{hex_decode, hex_encode};
const MASTER_KEY_LEN: usize = 32;
const TOTP_SLOT_HKDF_INFO: &[u8] = b"pond-unlock-slot-totp-v1";
const TOTP_CREDENTIAL_LABEL: &str = "koi-certmesh-unlock-totp";
#[derive(Debug, Serialize, Deserialize)]
pub struct SlotTable {
pub version: u32,
pub slots: Vec<UnlockSlot>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum UnlockSlot {
#[serde(rename = "passphrase")]
Passphrase {
wrapped_master_key: EncryptedKey,
},
#[serde(rename = "auto_unlock")]
AutoUnlock,
#[serde(rename = "totp")]
Totp {
#[serde(default)]
sealed: bool,
#[serde(skip_serializing_if = "Option::is_none")]
shared_secret_hex: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
encrypted_secret: Option<EncryptedKey>,
wrapped_master_key: EncryptedKey,
},
#[serde(rename = "fido2")]
Fido2 {
credential_id: String,
public_key: String,
rp_id: String,
sign_count: u32,
wrapped_master_key: EncryptedKey,
encrypted_slot_kek: EncryptedKey,
},
}
impl SlotTable {
pub fn new_with_passphrase(
master_key: &[u8; MASTER_KEY_LEN],
passphrase: &str,
) -> Result<Self, CryptoError> {
let wrapped = encrypt_bytes(master_key, passphrase)?;
Ok(Self {
version: 1,
slots: vec![UnlockSlot::Passphrase {
wrapped_master_key: wrapped,
}],
})
}
pub fn unwrap_with_passphrase(
&self,
passphrase: &str,
) -> Result<Zeroizing<[u8; MASTER_KEY_LEN]>, CryptoError> {
for slot in &self.slots {
if let UnlockSlot::Passphrase {
wrapped_master_key, ..
} = slot
{
let bytes = decrypt_bytes(wrapped_master_key, passphrase)?;
return bytes_to_master_key(&bytes);
}
}
Err(CryptoError::Decryption("no passphrase slot found".into()))
}
pub fn add_auto_unlock(&mut self) {
self.slots.retain(|s| !matches!(s, UnlockSlot::AutoUnlock));
self.slots.push(UnlockSlot::AutoUnlock);
}
pub fn remove_auto_unlock(&mut self) {
self.slots.retain(|s| !matches!(s, UnlockSlot::AutoUnlock));
}
pub fn has_auto_unlock(&self) -> bool {
self.slots
.iter()
.any(|s| matches!(s, UnlockSlot::AutoUnlock))
}
pub fn add_totp_slot(
&mut self,
master_key: &[u8; MASTER_KEY_LEN],
shared_secret: &[u8],
) -> Result<(), CryptoError> {
self.slots.retain(|s| !matches!(s, UnlockSlot::Totp { .. }));
let slot_kek = derive_totp_slot_kek(shared_secret);
let slot_kek_hex = Zeroizing::new(hex_encode(&*slot_kek));
let wrapped = encrypt_bytes(master_key, &slot_kek_hex)?;
let (sealed, encrypted_secret) =
match crate::tpm::seal_key_material(TOTP_CREDENTIAL_LABEL, shared_secret) {
Ok(()) => {
tracing::info!("TOTP shared secret sealed in platform credential store");
(true, None)
}
Err(_) => {
let fallback_key = get_or_create_fallback_key()?;
let fallback_hex = Zeroizing::new(hex_encode(&*fallback_key));
let enc = encrypt_bytes(shared_secret, &fallback_hex)?;
tracing::info!("TOTP shared secret encrypted with sealed fallback key");
(false, Some(enc))
}
};
self.slots.push(UnlockSlot::Totp {
sealed,
shared_secret_hex: None,
encrypted_secret,
wrapped_master_key: wrapped,
});
Ok(())
}
pub fn unwrap_with_totp(
&self,
code: &str,
) -> Result<Zeroizing<[u8; MASTER_KEY_LEN]>, CryptoError> {
for slot in &self.slots {
if let UnlockSlot::Totp {
sealed,
shared_secret_hex,
encrypted_secret,
wrapped_master_key,
} = slot
{
let secret_bytes = Zeroizing::new(if *sealed {
crate::tpm::unseal_key_material(TOTP_CREDENTIAL_LABEL).map_err(|e| {
CryptoError::Decryption(format!(
"failed to unseal TOTP secret from platform store: {e}"
))
})?
} else if let Some(enc) = encrypted_secret {
let fallback_key = get_or_create_fallback_key().map_err(|e| {
CryptoError::Decryption(format!(
"failed to retrieve TOTP fallback key: {e}"
))
})?;
let fallback_hex = Zeroizing::new(hex_encode(&*fallback_key));
decrypt_bytes(enc, &fallback_hex).map_err(|e| {
CryptoError::Decryption(format!(
"failed to decrypt TOTP secret with fallback key: {e}"
))
})?
} else if let Some(hex) = shared_secret_hex {
tracing::warn!(
"TOTP secret stored in plaintext (legacy format). \
Re-create the CA or rotate auth to migrate to encrypted storage."
);
hex_decode(hex).map_err(|e| {
CryptoError::Decryption(format!("invalid TOTP secret hex: {e}"))
})?
} else {
return Err(CryptoError::Decryption(
"TOTP slot has no recoverable secret".into(),
));
});
let secret = crate::totp::TotpSecret::from_bytes(secret_bytes.to_vec());
if !crate::totp::verify_code(&secret, code) {
return Err(CryptoError::Decryption("invalid TOTP code".into()));
}
let slot_kek = derive_totp_slot_kek(&secret_bytes);
drop(secret_bytes);
let slot_kek_hex = Zeroizing::new(hex_encode(&*slot_kek));
let bytes = decrypt_bytes(wrapped_master_key, &slot_kek_hex)?;
return bytes_to_master_key(&bytes);
}
}
Err(CryptoError::Decryption("no TOTP slot found".into()))
}
pub fn has_totp_slot(&self) -> bool {
self.slots
.iter()
.any(|s| matches!(s, UnlockSlot::Totp { .. }))
}
pub fn add_fido2_slot(
&mut self,
master_key: &[u8; MASTER_KEY_LEN],
credential_id: &[u8],
public_key: &[u8],
rp_id: &str,
) -> Result<(), CryptoError> {
use base64::Engine;
let b64 = base64::engine::general_purpose::STANDARD;
self.slots
.retain(|s| !matches!(s, UnlockSlot::Fido2 { .. }));
let mut slot_kek = Zeroizing::new([0u8; 32]);
rand::rng().fill_bytes(slot_kek.as_mut());
let slot_kek_hex = Zeroizing::new(hex_encode(slot_kek.as_ref()));
let wrapped = encrypt_bytes(master_key, &slot_kek_hex)?;
let cred_derived_key = derive_fido2_storage_key(credential_id);
let cred_derived_hex = Zeroizing::new(hex_encode(&*cred_derived_key));
let encrypted_slot_kek = encrypt_bytes(&*slot_kek, &cred_derived_hex)?;
self.slots.push(UnlockSlot::Fido2 {
credential_id: b64.encode(credential_id),
public_key: b64.encode(public_key),
rp_id: rp_id.to_string(),
sign_count: 0,
wrapped_master_key: wrapped,
encrypted_slot_kek,
});
Ok(())
}
pub fn unwrap_with_fido2(
&self,
credential_id: &[u8],
) -> Result<Zeroizing<[u8; MASTER_KEY_LEN]>, CryptoError> {
use base64::Engine;
let b64 = base64::engine::general_purpose::STANDARD;
let target_id = b64.encode(credential_id);
for slot in &self.slots {
if let UnlockSlot::Fido2 {
credential_id: stored_id,
encrypted_slot_kek,
wrapped_master_key,
..
} = slot
{
if stored_id == &target_id {
let cred_derived_key = derive_fido2_storage_key(credential_id);
let cred_derived_hex = Zeroizing::new(hex_encode(&*cred_derived_key));
let slot_kek =
Zeroizing::new(decrypt_bytes(encrypted_slot_kek, &cred_derived_hex)?);
let slot_kek_hex = Zeroizing::new(hex_encode(&slot_kek));
let bytes = decrypt_bytes(wrapped_master_key, &slot_kek_hex)?;
return bytes_to_master_key(&bytes);
}
}
}
Err(CryptoError::Decryption(
"no matching FIDO2 slot found".into(),
))
}
pub fn fido2_credential(&self) -> Option<Fido2SlotInfo> {
use base64::Engine;
let b64 = base64::engine::general_purpose::STANDARD;
for slot in &self.slots {
if let UnlockSlot::Fido2 {
credential_id,
public_key,
rp_id,
sign_count,
..
} = slot
{
let cred_bytes = match b64.decode(credential_id) {
Ok(bytes) => bytes,
Err(e) => {
tracing::warn!(error = %e, "FIDO2 credential_id base64 decode failed");
return None;
}
};
let pk_bytes = match b64.decode(public_key) {
Ok(bytes) => bytes,
Err(e) => {
tracing::warn!(error = %e, "FIDO2 public_key base64 decode failed");
return None;
}
};
return Some(Fido2SlotInfo {
credential_id: cred_bytes,
public_key: pk_bytes,
rp_id: rp_id.clone(),
sign_count: *sign_count,
});
}
}
None
}
pub fn has_fido2_slot(&self) -> bool {
self.slots
.iter()
.any(|s| matches!(s, UnlockSlot::Fido2 { .. }))
}
pub fn update_fido2_sign_count(&mut self, credential_id: &[u8], new_count: u32) {
use base64::Engine;
let b64 = base64::engine::general_purpose::STANDARD;
let target_id = b64.encode(credential_id);
for slot in &mut self.slots {
if let UnlockSlot::Fido2 {
credential_id: stored_id,
sign_count,
..
} = slot
{
if stored_id == &target_id {
*sign_count = new_count;
}
}
}
}
pub fn available_methods(&self) -> Vec<&'static str> {
let mut methods = Vec::new();
for slot in &self.slots {
match slot {
UnlockSlot::Passphrase { .. } => methods.push("passphrase"),
UnlockSlot::AutoUnlock => methods.push("auto_unlock"),
UnlockSlot::Totp { .. } => methods.push("totp"),
UnlockSlot::Fido2 { .. } => methods.push("fido2"),
}
}
methods
}
pub fn save(&self, path: &std::path::Path) -> Result<(), CryptoError> {
let json = serde_json::to_string_pretty(self)
.map_err(|e| CryptoError::Serialization(e.to_string()))?;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
crate::keys::write_secret_file(path, json.as_bytes())?;
tracing::debug!(path = %path.display(), "Slot table saved");
Ok(())
}
pub fn load(path: &std::path::Path) -> Result<Self, CryptoError> {
let json = std::fs::read_to_string(path)?;
let table: Self =
serde_json::from_str(&json).map_err(|e| CryptoError::Serialization(e.to_string()))?;
Ok(table)
}
}
#[derive(Debug, Clone)]
pub struct Fido2SlotInfo {
pub credential_id: Vec<u8>,
pub public_key: Vec<u8>,
pub rp_id: String,
pub sign_count: u32,
}
pub fn generate_master_key() -> Zeroizing<[u8; MASTER_KEY_LEN]> {
let mut key = Zeroizing::new([0u8; MASTER_KEY_LEN]);
rand::rng().fill_bytes(key.as_mut());
key
}
fn derive_totp_slot_kek(shared_secret: &[u8]) -> Zeroizing<[u8; 32]> {
let mut hasher = Sha256::new();
hasher.update(shared_secret);
hasher.update(TOTP_SLOT_HKDF_INFO);
let result = hasher.finalize();
let mut kek = Zeroizing::new([0u8; 32]);
kek.copy_from_slice(&result);
kek
}
const TOTP_FALLBACK_KEY_LABEL: &str = "koi-certmesh-totp-fallback-key";
fn get_or_create_fallback_key() -> Result<Zeroizing<[u8; 32]>, CryptoError> {
if let Ok(bytes) = crate::tpm::unseal_key_material(TOTP_FALLBACK_KEY_LABEL) {
if bytes.len() == 32 {
let mut key = Zeroizing::new([0u8; 32]);
key.copy_from_slice(&bytes);
return Ok(key);
}
}
let mut key = Zeroizing::new([0u8; 32]);
rand::rng().fill_bytes(key.as_mut());
crate::tpm::seal_key_material(TOTP_FALLBACK_KEY_LABEL, &*key).map_err(|e| {
CryptoError::Encryption(format!(
"cannot seal TOTP fallback key in platform credential store: {e}"
))
})?;
let confirmed = crate::tpm::unseal_key_material(TOTP_FALLBACK_KEY_LABEL)
.map_err(|e| CryptoError::Encryption(format!("cannot confirm TOTP fallback key: {e}")))?;
if confirmed.len() == 32 {
let mut k = Zeroizing::new([0u8; 32]);
k.copy_from_slice(&confirmed);
Ok(k)
} else {
Ok(key)
}
}
fn derive_fido2_storage_key(credential_id: &[u8]) -> Zeroizing<[u8; 32]> {
let mut hasher = Sha256::new();
hasher.update(b"pond-fido2-storage-key-v1");
hasher.update(credential_id);
let result = hasher.finalize();
let mut key = Zeroizing::new([0u8; 32]);
key.copy_from_slice(&result);
key
}
fn bytes_to_master_key(bytes: &[u8]) -> Result<Zeroizing<[u8; MASTER_KEY_LEN]>, CryptoError> {
if bytes.len() != MASTER_KEY_LEN {
return Err(CryptoError::Decryption(format!(
"master key has wrong length: expected {MASTER_KEY_LEN}, got {}",
bytes.len()
)));
}
let mut key = Zeroizing::new([0u8; MASTER_KEY_LEN]);
key.copy_from_slice(bytes);
Ok(key)
}
pub fn migrate_to_envelope(
old_encrypted: &EncryptedKey,
passphrase: &str,
) -> Result<(EncryptedKey, SlotTable, Zeroizing<[u8; MASTER_KEY_LEN]>), CryptoError> {
let plaintext = decrypt_bytes(old_encrypted, passphrase)?;
let master_key = generate_master_key();
let master_key_hex = Zeroizing::new(hex_encode(master_key.as_ref()));
let new_encrypted = encrypt_bytes(&plaintext, &master_key_hex)?;
let slot_table = SlotTable::new_with_passphrase(&master_key, passphrase)?;
Ok((new_encrypted, slot_table, master_key))
}
pub fn envelope_encrypt_new(
ca_key_der: &[u8],
passphrase: &str,
) -> Result<(EncryptedKey, SlotTable, Zeroizing<[u8; MASTER_KEY_LEN]>), CryptoError> {
let master_key = generate_master_key();
let master_key_hex = Zeroizing::new(hex_encode(master_key.as_ref()));
let encrypted = encrypt_bytes(ca_key_der, &master_key_hex)?;
let slot_table = SlotTable::new_with_passphrase(&master_key, passphrase)?;
Ok((encrypted, slot_table, master_key))
}
pub fn decrypt_with_master_key(
encrypted: &EncryptedKey,
master_key: &[u8; MASTER_KEY_LEN],
) -> Result<Vec<u8>, CryptoError> {
let master_key_hex = Zeroizing::new(hex_encode(master_key));
decrypt_bytes(encrypted, &master_key_hex)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn passphrase_slot_round_trip() {
let master_key = generate_master_key();
let table = SlotTable::new_with_passphrase(&master_key, "test-pass").unwrap();
let recovered = table.unwrap_with_passphrase("test-pass").unwrap();
assert_eq!(master_key, recovered);
}
#[test]
fn wrong_passphrase_fails() {
let master_key = generate_master_key();
let table = SlotTable::new_with_passphrase(&master_key, "correct").unwrap();
assert!(table.unwrap_with_passphrase("wrong").is_err());
}
#[cfg(feature = "keyring")]
#[test]
fn totp_slot_round_trip() {
let master_key = generate_master_key();
let mut table = SlotTable::new_with_passphrase(&master_key, "pass").unwrap();
let secret = crate::totp::generate_secret();
table.add_totp_slot(&master_key, secret.as_bytes()).unwrap();
let code = crate::totp::current_code(&secret).unwrap();
let recovered = table.unwrap_with_totp(&code).unwrap();
assert_eq!(master_key, recovered);
}
#[cfg(feature = "keyring")]
#[test]
fn totp_wrong_code_fails() {
let master_key = generate_master_key();
let mut table = SlotTable::new_with_passphrase(&master_key, "pass").unwrap();
let secret = crate::totp::generate_secret();
table.add_totp_slot(&master_key, secret.as_bytes()).unwrap();
assert!(table.unwrap_with_totp("000000").is_err());
}
#[test]
fn fido2_slot_round_trip() {
let master_key = generate_master_key();
let mut table = SlotTable::new_with_passphrase(&master_key, "pass").unwrap();
let cred_id = b"test-credential-id-12345";
let pub_key = b"fake-cose-public-key-data";
table
.add_fido2_slot(&master_key, cred_id, pub_key, "garden.local")
.unwrap();
let recovered = table.unwrap_with_fido2(cred_id).unwrap();
assert_eq!(master_key, recovered);
}
#[test]
fn fido2_wrong_credential_fails() {
let master_key = generate_master_key();
let mut table = SlotTable::new_with_passphrase(&master_key, "pass").unwrap();
table
.add_fido2_slot(&master_key, b"real-cred", b"pub-key", "garden.local")
.unwrap();
assert!(table.unwrap_with_fido2(b"wrong-cred").is_err());
}
#[test]
fn envelope_encrypt_new_round_trip() {
let plaintext = b"test CA private key DER bytes";
let (encrypted, table, master_key) =
envelope_encrypt_new(plaintext, "my-passphrase").unwrap();
let recovered_mk = table.unwrap_with_passphrase("my-passphrase").unwrap();
assert_eq!(master_key, recovered_mk);
let recovered = decrypt_with_master_key(&encrypted, &master_key).unwrap();
assert_eq!(&recovered, plaintext);
}
#[test]
fn migrate_preserves_ca_key() {
let ca_key_der = b"simulated CA private key bytes!!";
let old_encrypted = encrypt_bytes(ca_key_der, "old-pass").unwrap();
let (new_encrypted, table, master_key) =
migrate_to_envelope(&old_encrypted, "old-pass").unwrap();
let recovered_mk = table.unwrap_with_passphrase("old-pass").unwrap();
assert_eq!(master_key, recovered_mk);
let recovered = decrypt_with_master_key(&new_encrypted, &master_key).unwrap();
assert_eq!(&recovered, ca_key_der);
}
#[test]
fn auto_unlock_marker() {
let master_key = generate_master_key();
let mut table = SlotTable::new_with_passphrase(&master_key, "pass").unwrap();
assert!(!table.has_auto_unlock());
table.add_auto_unlock();
assert!(table.has_auto_unlock());
table.remove_auto_unlock();
assert!(!table.has_auto_unlock());
}
#[cfg(feature = "keyring")]
#[test]
fn available_methods_lists_all_slots() {
let master_key = generate_master_key();
let mut table = SlotTable::new_with_passphrase(&master_key, "pass").unwrap();
assert_eq!(table.available_methods(), vec!["passphrase"]);
table.add_auto_unlock();
let secret = crate::totp::generate_secret();
table.add_totp_slot(&master_key, secret.as_bytes()).unwrap();
table
.add_fido2_slot(&master_key, b"cred", b"pk", "rp")
.unwrap();
let methods = table.available_methods();
assert!(methods.contains(&"passphrase"));
assert!(methods.contains(&"auto_unlock"));
assert!(methods.contains(&"totp"));
assert!(methods.contains(&"fido2"));
}
#[cfg(feature = "keyring")]
#[test]
fn slot_table_serialization_round_trip() {
let master_key = generate_master_key();
let mut table = SlotTable::new_with_passphrase(&master_key, "pass").unwrap();
table.add_auto_unlock();
let secret = crate::totp::generate_secret();
table.add_totp_slot(&master_key, secret.as_bytes()).unwrap();
let json = serde_json::to_string_pretty(&table).unwrap();
let loaded: SlotTable = serde_json::from_str(&json).unwrap();
assert_eq!(loaded.version, 1);
assert_eq!(loaded.slots.len(), 3);
let recovered = loaded.unwrap_with_passphrase("pass").unwrap();
assert_eq!(master_key, recovered);
}
}