rialo-types 0.12.2

// Copyright (c) Subzero Labs, Inc.
// SPDX-License-Identifier: Apache-2.0

//! Consensus-related cryptographic types.
//!
//! These types wrap fastcrypto primitives and are used for validator identity
//! and block signing in the consensus protocol. They provide type safety over
//! raw byte arrays while preserving the underlying cryptographic validation.

use std::sync::atomic::{AtomicU64, Ordering};

use fastcrypto::{
    bls12381, ed25519,
    encoding::{self, Encoding},
    error::FastCryptoError,
    traits::{KeyPair as _, Signer as _, ToFromBytes, VerifyingKey},
};
use rand::SeedableRng;
use serde::{Deserialize, Serialize};

// =============================================================================
// Authority Public Key (BLS12-381)
// =============================================================================

/// Authority key represents the identity of an authority in the consensus protocol.
///
/// This is a BLS12-381 public key (96 bytes) used for validator identity.
/// It is primarily used for identity verification and not for cryptographic operations.
///
/// NOTE: No `Hash` derive - `BLS12381PublicKey` may not support it.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct AuthorityPublicKey(bls12381::min_sig::BLS12381PublicKey);

impl AuthorityPublicKey {
    /// Size of the BLS12-381 public key in bytes.
    pub const LENGTH: usize = bls12381::BLS_G2_LENGTH;

    /// Creates a new `AuthorityPublicKey` from a fastcrypto BLS12-381 public key.
    pub fn new(key: bls12381::min_sig::BLS12381PublicKey) -> Self {
        Self(key)
    }

    /// Creates an `AuthorityPublicKey` from raw bytes.
    ///
    /// Returns an error if the bytes do not represent a valid BLS12-381 public key.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, FastCryptoError> {
        Ok(Self(bls12381::min_sig::BLS12381PublicKey::from_bytes(
            bytes,
        )?))
    }

    /// Generate a deterministic unique AuthorityPublicKey for testing/fixtures.
    pub fn new_unique() -> Self {
        static I: AtomicU64 = AtomicU64::new(1);

        let seed = I.fetch_add(1, Ordering::Relaxed);
        let mut rng = rand::rngs::StdRng::seed_from_u64(seed);
        Self(
            bls12381::min_sig::BLS12381KeyPair::generate(&mut rng)
                .public()
                .clone(),
        )
    }

    /// Returns a reference to the inner fastcrypto public key.
    pub fn inner(&self) -> &bls12381::min_sig::BLS12381PublicKey {
        &self.0
    }

    /// Returns the length of the key in bytes.
    pub const fn len() -> usize {
        Self::LENGTH
    }

    /// Returns the key as a fixed-size byte array.
    ///
    /// For a slice, use `&key.to_bytes()` or `key.to_bytes().as_ref()`.
    pub fn to_bytes(&self) -> [u8; Self::LENGTH] {
        let mut bytes = [0u8; Self::LENGTH];
        bytes.copy_from_slice(self.0.as_bytes());
        bytes
    }
}

// =============================================================================
// Protocol Public Key (Ed25519)
// =============================================================================

/// Protocol key is used for signing blocks and verifying block signatures.
///
/// This is an Ed25519 public key (32 bytes) used in the consensus protocol
/// for block verification.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
pub struct ProtocolPublicKey(ed25519::Ed25519PublicKey);

impl ProtocolPublicKey {
    /// Size of the Ed25519 public key in bytes.
    pub const LENGTH: usize = 32;

    /// Creates a new `ProtocolPublicKey` from a fastcrypto Ed25519 public key.
    pub fn new(key: ed25519::Ed25519PublicKey) -> Self {
        Self(key)
    }

    /// Creates a `ProtocolPublicKey` from raw bytes.
    ///
    /// Returns an error if the bytes do not represent a valid Ed25519 public key.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, FastCryptoError> {
        Ok(Self(ed25519::Ed25519PublicKey::from_bytes(bytes)?))
    }

    /// Verifies a signature against this public key.
    pub fn verify(
        &self,
        message: &[u8],
        signature: &ProtocolKeySignature,
    ) -> Result<(), FastCryptoError> {
        self.0.verify(message, signature.inner())
    }

    /// Returns the key as a fixed-size byte array.
    pub fn to_bytes(&self) -> [u8; Self::LENGTH] {
        self.0 .0.to_bytes()
    }
}

