exochain-identity 0.2.0-beta

// Copyright 2026 Exochain Foundation
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at:
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

//! DID-based signature verification and key rotation.
//!
//! Bridges the gap between key management (`key_management`) and DID documents
//! (`did`) by providing:
//! - [`KeyVault`] trait for abstracting TEE/HSM key storage
//! - [`verify_did_signature`] for verifying signatures against DID document
//!   verification methods with multibase key decoding
//! - [`rotate_verification_key`] for proper lifecycle management of verification
//!   methods (deactivate old, add new with version increment)

use exo_core::{Did, PublicKey, crypto};

use crate::did::{DidDocument, VerificationMethod};

const ED25519_VERIFICATION_KEY_TYPE: &str = "Ed25519VerificationKey2020";

/// Errors specific to DID verification operations.
#[derive(Debug, thiserror::Error)]
pub enum DidVerificationError {
    /// Verification method not found by key ID.
    #[error("verification method not found: {0}")]
    MethodNotFound(String),

    /// Key has been revoked or deactivated.
    #[error("verification method revoked: {0}")]
    MethodRevoked(String),

    /// Verification method is not controlled by, rooted in, or key-bound to
    /// the DID document that presents it.
    #[error("verification method not bound to DID document: {0}")]
    MethodNotDocumentBound(String),

    /// Cryptographic operation failed (e.g., invalid multibase encoding,
    /// wrong key length, unsupported multibase prefix).
    #[error("cryptographic error: {0}")]
    CryptoError(String),

    /// Signature verification failed.
    #[error("invalid signature")]
    InvalidSignature,
}

fn decode_ed25519_multibase_public_key(encoded: &str) -> Result<PublicKey, DidVerificationError> {
    let pub_key_bytes = if let Some(encoded_key) = encoded.strip_prefix('z') {
        bs58::decode(encoded_key)
            .into_vec()
            .map_err(|e| DidVerificationError::CryptoError(format!("base58 decode: {e}")))?
    } else {
        return Err(DidVerificationError::CryptoError(
            "unsupported multibase prefix (expected 'z' for base58btc)".to_string(),
        ));
    };

    let pub_key_array: [u8; 32] = pub_key_bytes.try_into().map_err(|_| {
        DidVerificationError::CryptoError("public key must be 32 bytes".to_string())
    })?;

    Ok(PublicKey::from_bytes(pub_key_array))
}

/// Validate that an active Ed25519 DID verification method is controlled by
/// the document DID, rooted at the document fragment namespace, and backed by
/// key material declared in the document's `public_keys`.
///
/// This prevents a registered DID document from proving possession of one
/// public key while presenting a different active authentication key.
pub fn validate_verification_method_document_binding(
    doc: &DidDocument,
    method: &VerificationMethod,
) -> Result<PublicKey, DidVerificationError> {
    if doc.revoked {
        return Err(DidVerificationError::MethodRevoked(
            doc.id.as_str().to_owned(),
        ));
    }

    if !method.active || method.revoked_at.is_some() {
        return Err(DidVerificationError::MethodRevoked(method.id.clone()));
    }

    if method.key_type != ED25519_VERIFICATION_KEY_TYPE {
        return Err(DidVerificationError::MethodNotDocumentBound(format!(
            "{} has unsupported key type {}",
            method.id, method.key_type
        )));
    }

    if method.controller != doc.id {
        return Err(DidVerificationError::MethodNotDocumentBound(format!(
            "{} controller {} does not match document DID {}",
            method.id, method.controller, doc.id
        )));
    }

    let document_fragment_prefix = format!("{}#", doc.id);
    if !method.id.starts_with(&document_fragment_prefix) {
        return Err(DidVerificationError::MethodNotDocumentBound(format!(
            "{} is not rooted under document DID {}",
            method.id, doc.id
        )));
    }

    let public_key = decode_ed25519_multibase_public_key(&method.public_key_multibase)?;
    if !doc
        .public_keys
        .iter()
        .any(|declared| declared == &public_key)
    {
        return Err(DidVerificationError::MethodNotDocumentBound(format!(
            "{} key is not declared in the DID document public_keys",
            method.id
        )));
    }

    Ok(public_key)
}

