1use core::marker::PhantomData;
 2
 3use prover::RpoSignatureProver;
 4use rand::{distributions::Standard, prelude::Distribution, thread_rng, RngCore, SeedableRng};
 5use rand_chacha::ChaCha20Rng;
 6use winter_crypto::{ElementHasher, Hasher, SaltedMerkleTree};
 7use winter_math::fields::f64::BaseElement;
 8use winter_prover::{Proof, ProofOptions, Prover};
 9use winter_verifier::{verify, AcceptableOptions, VerifierError};
10
11use crate::{
12    hash::{rpo::Rpo256, DIGEST_SIZE},
13    rand::RpoRandomCoin,
14};
15
16mod air;
17pub use air::{PublicInputs, RescueAir};
18mod prover;
19
20/// Represents an abstract STARK-based signature scheme with knowledge of RPO pre-image as
21/// the hard relation.
22pub struct RpoSignatureScheme<H: ElementHasher> {
23    options: ProofOptions,
24    _h: PhantomData<H>,
25}
26
27impl<H: ElementHasher<BaseField = BaseElement> + Sync> RpoSignatureScheme<H>
28where
29    Standard: Distribution<<H as Hasher>::Digest>,
30{
31    pub fn new(options: ProofOptions) -> Self {
32        RpoSignatureScheme { options, _h: PhantomData }
33    }
34
35    pub fn sign(&self, sk: [BaseElement; DIGEST_SIZE], msg: [BaseElement; DIGEST_SIZE]) -> Proof {
36        // create a prover
37        let prover = RpoSignatureProver::<H>::new(msg, self.options.clone());
38
39        // generate execution trace
40        let trace = prover.build_trace(sk);
41
42        // generate the initial seed for the PRNG used for zero-knowledge
43        let mut seed = <ChaCha20Rng as SeedableRng>::Seed::default();
44        let mut rng = thread_rng();
45        rng.fill_bytes(&mut seed);
46
47        // generate the proof
48        prover.prove(trace, Some(seed)).expect("failed to generate the signature")
49    }
50
51    pub fn verify(
52        &self,
53        pub_key: [BaseElement; DIGEST_SIZE],
54        msg: [BaseElement; DIGEST_SIZE],
55        proof: Proof,
56    ) -> Result<(), VerifierError> {
57        // we make sure that the parameters used in generating the proof match the expected ones
58        if *proof.options() != self.options {
59            return Err(VerifierError::UnacceptableProofOptions);
60        }
61        let pub_inputs = PublicInputs { pub_key, msg };
62        let acceptable_options = AcceptableOptions::OptionSet(vec![proof.options().clone()]);
63        verify::<RescueAir, Rpo256, RpoRandomCoin, SaltedMerkleTree<Rpo256, ChaCha20Rng>>(
64            proof,
65            pub_inputs,
66            &acceptable_options,
67        )
68    }
69}