rialo-types 0.5.0-alpha.0

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

//! Epoch-complete message and attestation threshold types.

use serde::{Deserialize, Serialize};

use super::{
    attestation_framework::{Attestation, INTENT_PREFIX_LENGTH},
    handover::{HandoverAttestation, HandoverRecord},
    headers::RecordKind,
    snapshot::{SnapshotAttestation, SnapshotRecord},
    RecordDigest, BLAKE3_HASH_SIZE, ED25519_SIGNATURE_SIZE,
};
use crate::consensus_crypto::{AuthorityPublicKey, ProtocolPublicKey};

/// Intent prefix for epoch-complete signing messages, sourced from [`RecordKind::EpochComplete`].
///
/// Wire-compatible with `bcs(IntentMessage(Intent::consensus_app(IntentScope::EpochComplete), _))`.
///
/// Guard tests:
/// - `test_epoch_complete_intent_prefix_matches_shared_crypto` (enum value sync, via `headers::tests`)
/// - `test_epoch_complete_signing_message_wire_compatible_with_bcs` (full BCS wire compat)
const EPOCH_COMPLETE_INTENT_PREFIX: [u8; INTENT_PREFIX_LENGTH] =
    RecordKind::EpochComplete.intent_prefix();

/// Size of the signing message for epoch-complete attestations:
/// 3-byte intent prefix + 32-byte handover_digest + 32-byte snapshot_digest = 67 bytes.
pub const EPOCH_COMPLETE_SIGNING_SIZE: usize =
    INTENT_PREFIX_LENGTH + BLAKE3_HASH_SIZE + BLAKE3_HASH_SIZE;

/// A validator's unified attestation that an epoch transition is complete.
///
/// Submitted after a validator observes both a handover and a snapshot for the
/// same epoch transition (e−1 → e). Bundles the two individual attestations
/// with a unified signature over both digests. Submitted as an
/// `AdminTransaction::EpochComplete` variant (added in a later EP)
/// through the *new* epoch's consensus.
///
/// Unlike single-digest `Attestation<R>`, this does NOT implement
/// `AttestableRecord` because it signs two digests (handover + snapshot).
#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
pub struct EpochCompleteMessage {
    /// The validator's handover attestation for this epoch transition.
    pub handover_attestation: HandoverAttestation,

    /// The validator's snapshot attestation for this epoch transition.
    pub snapshot_attestation: SnapshotAttestation,

    /// Ed25519 signature over `EPOCH_COMPLETE_INTENT_PREFIX || handover_digest || snapshot_digest`.
    #[serde(with = "serde_big_array::BigArray")]
    pub unified_signature: [u8; ED25519_SIGNATURE_SIZE],
}

impl EpochCompleteMessage {
    /// Returns the 67-byte signing message:
    /// `EPOCH_COMPLETE_INTENT_PREFIX (3) || handover_digest (32) || snapshot_digest (32)`.
    pub fn signing_message(
        handover_digest: &RecordDigest,
        snapshot_digest: &RecordDigest,
    ) -> [u8; EPOCH_COMPLETE_SIGNING_SIZE] {
        let mut msg = [0u8; EPOCH_COMPLETE_SIGNING_SIZE];
        msg[..INTENT_PREFIX_LENGTH].copy_from_slice(&EPOCH_COMPLETE_INTENT_PREFIX);
        msg[INTENT_PREFIX_LENGTH..INTENT_PREFIX_LENGTH + BLAKE3_HASH_SIZE]
            .copy_from_slice(handover_digest);
        msg[INTENT_PREFIX_LENGTH + BLAKE3_HASH_SIZE..].copy_from_slice(snapshot_digest);
        msg
    }

    /// The validator's authority public key (derived from the handover attestation).
    pub fn authority_key(&self) -> &AuthorityPublicKey {
        self.handover_attestation.authority_key()
    }

    /// The validator's protocol public key (derived from the handover attestation).
    pub fn protocol_key(&self) -> &ProtocolPublicKey {
        self.handover_attestation.protocol_key()
    }

    /// Builds and signs an `EpochCompleteMessage`.
    ///
    /// Signs the handover and snapshot records individually, then produces
    /// the unified 67-byte signature over both digests. Asserts that both
    /// attestations use the same keys and self-verifies the unified signature.
    ///
    /// # Errors
    /// Returns `Err` if any of the three signatures fail self-verification,
    /// or if the attestation keys are inconsistent.
    pub fn sign(
        handover_record: &HandoverRecord,
        snapshot_record: &SnapshotRecord,
        protocol_keypair: &crate::ProtocolKeyPair,
        authority_key: &AuthorityPublicKey,
    ) -> Result<Self, &'static str> {
        let handover_attestation =
            Attestation::<HandoverRecord>::sign(handover_record, protocol_keypair, authority_key)
                .map_err(|_| "EpochCompleteMessage: handover attestation signing failed")?;
        let snapshot_attestation =
            Attestation::<SnapshotRecord>::sign(snapshot_record, protocol_keypair, authority_key)
                .map_err(|_| "EpochCompleteMessage: snapshot attestation signing failed")?;

