quicproquo/crates/quicproquo-kt/src/revocation.rs

//! Key revocation tracking for the Key Transparency log.
//!
//! Revocation entries are appended to the same Merkle log as regular key
//! bindings, using a distinct leaf hash prefix to differentiate them. A
//! separate in-memory index tracks which identity keys have been revoked,
//! enabling O(1) revocation checks.
//!
//! ## Revocation leaf hash
//!
//! ```text
//! SHA-256(0x02 || SHA-256(identity_key || 0x00 || reason_bytes))
//! ```
//!
//! The 0x02 prefix domain-separates revocation leaves from binding leaves (0x00)
//! and internal nodes (0x01).

use std::collections::HashMap;

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::{KtError, MerkleLog};

/// Domain-separation prefix for revocation leaves.
const REVOCATION_PREFIX: u8 = 0x02;

/// Reason for key revocation.
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum RevocationReason {
    /// Key material was compromised.
    Compromised,
    /// Key was superseded by a new key.
    Superseded,
    /// User-initiated revocation (e.g. account deletion).
    UserRevoked,
}

impl RevocationReason {
    fn as_bytes(&self) -> &[u8] {
        match self {
            RevocationReason::Compromised => b"compromised",
            RevocationReason::Superseded => b"superseded",
            RevocationReason::UserRevoked => b"user_revoked",
        }
    }

    /// Parse from a string tag.
    pub fn from_tag(tag: &str) -> Option<Self> {
        match tag {
            "compromised" => Some(RevocationReason::Compromised),
            "superseded" => Some(RevocationReason::Superseded),
            "user_revoked" => Some(RevocationReason::UserRevoked),
            _ => None,
        }
    }

    /// Return the string tag for serialization.
    pub fn as_tag(&self) -> &str {
        match self {
            RevocationReason::Compromised => "compromised",
            RevocationReason::Superseded => "superseded",
            RevocationReason::UserRevoked => "user_revoked",
        }
    }
}

/// A record of a key revocation.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RevocationEntry {
    /// The 32-byte identity key that was revoked.
    pub identity_key: Vec<u8>,
    /// Reason for revocation.
    pub reason: RevocationReason,
    /// Timestamp (ms since UNIX epoch) when the revocation was recorded.
    pub timestamp_ms: u64,
    /// Index of the revocation leaf in the Merkle log.
    pub leaf_index: u64,
}

/// Tracks revoked identity keys alongside the Merkle log.
///
/// Revocation entries are appended to the Merkle log (with a distinct prefix)
/// and indexed in-memory by identity key for O(1) lookup.
#[derive(Default, Serialize, Deserialize, Clone)]
pub struct RevocationLog {
    /// Revocation entries in append order.
    entries: Vec<RevocationEntry>,
    /// Index from identity_key bytes to entry index for O(1) lookup.
    #[serde(skip)]
    index: HashMap<Vec<u8>, usize>,
}

impl RevocationLog {
    /// Create an empty revocation log.
    pub fn new() -> Self {
        Self::default()
    }

    /// Rebuild the in-memory index from the entries list.
    ///
    /// Must be called after deserialization.
    pub fn rebuild_index(&mut self) {
        self.index.clear();
        for (i, entry) in self.entries.iter().enumerate() {
            self.index.insert(entry.identity_key.clone(), i);
        }
    }

    /// Record a key revocation, appending a revocation leaf to the Merkle log.
    ///
    /// Returns the leaf index in the Merkle log, or an error if the key is
    /// already revoked.
    pub fn revoke(
        &mut self,
        kt_log: &mut MerkleLog,
        identity_key: &[u8],
        reason: RevocationReason,
        timestamp_ms: u64,
    ) -> Result<u64, KtError> {
        if self.index.contains_key(identity_key) {
            return Err(KtError::AlreadyRevoked);
        }

        // Compute revocation leaf hash and append to the Merkle log.
        let leaf = revocation_leaf_hash(identity_key, &reason);
        let leaf_index = kt_log.append_raw(leaf);

        let entry = RevocationEntry {
            identity_key: identity_key.to_vec(),
            reason,
            timestamp_ms,
            leaf_index,
        };

        let entry_idx = self.entries.len();
        self.entries.push(entry);
        self.index.insert(identity_key.to_vec(), entry_idx);

        Ok(leaf_index)
    }