/// Abstract key vault interface for secure key storage.
///
/// In production, implementations would interface with a TEE (Trusted
/// Execution Environment) or HSM (Hardware Security Module). For testing,
/// an in-memory implementation suffices.
pub trait KeyVault {
    /// Retrieve a public key for a DID at a specific version.
    fn get_public_key(&self, did: &Did, version: u64) -> Result<PublicKey, DidVerificationError>;

    /// Store a public key for a DID at a specific version.
    fn store_public_key(
        &mut self,
        did: &Did,
        key: PublicKey,
        version: u64,
    ) -> Result<(), DidVerificationError>;
}

/// Verify a signature against a DID document's verification methods.
///
/// Resolves the verification method by `key_id` (e.g., `"did:exo:123#key-1"`),
/// checks that the key is active, decodes the multibase base58btc public key,
/// and verifies the signature over the provided message.
///
/// # Errors
///
/// Returns [`DidVerificationError::MethodNotFound`] if the key ID doesn't
/// match any verification method.
/// Returns [`DidVerificationError::MethodRevoked`] if the key is inactive.
/// Returns [`DidVerificationError::InvalidSignature`] if verification fails.
pub fn verify_did_signature(
    doc: &DidDocument,
    key_id: &str,
    message: &[u8],
    signature: &exo_core::Signature,
) -> Result<(), DidVerificationError> {
    let method = doc
        .verification_methods
        .iter()
        .find(|m| m.id == key_id)
        .ok_or_else(|| DidVerificationError::MethodNotFound(key_id.to_string()))?;

    let public_key = validate_verification_method_document_binding(doc, method)?;

    if crypto::verify(message, signature, &public_key) {
        Ok(())
    } else {
        Err(DidVerificationError::InvalidSignature)
    }
}

/// Rotate a verification key in a DID document.
///
/// Deactivates the old verification method identified by `old_key_id`,
/// creates a new verification method with an incremented version, and
/// appends it to the document's verification methods.
///
/// # Arguments
///
/// * `doc` — The DID document to mutate.
/// * `old_key_id` — ID of the verification method to deactivate.
/// * `new_public_key` — Raw 32-byte Ed25519 public key for the new method.
/// * `controller` — DID that controls the new key.
/// * `current_time_ms` — Current wall-clock time in milliseconds (for lifecycle tracking).
///
/// # Returns
///
/// The newly created [`VerificationMethod`].
pub fn rotate_verification_key(
    doc: &mut DidDocument,
    old_key_id: &str,
    new_public_key: &[u8; 32],
    controller: &Did,
    current_time_ms: u64,
) -> Result<VerificationMethod, DidVerificationError> {
    // Find the old method
    let old_method_idx = doc
        .verification_methods
        .iter()
        .position(|m| m.id == old_key_id)
        .ok_or_else(|| DidVerificationError::MethodNotFound(old_key_id.to_string()))?;

    let old_version = doc.verification_methods[old_method_idx].version;
    let new_version = old_version.checked_add(1).ok_or_else(|| {
        DidVerificationError::CryptoError(format!(
            "verification method version overflow for key {old_key_id}"
        ))
    })?;

    // Deactivate old key
    doc.verification_methods[old_method_idx].active = false;
    doc.verification_methods[old_method_idx].revoked_at = Some(current_time_ms);

    // Create new method with incremented version
    let new_id = format!("{}#key-{}", doc.id, new_version);
    let multibase = format!("z{}", bs58::encode(new_public_key).into_string());

    let new_method = VerificationMethod {
        id: new_id,
        key_type: "Ed25519VerificationKey2020".to_string(),
        controller: controller.clone(),
        public_key_multibase: multibase,
        version: new_version,
        active: true,
        valid_from: current_time_ms,
        revoked_at: None,
    };

    doc.public_keys.clear();
    doc.public_keys.push(PublicKey::from_bytes(*new_public_key));
    doc.verification_methods.push(new_method.clone());
    doc.updated = exo_core::Timestamp::new(current_time_ms, 0);

    Ok(new_method)
}

// ===========================================================================
// Tests
// ===========================================================================