// =============================================================================
// Protocol Key Signature (Ed25519)
// =============================================================================

/// Signature produced by a protocol key.
///
/// This is an Ed25519 signature (64 bytes) used for block verification.
#[derive(Serialize, Deserialize)]
pub struct ProtocolKeySignature(ed25519::Ed25519Signature);

impl ProtocolKeySignature {
    /// Size of the Ed25519 signature in bytes.
    pub const LENGTH: usize = 64;

    /// Creates a `ProtocolKeySignature` from a fastcrypto Ed25519 signature.
    pub fn from_inner(sig: ed25519::Ed25519Signature) -> Self {
        Self(sig)
    }

    /// Creates a `ProtocolKeySignature` from raw bytes.
    ///
    /// Returns an error if the bytes do not represent a valid Ed25519 signature.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, FastCryptoError> {
        Ok(Self(ed25519::Ed25519Signature::from_bytes(bytes)?))
    }

    /// Returns the signature as a fixed-size byte array.
    ///
    /// For a slice, use `&sig.to_bytes()` or `sig.to_bytes().as_ref()`.
    pub fn to_bytes(&self) -> [u8; Self::LENGTH] {
        let mut bytes = [0u8; Self::LENGTH];
        bytes.copy_from_slice(self.0.as_bytes());
        bytes
    }

    /// Returns a reference to the inner fastcrypto signature.
    pub fn inner(&self) -> &ed25519::Ed25519Signature {
        &self.0
    }
}

/// Alias used when an Ed25519 signature is produced by a [`NetworkKeyPair`] rather
/// than a [`ProtocolKeyPair`]. The underlying cryptographic type is identical
/// (`ed25519::Ed25519Signature`); the alias exists solely for readability at call
/// sites that operate on network-key signatures (e.g. `GenericLocalEndorsement` in
/// `rialo_protocol_types::admin`).
pub type GenesisSignature = ProtocolKeySignature;

// =============================================================================
// Network Public Key (Ed25519)
// =============================================================================

/// Network key is used for TLS and as the network identity of the authority.
///
/// This is an Ed25519 public key (32 bytes) used for network-level identity
/// and secure communication.
#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct NetworkPublicKey(ed25519::Ed25519PublicKey);

impl NetworkPublicKey {
    /// Size of the Ed25519 public key in bytes.
    pub const LENGTH: usize = 32;

    /// Creates a new `NetworkPublicKey` from a fastcrypto Ed25519 public key.
    pub fn new(key: ed25519::Ed25519PublicKey) -> Self {
        Self(key)
    }

    /// Creates a `NetworkPublicKey` from raw bytes.
    ///
    /// Returns an error if the bytes do not represent a valid Ed25519 public key.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, FastCryptoError> {
        Ok(Self(ed25519::Ed25519PublicKey::from_bytes(bytes)?))
    }

    /// Consumes self and returns the inner fastcrypto public key.
    pub fn into_inner(self) -> ed25519::Ed25519PublicKey {
        self.0
    }

    /// Returns the key as a fixed-size byte array.
    pub fn to_bytes(&self) -> [u8; Self::LENGTH] {
        self.0 .0.to_bytes()
    }

    /// Verifies an Ed25519 signature against this network public key.
    pub fn verify(
        &self,
        message: &[u8],
        signature: &GenesisSignature,
    ) -> Result<(), FastCryptoError> {
        self.0.verify(message, signature.inner())
    }
}

// =============================================================================
// Network Key Pair (Ed25519)
// =============================================================================

/// Private key for network identity.
///
/// Be very careful when serializing the private key. It should be encrypted before being stored.
#[derive(Serialize, Deserialize)]
pub struct NetworkPrivateKey(ed25519::Ed25519PrivateKey);

impl NetworkPrivateKey {
    /// Consumes self and returns the inner fastcrypto private key.
    pub fn into_inner(self) -> ed25519::Ed25519PrivateKey {
        self.0
    }
}

