1use alloc::vec::Vec;
 2use core::marker::PhantomData;
 3
 4use prover::RpoSignatureProver;
 5use rand_chacha::ChaCha20Rng;
 6use winter_crypto::{ElementHasher, SaltedMerkleTree};
 7use winter_math::fields::f64::BaseElement;
 8use winter_prover::{Proof, ProofOptions, Prover};
 9use winter_utils::Serializable;
10use winter_verifier::{verify, AcceptableOptions, VerifierError};
11
12use crate::{
13    hash::{rpo::Rpo256, DIGEST_SIZE},
14    rand::RpoRandomCoin,
15};
16
17mod air;
18pub use air::{PublicInputs, RescueAir};
19mod prover;
20
21/// Represents an abstract STARK-based signature scheme with knowledge of RPO pre-image as
22/// the hard relation.
23pub struct RpoSignatureScheme<H: ElementHasher> {
24    options: ProofOptions,
25    _h: PhantomData<H>,
26}
27
28impl<H: ElementHasher<BaseField = BaseElement> + Sync> RpoSignatureScheme<H> {
29    pub fn new(options: ProofOptions) -> Self {
30        RpoSignatureScheme { options, _h: PhantomData }
31    }
32
33    pub fn sign(&self, sk: [BaseElement; DIGEST_SIZE], msg: [BaseElement; DIGEST_SIZE]) -> Proof {
34        // create a prover
35        let prover = RpoSignatureProver::<H>::new(msg, self.options.clone());
36
37        // generate execution trace
38        let trace = prover.build_trace(sk);
39
40        // generate the initial seed for the PRNG used for zero-knowledge
41        let seed: [u8; 32] = generate_seed(sk, msg);
42
43        // generate the proof
44        prover.prove(trace, Some(seed)).expect("failed to generate the signature")
45    }
46
47    pub fn verify(
48        &self,
49        pub_key: [BaseElement; DIGEST_SIZE],
50        msg: [BaseElement; DIGEST_SIZE],
51        proof: Proof,
52    ) -> Result<(), VerifierError> {
53        // we make sure that the parameters used in generating the proof match the expected ones
54        if *proof.options() != self.options {
55            return Err(VerifierError::UnacceptableProofOptions);
56        }
57        let pub_inputs = PublicInputs { pub_key, msg };
58        let acceptable_options = AcceptableOptions::OptionSet(vec![proof.options().clone()]);
59        verify::<RescueAir, Rpo256, RpoRandomCoin, SaltedMerkleTree<Rpo256, ChaCha20Rng>>(
60            proof,
61            pub_inputs,
62            &acceptable_options,
63        )
64    }
65}
66
67/// Deterministically generates a seed for seeding the PRNG used for zero-knowledge.
68///
69/// This uses the argument described in [RFC 6979](https://datatracker.ietf.org/doc/html/rfc6979#section-3.5)
70/// § 3.5 where the concatenation of the private key and the hashed message, i.e., sk || H(m), is
71/// used in order to construct the initial seed of a PRNG.
72///
73/// Note that we hash in also a context string in order to domain separate between different
74/// instantiations of the signature scheme.
75#[inline]
76pub fn generate_seed(sk: [BaseElement; DIGEST_SIZE], msg: [BaseElement; DIGEST_SIZE]) -> [u8; 32] {
77    let context_bytes = "
78    Seed for PRNG used for Zero-knowledge in RPO-STARK signature scheme:
79        1. Version: Conjectured security
80        2. FRI queries: 30
81        3. Blowup factor: 8
82        4. Grinding bits: 12
83        5. Field extension degree: 2
84        6. FRI folding factor: 4
85        7. FRI remainder polynomial max degree: 7
86    "
87    .to_bytes();
88    let sk_bytes = sk.to_bytes();
89    let msg_bytes = msg.to_bytes();
90
91    let total_length = context_bytes.len() + sk_bytes.len() + msg_bytes.len();
92    let mut buffer = Vec::with_capacity(total_length);
93    buffer.extend_from_slice(&context_bytes);
94    buffer.extend_from_slice(&sk_bytes);
95    buffer.extend_from_slice(&msg_bytes);
96
97    blake3::hash(&buffer).into()
98}