#[cfg(test)]
mod tests {
    use exo_core::{
        Timestamp,
        crypto::{generate_keypair, sign},
    };

    use super::*;

    fn test_did() -> Did {
        Did::new("did:exo:test-verification").expect("valid")
    }

    fn make_doc_with_verification(did: Did, pk: PublicKey) -> DidDocument {
        let multibase = format!("z{}", bs58::encode(pk.as_bytes()).into_string());
        DidDocument {
            id: did.clone(),
            public_keys: vec![pk],
            authentication: vec![],
            verification_methods: vec![VerificationMethod {
                id: format!("{}#key-1", did),
                key_type: "Ed25519VerificationKey2020".to_string(),
                controller: did,
                public_key_multibase: multibase,
                version: 1,
                active: true,
                valid_from: 1000,
                revoked_at: None,
            }],
            hybrid_verification_methods: vec![],
            service_endpoints: vec![],
            created: Timestamp::new(1000, 0),
            updated: Timestamp::new(1000, 0),
            revoked: false,
        }
    }

    #[test]
    fn verify_valid_signature() {
        let (pk, sk) = generate_keypair();
        let did = test_did();
        let doc = make_doc_with_verification(did.clone(), pk);

        let message = b"hello world";
        let signature = sign(message, &sk);

        let key_id = format!("{}#key-1", did);
        assert!(verify_did_signature(&doc, &key_id, message, &signature).is_ok());
    }

    #[test]
    fn verify_wrong_signature_fails() {
        let (pk, _sk) = generate_keypair();
        let (_pk2, sk2) = generate_keypair();
        let did = test_did();
        let doc = make_doc_with_verification(did.clone(), pk);

        let message = b"hello world";
        let wrong_sig = sign(message, &sk2);

        let key_id = format!("{}#key-1", did);
        let err = verify_did_signature(&doc, &key_id, message, &wrong_sig).unwrap_err();
        assert!(matches!(err, DidVerificationError::InvalidSignature));
    }

    #[test]
    fn verify_unknown_key_id_fails() {
        let (pk, sk) = generate_keypair();
        let did = test_did();
        let doc = make_doc_with_verification(did, pk);

        let message = b"test";
        let signature = sign(message, &sk);

        let err =
            verify_did_signature(&doc, "nonexistent#key-99", message, &signature).unwrap_err();
        assert!(matches!(err, DidVerificationError::MethodNotFound(_)));
    }

    #[test]
    fn verify_revoked_key_fails() {
        let (pk, sk) = generate_keypair();
        let did = test_did();
        let mut doc = make_doc_with_verification(did.clone(), pk);

        // Revoke the key
        doc.verification_methods[0].active = false;

        let message = b"test";
        let signature = sign(message, &sk);

        let key_id = format!("{}#key-1", did);
        let err = verify_did_signature(&doc, &key_id, message, &signature).unwrap_err();
        assert!(matches!(err, DidVerificationError::MethodRevoked(_)));
    }

    #[test]
    fn verify_rejects_method_key_not_declared_by_document() {
        let (declared_pk, _) = generate_keypair();
        let (method_pk, method_sk) = generate_keypair();
        let did = test_did();
        let mut doc = make_doc_with_verification(did.clone(), declared_pk);
        doc.verification_methods[0].public_key_multibase =
            format!("z{}", bs58::encode(method_pk.as_bytes()).into_string());

        let message = b"method key must be document-bound";
        let signature = sign(message, &method_sk);
        let key_id = format!("{}#key-1", did);

        let err = verify_did_signature(&doc, &key_id, message, &signature).unwrap_err();
        assert!(
            err.to_string().contains("not declared"),
            "verification should reject active methods whose keys are not declared by the DID document: {err}"
        );
    }

    #[test]
    fn verify_bad_multibase_prefix_fails() {
        let (pk, sk) = generate_keypair();
        let did = test_did();
        let mut doc = make_doc_with_verification(did.clone(), pk);

        // Set unsupported multibase prefix
        doc.verification_methods[0].public_key_multibase =
            format!("m{}", bs58::encode(pk.as_bytes()).into_string());

        let message = b"test";
        let signature = sign(message, &sk);

        let key_id = format!("{}#key-1", did);
        let err = verify_did_signature(&doc, &key_id, message, &signature).unwrap_err();
        assert!(matches!(err, DidVerificationError::CryptoError(_)));
    }

