arcium-primitives 0.4.2

use crate::{
    random::{CryptoRngCore, Seed},
    types::{identifiers::ProtocolInfo, SessionId},
};

pub mod folklore;
pub mod merlin;
pub mod protocol_transcript;

pub use folklore::FolkloreTranscript;
pub use merlin::MerlinTranscript;
pub use protocol_transcript::ProtocolTranscript;

/// A trait for succinct transcripts (via some form of hashing) used in cryptographic protocols.
/// The transcript should be able to:
/// - Absorb messages with associated labels.
/// - Produce pseudorandom outputs (e.g., via a CSPRNG) based on the absorbed messages.
/// - Ensure that the order and content of messages affect the outputs, providing domain separation.
pub trait Transcript {
    /// The type of the RNG derived from the transcript.
    type Rng: CryptoRngCore;

    /// Create a new transcript for a specific protocol with a session ID for domain separation.
    fn new(protocol_info: &'static ProtocolInfo, session_id: &SessionId) -> Self;

    /// Append a message with a label to the transcript.
    fn append_with<S: AsRef<[u8]>>(&mut self, label: &'static [u8], message: &S);

    /// Append multiple messages with a common label to the transcript.
    fn append_many_with<S: AsRef<[u8]>>(&mut self, label: &'static [u8], messages: &[S]);

    /// Extract pseudorandom bytes based on the transcript state.
    /// Note: two consecutive extractions must yield different outputs.
    fn extract(&mut self, label: &'static [u8]) -> Seed;

    /// Derive a CSPRNG from the transcript state with a specific label. Allows
    /// arbitrary-length output generation
    fn extract_rng(&mut self, label: &'static [u8]) -> Self::Rng;
}

/// A trait for transcripts that do not use labels when appending messages.
/// WARNING! This should only be implemented by transcripts that are designed
/// to automatically manage the labels internally at runtime.
pub trait AutoTranscript: Transcript + Sized {
    /// The type of the label used for messages in this transcript.
    type Label: AsRef<[u8]>;

    #[inline]
    /// Create a new transcript for a specific protocol with a session ID for domain separation.
    fn new(protocol_info: &'static ProtocolInfo, session_id: &SessionId) -> Self {
        Transcript::new(protocol_info, session_id)
    }

    /// Get the label that was used for the last appended message.
    /// If no messages have been appended yet, returns the initial label.
    fn get_current_label(&self) -> Self::Label;

    /// Get the label that will be used for the next appended message.
    fn next_label(&mut self) -> Self::Label;

    /// Append a message to the transcript with an automatically generated label.
    fn append<S: AsRef<[u8]>>(&mut self, message: &S);

    /// Append multiple messages to the transcript with a common automatically generated label.
    fn append_many<S: AsRef<[u8]>>(&mut self, values: &[S]);

    #[inline]
    /// Extract pseudorandom bytes based on the transcript state.
    /// Note: two consecutive extractions must yield different outputs.
    fn extract(&mut self, label: &'static [u8]) -> Seed {
        Transcript::extract(self, label)
    }

    #[inline]
    /// Derive a CSPRNG from the transcript state with a specific label. Allows
    /// arbitrary-length output generation
    fn extract_rng(&mut self, label: &'static [u8]) -> Self::Rng {
        Transcript::extract_rng(self, label)
    }
}

#[cfg(test)]
mod tests {
    use rand::RngCore;

    use crate::{
        random::{test_rng, Random},
        transcripts::Transcript,
        types::{identifiers::ProtocolInfo, SessionId},
    };

    // Appending a message and extracting it in two identical transcripts should
    // yield the same result.
    pub fn test_happypath_same_transcripts_yield_same_extracts<T: Transcript>() {
        let message = b"hello world";
        let session_id = SessionId::random(&mut test_rng());

        static TEST_PROTOCOL_INFO: ProtocolInfo = ProtocolInfo::new("test_protocol", 42);
        let mut t1 = T::new(&TEST_PROTOCOL_INFO, &session_id);
        t1.append_with(b"message_1", &message);
        let mut t2 = T::new(&TEST_PROTOCOL_INFO, &session_id);
        t2.append_with(b"message_1", &message);

        let dest1 = t1.extract(b"value_1");
        let dest2 = t2.extract(b"value_1");
        assert_eq!(dest1, dest2);

        let mut bytes1 = [0u8; 100];
        let mut bytes2 = [0u8; 100];
        t1.extract_rng(b"stream").fill_bytes(&mut bytes1);
        t2.extract_rng(b"stream").fill_bytes(&mut bytes2);
        assert_eq!(bytes1, bytes2);
    }

    // Appending a message and extracting it in two different transcripts should
    // yield different results.
    pub fn test_unhappypath_different_transcripts_yield_different_extracts<T: Transcript>() {
        // Two transcripts that differ only in their init labels.
        let message = b"hello world";
        let session_id = SessionId::random(&mut test_rng());

        static PROTOCOL_1: ProtocolInfo = ProtocolInfo::new("test_protocol_1", 42);
        let mut t1 = T::new(&PROTOCOL_1, &session_id);
        t1.append_with(b"message_1", &message);

        static PROTOCOL_2: ProtocolInfo = ProtocolInfo::new("test_protocol_2", 43);
        let mut t2 = T::new(&PROTOCOL_2, &session_id);
        t2.append_with(b"message_1", &message);

        let dest1 = t1.extract(b"value_1");
        let dest2 = t2.extract(b"value_1");
        assert_ne!(dest1, dest2);

        let mut bytes1 = [0u8; 100];
        let mut bytes2 = [0u8; 100];
        t1.extract_rng(b"stream").fill_bytes(&mut bytes1);
        t2.extract_rng(b"stream").fill_bytes(&mut bytes2);
        assert_ne!(bytes1, bytes2);

        // Two transcripts that differ in their message ordering.
        let message1 = b"hello world";
        let message2 = b"goodbye world";

        let mut t1 = T::new(&PROTOCOL_1, &session_id);
        t1.append_with(b"message_1", &message1);
        t1.append_with(b"message_2", &message2);
        let mut t2 = T::new(&PROTOCOL_1, &session_id);
        t2.append_with(b"message_2", &message2);
        t2.append_with(b"message_1", &message1);

        let dest1 = t1.extract(b"value_1");
        let dest2 = t2.extract(b"value_1");
        assert_ne!(dest1, dest2);

        let mut bytes1 = [0u8; 100];
        let mut bytes2 = [0u8; 100];
        t1.extract_rng(b"stream").fill_bytes(&mut bytes1);
        t2.extract_rng(b"stream").fill_bytes(&mut bytes2);
        assert_ne!(bytes1, bytes2);
    }

    #[test]
    fn test_folklore_transcript() {
        test_happypath_same_transcripts_yield_same_extracts::<super::folklore::FolkloreTranscript>(
        );
        test_unhappypath_different_transcripts_yield_different_extracts::<
            super::folklore::FolkloreTranscript,
        >();
    }

    #[test]
    fn test_merlin_transcript() {
        test_happypath_same_transcripts_yield_same_extracts::<super::merlin::MerlinTranscript>();
        test_unhappypath_different_transcripts_yield_different_extracts::<
            super::merlin::MerlinTranscript,
        >();
    }
}