        if handover_record.dest_epoch() != snapshot_record.epoch() {
            return Err("EpochCompleteMessage: handover dest_epoch does not match snapshot epoch");
        }
        if handover_attestation.authority_key() != snapshot_attestation.authority_key() {
            return Err("EpochCompleteMessage: attestation authority keys are inconsistent");
        }
        if handover_attestation.protocol_key() != snapshot_attestation.protocol_key() {
            return Err("EpochCompleteMessage: attestation protocol keys are inconsistent");
        }

        let handover_digest = handover_record.record_digest();
        let snapshot_digest = snapshot_record.record_digest();
        let msg = Self::signing_message(&handover_digest, &snapshot_digest);

        let sig = protocol_keypair.sign(&msg);
        let mut unified_signature = [0u8; ED25519_SIGNATURE_SIZE];
        unified_signature.copy_from_slice(&sig.to_bytes());

        if protocol_keypair.public().verify(&msg, &sig).is_err() {
            return Err("EpochCompleteMessage unified signature self-verification failed");
        }

        Ok(Self {
            handover_attestation,
            snapshot_attestation,
            unified_signature,
        })
    }

    /// Validates internal consistency:
    /// - handover attestation authority_key == snapshot attestation authority_key
    /// - handover attestation protocol_key == snapshot attestation protocol_key
    /// - handover attestation Ed25519 signature verifies
    /// - snapshot attestation Ed25519 signature verifies
    /// - unified signature verifies against the (shared) protocol key
    pub fn validate_internal_consistency(&self) -> Result<(), &'static str> {
        if self.handover_attestation.authority_key() != self.snapshot_attestation.authority_key() {
            return Err("EpochCompleteMessage: authority_key mismatch between handover and snapshot attestations");
        }
        if self.handover_attestation.protocol_key() != self.snapshot_attestation.protocol_key() {
            return Err("EpochCompleteMessage: protocol_key mismatch between handover and snapshot attestations");
        }
        self.handover_attestation
            .verify_internal_consistency()
            .map_err(|_| {
                "EpochCompleteMessage: handover attestation signature verification failed"
            })?;
        self.snapshot_attestation
            .verify_internal_consistency()
            .map_err(|_| {
                "EpochCompleteMessage: snapshot attestation signature verification failed"
            })?;
        self.verify_unified_signature()
    }

    /// Verifies the unified Ed25519 signature against the embedded protocol key.
    ///
    /// Like [`Attestation::verify_internal_consistency()`], this only proves
    /// internal consistency — callers MUST separately verify that the protocol key
    /// belongs to a committee-registered validator.
    pub fn verify_unified_signature(&self) -> Result<(), &'static str> {
        let handover_digest = self.handover_attestation.record_digest();
        let snapshot_digest = self.snapshot_attestation.record_digest();
        let msg = Self::signing_message(handover_digest, snapshot_digest);
        let sig = crate::ProtocolKeySignature::from_bytes(&self.unified_signature)
            .map_err(|_| "EpochCompleteMessage: invalid unified signature bytes")?;
        self.protocol_key()
            .verify(&msg, &sig)
            .map_err(|_| "EpochCompleteMessage: unified signature verification failed")
    }
}

impl std::fmt::Display for EpochCompleteMessage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "EpochCompleteMessage(authority={}…, handover_digest={}…, snapshot_digest={}…)",
            hex::encode(&self.authority_key().to_bytes()[..4]),
            hex::encode(&self.handover_attestation.record_digest()[..4]),
            hex::encode(&self.snapshot_attestation.record_digest()[..4]),
        )
    }
}

#[cfg(test)]
mod tests {
    use super::{
        super::test_helpers::{test_keys::*, test_records::*},
        *,
    };

    fn create_test_epoch_complete_message(
        seed: u64,
    ) -> (EpochCompleteMessage, crate::ProtocolKeyPair) {
        let (authority_key, protocol_keypair, _) = create_test_keys_with_keypair(seed);
        let handover_record = create_test_handover_record(100, 0, 1, [0u8; 32]);
        let snapshot_record = create_test_snapshot_record(100, 1);
        let msg = EpochCompleteMessage::sign(
            &handover_record,
            &snapshot_record,
            &protocol_keypair,
            &authority_key,
        )
        .expect("EpochCompleteMessage::sign() should succeed");
        (msg, protocol_keypair)
    }