    #[test]
    fn rotate_key_success() {
        let (pk, _sk) = generate_keypair();
        let did = test_did();
        let mut doc = make_doc_with_verification(did.clone(), pk);

        let (new_pk, _new_sk) = generate_keypair();
        let new_method = rotate_verification_key(
            &mut doc,
            &format!("{}#key-1", did),
            new_pk.as_bytes(),
            &did,
            2000,
        )
        .expect("rotation should succeed");

        // Old key deactivated
        assert!(!doc.verification_methods[0].active);
        assert_eq!(doc.verification_methods[0].revoked_at, Some(2000));

        // New key active
        assert_eq!(new_method.version, 2);
        assert!(new_method.active);
        assert_eq!(doc.public_keys, vec![new_pk]);
        assert_eq!(doc.verification_methods.len(), 2);
        assert_eq!(doc.updated.physical_ms, 2000);
    }

    #[test]
    fn rotate_unknown_key_fails() {
        let (pk, _sk) = generate_keypair();
        let did = test_did();
        let mut doc = make_doc_with_verification(did.clone(), pk);

        let (new_pk, _) = generate_keypair();
        let err = rotate_verification_key(
            &mut doc,
            "nonexistent#key-99",
            new_pk.as_bytes(),
            &did,
            2000,
        )
        .unwrap_err();
        assert!(matches!(err, DidVerificationError::MethodNotFound(_)));
    }

    #[test]
    fn rotate_key_version_overflow_fails_without_mutating_document() {
        let (pk, _sk) = generate_keypair();
        let did = test_did();
        let mut doc = make_doc_with_verification(did.clone(), pk);
        doc.verification_methods[0].version = u64::MAX;
        let original_doc = doc.clone();

        let (new_pk, _) = generate_keypair();
        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            rotate_verification_key(
                &mut doc,
                &format!("{}#key-1", did),
                new_pk.as_bytes(),
                &did,
                2000,
            )
        }));

        assert!(
            result.is_ok(),
            "version overflow must return an error instead of panicking"
        );
        let rotation_result = match result {
            Ok(rotation_result) => rotation_result,
            Err(_) => unreachable!("asserted above"),
        };
        assert!(
            matches!(
                rotation_result,
                Err(DidVerificationError::CryptoError(ref reason))
                    if reason.contains("version overflow")
            ),
            "unexpected rotation result: {rotation_result:?}"
        );
        assert_eq!(
            doc, original_doc,
            "failed rotation must not revoke the old key or append a replacement"
        );
    }

    #[test]
    fn multibase_decoding_uses_prefix_stripping_instead_of_byte_slicing() {
        let source = include_str!("did_verification.rs");
        let direct_slice = concat!("public_key_multibase", "[1..]");
        assert!(
            source.contains("strip_prefix('z')"),
            "multibase decoding should strip the ASCII prefix without direct byte slicing"
        );
        assert!(
            !source.contains(direct_slice),
            "multibase decoding must not slice the key string by byte index"
        );
    }

    #[test]
    fn verify_after_rotation() {
        let (pk, _sk) = generate_keypair();
        let did = test_did();
        let mut doc = make_doc_with_verification(did.clone(), pk);

        let (new_pk, new_sk) = generate_keypair();
        let new_method = rotate_verification_key(
            &mut doc,
            &format!("{}#key-1", did),
            new_pk.as_bytes(),
            &did,
            2000,
        )
        .expect("rotation");

        // Verify with new key works
        let message = b"post-rotation message";
        let signature = sign(message, &new_sk);
        assert!(verify_did_signature(&doc, &new_method.id, message, &signature).is_ok());

        // Verify with old key fails (revoked)
        let old_key_id = format!("{}#key-1", did);
        let err = verify_did_signature(&doc, &old_key_id, message, &signature).unwrap_err();
        assert!(matches!(err, DidVerificationError::MethodRevoked(_)));
    }
}