/// Keypair for network identity (Ed25519).
///
/// Be very careful when serializing the key pair. It should be encrypted before being stored.
pub struct NetworkKeyPair(ed25519::Ed25519KeyPair);

impl NetworkKeyPair {
    /// Creates a new `NetworkKeyPair` from a fastcrypto Ed25519 keypair.
    pub fn new(keypair: ed25519::Ed25519KeyPair) -> Self {
        Self(keypair)
    }

    /// Generates a new random `NetworkKeyPair`.
    pub fn generate<R: rand::Rng + fastcrypto::traits::AllowedRng>(rng: &mut R) -> Self {
        Self(ed25519::Ed25519KeyPair::generate(rng))
    }

    /// Returns the public key for this keypair.
    pub fn public(&self) -> NetworkPublicKey {
        NetworkPublicKey::new(self.0.public().clone())
    }

    /// Signs a message with this keypair.
    pub fn sign(&self, message: &[u8]) -> GenesisSignature {
        GenesisSignature::from_inner(self.0.sign(message))
    }

    /// Consumes self and returns the private key.
    pub fn private_key(self) -> NetworkPrivateKey {
        NetworkPrivateKey(self.0.copy().private())
    }

    /// Consumes self and returns the private key bytes.
    pub fn private_key_bytes(self) -> [u8; 32] {
        self.0.private().0.to_bytes()
    }

    /// Export the key pair to a Base64 encoded string.
    pub fn to_base64(&self) -> String {
        let secret_bytes = self.0.copy().private().0.to_bytes();
        let mut key_bytes = Vec::new();
        key_bytes.extend_from_slice(&secret_bytes);
        encoding::Base64::encode(&key_bytes)
    }

    /// Import a key pair from a Base64 encoded string.
    pub fn from_base64(encoded: &str) -> Result<Self, FastCryptoError> {
        let key_bytes = encoding::Base64::decode(encoded)?;
        let secret = ed25519::Ed25519PrivateKey::from_bytes(&key_bytes)?;
        let keypair = ed25519::Ed25519KeyPair::from(secret);
        Ok(NetworkKeyPair(keypair))
    }
}

impl Clone for NetworkKeyPair {
    fn clone(&self) -> Self {
        Self(self.0.copy())
    }
}

impl Serialize for NetworkKeyPair {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_str(&self.to_base64())
    }
}

impl<'de> Deserialize<'de> for NetworkKeyPair {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        use serde::de::Error;
        let s = String::deserialize(deserializer)?;
        NetworkKeyPair::from_base64(&s)
            .map_err(|e| Error::custom(format!("Failed to deserialize NetworkKeyPair: {e}")))
    }
}

// =============================================================================
// Protocol Key Pair (Ed25519)
// =============================================================================

/// Keypair for protocol signing (Ed25519).
///
/// Be very careful when serializing the key pair. It should be encrypted before being stored.
pub struct ProtocolKeyPair(ed25519::Ed25519KeyPair);

impl ProtocolKeyPair {
    /// Creates a new `ProtocolKeyPair` from a fastcrypto Ed25519 keypair.
    pub fn new(keypair: ed25519::Ed25519KeyPair) -> Self {
        Self(keypair)
    }

    /// Generates a new random `ProtocolKeyPair`.
    pub fn generate<R: rand::Rng + fastcrypto::traits::AllowedRng>(rng: &mut R) -> Self {
        Self(ed25519::Ed25519KeyPair::generate(rng))
    }

    /// Returns the public key for this keypair.
    pub fn public(&self) -> ProtocolPublicKey {
        ProtocolPublicKey::new(self.0.public().clone())
    }

    /// Signs a message with this keypair.
    pub fn sign(&self, message: &[u8]) -> ProtocolKeySignature {
        ProtocolKeySignature::from_inner(self.0.sign(message))
    }

    /// Export the key pair to a Base64 encoded string.
    pub fn to_base64(&self) -> String {
        let secret_bytes = self.0.copy().private().0.to_bytes();
        let mut key_bytes = Vec::new();
        key_bytes.extend_from_slice(&secret_bytes);
        encoding::Base64::encode(&key_bytes)
    }