    #[test]
    fn test_epoch_complete_intent_prefix_matches_shared_crypto() {
        use shared_crypto::intent::{AppId, IntentScope, IntentVersion};

        assert_eq!(
            EPOCH_COMPLETE_INTENT_PREFIX,
            [
                IntentScope::EpochComplete as u8,
                IntentVersion::V0 as u8,
                AppId::Consensus as u8,
            ],
            "EPOCH_COMPLETE_INTENT_PREFIX is out of sync with shared-crypto enums"
        );
    }

    #[test]
    fn test_epoch_complete_signing_message_wire_compatible_with_bcs() {
        use shared_crypto::intent::{Intent, IntentMessage, IntentScope};

        let handover_digest = [0xAAu8; BLAKE3_HASH_SIZE];
        let snapshot_digest = [0xBBu8; BLAKE3_HASH_SIZE];

        let msg = EpochCompleteMessage::signing_message(&handover_digest, &snapshot_digest);

        // BCS wire compat: the first 3 bytes must match BCS-serialized intent prefix.
        let intent_msg = IntentMessage::new(
            Intent::consensus_app(IntentScope::EpochComplete),
            handover_digest,
        );
        let bcs_bytes = bcs::to_bytes(&intent_msg).expect("BCS serialization should succeed");

        assert_eq!(
            msg[..INTENT_PREFIX_LENGTH],
            bcs_bytes[..INTENT_PREFIX_LENGTH],
            "intent prefix bytes must be byte-identical to bcs(IntentMessage(EpochComplete, _))"
        );
    }

    #[test]
    fn test_epoch_complete_signing_message_format_and_size() {
        let handover_digest = [0xAAu8; BLAKE3_HASH_SIZE];
        let snapshot_digest = [0xBBu8; BLAKE3_HASH_SIZE];

        let msg = EpochCompleteMessage::signing_message(&handover_digest, &snapshot_digest);

        assert_eq!(msg.len(), EPOCH_COMPLETE_SIGNING_SIZE);

        assert_eq!(msg[..INTENT_PREFIX_LENGTH], EPOCH_COMPLETE_INTENT_PREFIX);
        assert_eq!(
            msg[INTENT_PREFIX_LENGTH..INTENT_PREFIX_LENGTH + BLAKE3_HASH_SIZE],
            handover_digest
        );
        assert_eq!(
            msg[INTENT_PREFIX_LENGTH + BLAKE3_HASH_SIZE..],
            snapshot_digest
        );
    }

    #[test]
    fn test_epoch_complete_validate_internal_consistency_success() {
        let (msg, _) = create_test_epoch_complete_message(5001);
        assert!(msg.validate_internal_consistency().is_ok());
    }

    #[test]
    fn test_epoch_complete_validate_internal_consistency_rejects_tampered_signature() {
        let (mut msg, _) = create_test_epoch_complete_message(5002);
        msg.unified_signature[0] ^= 0xFF;
        assert!(msg.validate_internal_consistency().is_err());
    }

    #[test]
    fn test_validate_internal_consistency_rejects_tampered_individual_attestation() {
        let (mut msg, _) = create_test_epoch_complete_message(5050);
        msg.handover_attestation.tamper_signature(0, 0xFF);
        assert!(msg.validate_internal_consistency().is_err());
    }

    #[test]
    fn test_epoch_complete_message_display() {
        let (msg, _) = create_test_epoch_complete_message(5030);
        let display = format!("{msg}");
        assert!(display.contains("EpochCompleteMessage"));
        assert!(display.contains("authority="));
        assert!(display.contains("handover_digest="));
        assert!(display.contains("snapshot_digest="));
    }

    #[test]
    fn test_epoch_complete_sign_verify_roundtrip() {
        let (msg, _) = create_test_epoch_complete_message(5040);
        assert!(msg.verify_unified_signature().is_ok());
        assert!(msg.validate_internal_consistency().is_ok());
        msg.handover_attestation
            .verify_internal_consistency()
            .expect("handover attestation should verify");
        msg.snapshot_attestation
            .verify_internal_consistency()
            .expect("snapshot attestation should verify");
    }

    #[test]
    fn test_epoch_complete_verify_rejects_tampered_unified_signature() {
        let (mut msg, _) = create_test_epoch_complete_message(5041);
        msg.unified_signature[0] ^= 0xFF;
        assert!(msg.verify_unified_signature().is_err());
    }

    #[test]
    fn test_epoch_complete_sign_rejects_epoch_mismatch() {
        let (authority_key, protocol_keypair, _) = create_test_keys_with_keypair(5060);
        let handover_record = create_test_handover_record(100, 0, 1, [0u8; 32]);
        let snapshot_record = create_test_snapshot_record(100, 5);
        let result = EpochCompleteMessage::sign(
            &handover_record,
            &snapshot_record,
            &protocol_keypair,
            &authority_key,
        );
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .contains("dest_epoch does not match snapshot epoch"));
    }
}