    /// Check if an identity key has been revoked.
    pub fn is_revoked(&self, identity_key: &[u8]) -> bool {
        self.index.contains_key(identity_key)
    }

    /// Get the revocation entry for an identity key, if revoked.
    pub fn get(&self, identity_key: &[u8]) -> Option<&RevocationEntry> {
        self.index
            .get(identity_key)
            .map(|&idx| &self.entries[idx])
    }

    /// Return all revocation entries in append order.
    pub fn entries(&self) -> &[RevocationEntry] {
        &self.entries
    }

    /// Number of revoked keys.
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    /// Return `true` if no keys have been revoked.
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// Serialise the revocation log to bytes (bincode).
    pub fn to_bytes(&self) -> Result<Vec<u8>, KtError> {
        bincode::serialize(self).map_err(|e| KtError::Serialisation(e.to_string()))
    }

    /// Deserialise from bytes and rebuild the in-memory index.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, KtError> {
        let mut log: Self =
            bincode::deserialize(bytes).map_err(|e| KtError::Serialisation(e.to_string()))?;
        log.rebuild_index();
        Ok(log)
    }
}

/// Compute the leaf hash for a revocation entry.
///
/// `SHA-256(0x02 || SHA-256(identity_key || 0x00 || reason_bytes))`
pub fn revocation_leaf_hash(identity_key: &[u8], reason: &RevocationReason) -> [u8; 32] {
    let mut inner = Sha256::new();
    inner.update(identity_key);
    inner.update([0x00]);
    inner.update(reason.as_bytes());
    let inner_digest: [u8; 32] = inner.finalize().into();

    let mut outer = Sha256::new();
    outer.update([REVOCATION_PREFIX]);
    outer.update(inner_digest);
    outer.finalize().into()
}

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

    #[test]
    fn revoke_and_check() {
        let mut kt = MerkleLog::new();
        let mut revlog = RevocationLog::new();

        // Append a normal binding first.
        kt.append("alice", &[1u8; 32]);

        // Revoke alice's key.
        let leaf_idx = revlog
            .revoke(&mut kt, &[1u8; 32], RevocationReason::Compromised, 1000)
            .unwrap();
        assert_eq!(leaf_idx, 1); // second leaf in the log

        assert!(revlog.is_revoked(&[1u8; 32]));
        assert!(!revlog.is_revoked(&[2u8; 32]));

        let entry = revlog.get(&[1u8; 32]).unwrap();
        assert_eq!(entry.reason, RevocationReason::Compromised);
        assert_eq!(entry.timestamp_ms, 1000);
    }

    #[test]
    fn double_revoke_fails() {
        let mut kt = MerkleLog::new();
        let mut revlog = RevocationLog::new();

        revlog
            .revoke(&mut kt, &[1u8; 32], RevocationReason::Compromised, 1000)
            .unwrap();
        let result = revlog.revoke(&mut kt, &[1u8; 32], RevocationReason::Superseded, 2000);
        assert!(matches!(result, Err(KtError::AlreadyRevoked)));
    }

    #[test]
    fn revocation_leaf_is_distinct_from_binding_leaf() {
        let binding_hash = crate::leaf_hash("alice", &[1u8; 32]);
        let revocation_hash =
            revocation_leaf_hash(&[1u8; 32], &RevocationReason::Compromised);
        assert_ne!(binding_hash, revocation_hash);
    }

    #[test]
    fn serialization_roundtrip() {
        let mut kt = MerkleLog::new();
        let mut revlog = RevocationLog::new();

        revlog
            .revoke(&mut kt, &[1u8; 32], RevocationReason::Compromised, 1000)
            .unwrap();
        revlog
            .revoke(&mut kt, &[2u8; 32], RevocationReason::Superseded, 2000)
            .unwrap();

        let bytes = revlog.to_bytes().unwrap();
        let restored = RevocationLog::from_bytes(&bytes).unwrap();

        assert_eq!(restored.len(), 2);
        assert!(restored.is_revoked(&[1u8; 32]));
        assert!(restored.is_revoked(&[2u8; 32]));
    }

    #[test]
    fn reason_tag_roundtrip() {
        for reason in &[
            RevocationReason::Compromised,
            RevocationReason::Superseded,
            RevocationReason::UserRevoked,
        ] {
            let tag = reason.as_tag();
            let parsed = RevocationReason::from_tag(tag).unwrap();
            assert_eq!(*reason, parsed);
        }
    }
}