    /// Import a key pair from a Base64 encoded string.
    pub fn from_base64(encoded: &str) -> Result<Self, FastCryptoError> {
        let key_bytes = encoding::Base64::decode(encoded)?;
        let secret = ed25519::Ed25519PrivateKey::from_bytes(&key_bytes)?;
        let keypair = ed25519::Ed25519KeyPair::from(secret);
        Ok(ProtocolKeyPair(keypair))
    }
}

impl Clone for ProtocolKeyPair {
    fn clone(&self) -> Self {
        Self(self.0.copy())
    }
}

impl Serialize for ProtocolKeyPair {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        serializer.serialize_str(&self.to_base64())
    }
}

impl<'de> Deserialize<'de> for ProtocolKeyPair {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::Deserializer<'de>,
    {
        use serde::de::Error;
        let s = String::deserialize(deserializer)?;
        ProtocolKeyPair::from_base64(&s)
            .map_err(|e| Error::custom(format!("Failed to deserialize ProtocolKeyPair: {e}")))
    }
}

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

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

    /// Deterministic `StdRng` seeds for key generation in this module’s tests.
    const TEST_RNG_SEED_BASE: u64 = 42;

    /// TDD Test: Verify that AuthorityPublicKey bytes roundtrip correctly.
    /// This is critical for hash stability - if from_bytes/to_bytes aren't inverses,
    /// deterministic_hash() will produce different results.
    #[test]
    fn test_authority_key_bytes_roundtrip() {
        // Generate a valid BLS key
        let mut rng = rand::rngs::StdRng::seed_from_u64(TEST_RNG_SEED_BASE);
        let keypair = bls12381::min_sig::BLS12381KeyPair::generate(&mut rng);
        let pubkey = AuthorityPublicKey::new(keypair.public().clone());

        // Get the bytes
        let bytes = pubkey.to_bytes();

        // Restore from bytes
        let restored = AuthorityPublicKey::from_bytes(&bytes).expect("from_bytes should succeed");

        // Verify roundtrip
        assert_eq!(
            pubkey.to_bytes(),
            restored.to_bytes(),
            "Bytes should be identical after roundtrip"
        );
        assert_eq!(pubkey, restored, "Keys should be equal after roundtrip");
    }

    /// TDD Test: Verify that ProtocolPublicKey bytes roundtrip correctly.
    #[test]
    fn test_protocol_key_bytes_roundtrip() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(TEST_RNG_SEED_BASE + 1);
        let keypair = ed25519::Ed25519KeyPair::generate(&mut rng);
        let pubkey = ProtocolPublicKey::new(keypair.public().clone());

        let bytes = pubkey.to_bytes();
        let restored = ProtocolPublicKey::from_bytes(&bytes).expect("from_bytes should succeed");

        assert_eq!(
            pubkey.to_bytes(),
            restored.to_bytes(),
            "Bytes should be identical after roundtrip"
        );
        assert_eq!(pubkey, restored, "Keys should be equal after roundtrip");
    }

    /// TDD Test: Verify that NetworkPublicKey bytes roundtrip correctly.
    #[test]
    fn test_network_key_bytes_roundtrip() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(TEST_RNG_SEED_BASE + 2);
        let keypair = ed25519::Ed25519KeyPair::generate(&mut rng);
        let pubkey = NetworkPublicKey::new(keypair.public().clone());

        let bytes = pubkey.to_bytes();
        let restored = NetworkPublicKey::from_bytes(&bytes).expect("from_bytes should succeed");

        assert_eq!(
            pubkey.to_bytes(),
            restored.to_bytes(),
            "Bytes should be identical after roundtrip"
        );
        assert_eq!(pubkey, restored, "Keys should be equal after roundtrip");
    }

    /// Test that invalid bytes are rejected.
    #[test]
    fn test_authority_key_invalid_bytes() {
        // Invalid length
        let short_bytes = [0u8; 32];
        assert!(AuthorityPublicKey::from_bytes(&short_bytes).is_err());

        // Correct length but invalid point (all zeros is not on the curve)
        let invalid_bytes = [0u8; 96];
        assert!(AuthorityPublicKey::from_bytes(&invalid_bytes).is_err());
    }

    /// Test that new_unique generates different keys.
    #[test]
    fn test_authority_key_new_unique() {
        let key1 = AuthorityPublicKey::new_unique();
        let key2 = AuthorityPublicKey::new_unique();
        assert_ne!(key1, key2, "new_unique should generate different keys");
    }

    /// TDD Test: Verify to_bytes() returns correct size array.
    /// This is critical for deterministic_hash() to produce consistent results.
    #[test]
    fn test_authority_key_bytes_consistency() {
        let key = AuthorityPublicKey::new_unique();

        let bytes = key.to_bytes();

        assert_eq!(bytes.len(), 96, "Authority key should be 96 bytes");

        // Verify roundtrip works
        let restored = AuthorityPublicKey::from_bytes(&bytes).expect("from_bytes should succeed");
        assert_eq!(key, restored, "Roundtrip should preserve key");
    }

    /// TDD Test: Verify protocol key bytes consistency.
    #[test]
    fn test_protocol_key_bytes_consistency() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(TEST_RNG_SEED_BASE + 3);
        let keypair = ed25519::Ed25519KeyPair::generate(&mut rng);
        let key = ProtocolPublicKey::new(keypair.public().clone());

        let bytes = key.to_bytes();
        assert_eq!(bytes.len(), 32, "Protocol key should be 32 bytes");
    }

    /// TDD Test: Verify network key bytes consistency.
    #[test]
    fn test_network_key_bytes_consistency() {
        let mut rng = rand::rngs::StdRng::seed_from_u64(TEST_RNG_SEED_BASE + 4);
        let keypair = ed25519::Ed25519KeyPair::generate(&mut rng);
        let key = NetworkPublicKey::new(keypair.public().clone());

        let bytes = key.to_bytes();
        assert_eq!(bytes.len(), 32, "Network key should be 32 bytes");
    }

    /// Test NetworkKeyPair export/import via base64.
    #[test]
    fn test_network_keypair_export_import() {
        let mut rng = rand::thread_rng();
        let keypair = NetworkKeyPair::generate(&mut rng);

        // Export the key pair
        let exported = keypair.to_base64();

        // Import the key pair back
        let imported = NetworkKeyPair::from_base64(&exported).expect("Failed to import key pair");

        // Check if the imported key pair matches the original
        assert_eq!(keypair.public().to_bytes(), imported.public().to_bytes());
        assert_eq!(keypair.private_key_bytes(), imported.private_key_bytes());
    }

    /// Test NetworkKeyPair invalid base64 import.
    #[test]
    fn test_network_keypair_import_invalid_base64() {
        let result = NetworkKeyPair::from_base64("invalid_base64_string");
        assert!(result.is_err());
    }

    /// Test ProtocolKeyPair export/import via base64.
    #[test]
    fn test_protocol_keypair_export_import() {
        let mut rng = rand::thread_rng();
        let keypair = ProtocolKeyPair::generate(&mut rng);

        // Export the key pair
        let exported = keypair.to_base64();

        // Import the key pair back
        let imported = ProtocolKeyPair::from_base64(&exported).expect("Failed to import key pair");

        // Check if the imported key pair matches the original
        assert_eq!(keypair.public().to_bytes(), imported.public().to_bytes());
    }

    /// Test ProtocolKeyPair signing and verification.
    #[test]
    fn test_protocol_keypair_sign_verify() {
        let mut rng = rand::thread_rng();
        let keypair = ProtocolKeyPair::generate(&mut rng);
        let message = b"test message";

        // Sign the message
        let signature = keypair.sign(message);

        // Verify the signature
        let result = keypair.public().verify(message, &signature);
        assert!(result.is_ok(), "Signature verification should succeed");

        // Verify with wrong message should fail
        let wrong_result = keypair.public().verify(b"wrong message", &signature);
        assert!(
            wrong_result.is_err(),
            "Signature verification should fail with wrong message"
        );
    }
}