feat: implement RPO STARK-based signature DSA (with zero knowledge) (#349)

src/dsa/mod.rs

@@ -1,3 +1,5 @@
 //! Digital signature schemes supported by default in the Miden VM.
 
 pub mod rpo_falcon512;
+
+pub mod rpo_stark;

src/dsa/rpo_stark/mod.rs

@@ -0,0 +1,24 @@
+mod signature;
+pub use signature::{PublicKey, SecretKey, Signature};
+
+mod stark;
+pub use stark::{PublicInputs, RescueAir};
+
+// TESTS
+// ================================================================================================
+
+#[cfg(test)]
+mod tests {
+    use super::SecretKey;
+    use crate::Word;
+
+    #[test]
+    fn test_signature() {
+        let sk = SecretKey::new(Word::default());
+
+        let message = Word::default();
+        let signature = sk.sign(message);
+        let pk = sk.public_key();
+        assert!(pk.verify(message, &signature))
+    }
+}

src/dsa/rpo_stark/signature/mod.rs

@@ -0,0 +1,173 @@
+use rand::{distributions::Uniform, prelude::Distribution, Rng};
+use winter_air::{FieldExtension, ProofOptions};
+use winter_math::{fields::f64::BaseElement, FieldElement};
+use winter_prover::Proof;
+use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
+
+use crate::{
+    dsa::rpo_stark::stark::RpoSignatureScheme,
+    hash::{rpo::Rpo256, DIGEST_SIZE},
+    StarkField, Word, ZERO,
+};
+
+// CONSTANTS
+// ================================================================================================
+
+/// Specifies the parameters of the STARK underlying the signature scheme. These parameters provide
+/// at least 102 bits of security under the conjectured security of the toy protocol in
+/// the ethSTARK paper [1].
+///
+/// [1]: https://eprint.iacr.org/2021/582
+pub const PROOF_OPTIONS: ProofOptions =
+    ProofOptions::new(30, 8, 12, FieldExtension::Quadratic, 4, 7, true);
+
+// PUBLIC KEY
+// ================================================================================================
+
+/// A public key for verifying signatures.
+///
+/// The public key is a [Word] (i.e., 4 field elements) that is the hash of the secret key.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct PublicKey(Word);
+
+impl PublicKey {
+    /// Returns the [Word] defining the public key.
+    pub fn inner(&self) -> Word {
+        self.0
+    }
+}
+
+impl PublicKey {
+    /// Verifies the provided signature against provided message and this public key.
+    pub fn verify(&self, message: Word, signature: &Signature) -> bool {
+        signature.verify(message, *self)
+    }
+}
+
+impl Serializable for PublicKey {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        self.0.write_into(target);
+    }
+}
+
+impl Deserializable for PublicKey {
+    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
+        let pk = <Word>::read_from(source)?;
+        Ok(Self(pk))
+    }
+}
+
+// SECRET KEY
+// ================================================================================================
+
+/// A secret key for generating signatures.
+///
+/// The secret key is a [Word] (i.e., 4 field elements).
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub struct SecretKey(Word);
+
+impl SecretKey {
+    /// Generates a secret key from OS-provided randomness.
+    pub fn new(word: Word) -> Self {
+        Self(word)
+    }
+
+    /// Generates a secret key from a [Word].
+    #[cfg(feature = "std")]
+    pub fn random() -> Self {
+        use rand::{rngs::StdRng, SeedableRng};
+
+        let mut rng = StdRng::from_entropy();
+        Self::with_rng(&mut rng)
+    }
+
+    /// Generates a secret_key using the provided random number generator `Rng`.
+    pub fn with_rng<R: Rng>(rng: &mut R) -> Self {
+        let mut sk = [ZERO; 4];
+        let uni_dist = Uniform::from(0..BaseElement::MODULUS);
+
+        for s in sk.iter_mut() {
+            let sampled_integer = uni_dist.sample(rng);
+            *s = BaseElement::new(sampled_integer);
+        }
+
+        Self(sk)
+    }
+
+    /// Computes the public key corresponding to this secret key.
+    pub fn public_key(&self) -> PublicKey {
+        let mut elements = [BaseElement::ZERO; 8];
+        elements[..DIGEST_SIZE].copy_from_slice(&self.0);
+        let pk = Rpo256::hash_elements(&elements);
+        PublicKey(pk.into())
+    }
+
+    /// Signs a message with this secret key.
+    pub fn sign(&self, message: Word) -> Signature {
+        let signature: RpoSignatureScheme<Rpo256> = RpoSignatureScheme::new(PROOF_OPTIONS);
+        let proof = signature.sign(self.0, message);
+        Signature { proof }
+    }
+}
+
+impl Serializable for SecretKey {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        self.0.write_into(target);
+    }
+}
+
+impl Deserializable for SecretKey {
+    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
+        let sk = <Word>::read_from(source)?;
+        Ok(Self(sk))
+    }
+}
+
+// SIGNATURE
+// ================================================================================================
+
+/// An RPO STARK-based signature over a message.
+///
+/// The signature is a STARK proof of knowledge of a pre-image given an image where the map is
+/// the RPO permutation, the pre-image is the secret key and the image is the public key.
+/// The current implementation follows the description in [1] but relies on the conjectured security
+/// of the toy protocol in the ethSTARK paper [2], which gives us using the parameter set
+/// given in `PROOF_OPTIONS` a signature with $102$ bits of average-case existential unforgeability
+/// security against $2^{113}$-query bound adversaries that can obtain up to $2^{64}$ signatures
+/// under the same public key.
+///
+/// [1]: https://eprint.iacr.org/2024/1553
+/// [2]: https://eprint.iacr.org/2021/582
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct Signature {
+    proof: Proof,
+}
+
+impl Signature {
+    /// Returns the STARK proof constituting the signature.
+    pub fn inner(&self) -> Proof {
+        self.proof.clone()
+    }
+
+    /// Returns true if this signature is a valid signature for the specified message generated
+    /// against the secret key matching the specified public key.
+    pub fn verify(&self, message: Word, pk: PublicKey) -> bool {
+        let signature: RpoSignatureScheme<Rpo256> = RpoSignatureScheme::new(PROOF_OPTIONS);
+
+        let res = signature.verify(pk.inner(), message, self.proof.clone());
+        res.is_ok()
+    }
+}
+
+impl Serializable for Signature {
+    fn write_into<W: ByteWriter>(&self, target: &mut W) {
+        self.proof.write_into(target);
+    }
+}
+
+impl Deserializable for Signature {
+    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
+        let proof = Proof::read_from(source)?;
+        Ok(Self { proof })
+    }
+}

src/dsa/rpo_stark/stark/air.rs

@@ -0,0 +1,198 @@
+use alloc::vec::Vec;
+
+use winter_math::{fields::f64::BaseElement, FieldElement, ToElements};
+use winter_prover::{
+    Air, AirContext, Assertion, EvaluationFrame, ProofOptions, TraceInfo,
+    TransitionConstraintDegree,
+};
+
+use crate::{
+    hash::{ARK1, ARK2, MDS, STATE_WIDTH},
+    Word, ZERO,
+};
+
+// CONSTANTS
+// ================================================================================================
+
+pub const HASH_CYCLE_LEN: usize = 8;
+
+// AIR
+// ================================================================================================
+
+pub struct RescueAir {
+    context: AirContext<BaseElement>,
+    pub_key: Word,
+}
+
+impl Air for RescueAir {
+    type BaseField = BaseElement;
+    type PublicInputs = PublicInputs;
+
+    type GkrProof = ();
+    type GkrVerifier = ();
+
+    // CONSTRUCTOR
+    // --------------------------------------------------------------------------------------------
+    fn new(trace_info: TraceInfo, pub_inputs: PublicInputs, options: ProofOptions) -> Self {
+        let degrees = vec![
+            // Apply RPO rounds.
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+            TransitionConstraintDegree::new(7),
+        ];
+        assert_eq!(STATE_WIDTH, trace_info.width());
+        let context = AirContext::new(trace_info, degrees, 12, options);
+        let context = context.set_num_transition_exemptions(1);
+        RescueAir { context, pub_key: pub_inputs.pub_key }
+    }
+
+    fn context(&self) -> &AirContext<Self::BaseField> {
+        &self.context
+    }
+
+    fn evaluate_transition<E: FieldElement + From<Self::BaseField>>(
+        &self,
+        frame: &EvaluationFrame<E>,
+        periodic_values: &[E],
+        result: &mut [E],
+    ) {
+        let current = frame.current();
+        let next = frame.next();
+        // expected state width is 12 field elements
+        debug_assert_eq!(STATE_WIDTH, current.len());
+        debug_assert_eq!(STATE_WIDTH, next.len());
+
+        enforce_rpo_round(frame, result, periodic_values);
+    }
+
+    fn get_assertions(&self) -> Vec<Assertion<Self::BaseField>> {
+        let initial_step = 0;
+        let last_step = self.trace_length() - 1;
+        vec![
+            // Assert that the capacity as well as the second half of the rate portion of the state
+            // are initialized to `ZERO`.The first half of the rate is unconstrained as it will
+            // contain the secret key
+            Assertion::single(0, initial_step, Self::BaseField::ZERO),
+            Assertion::single(1, initial_step, Self::BaseField::ZERO),
+            Assertion::single(2, initial_step, Self::BaseField::ZERO),
+            Assertion::single(3, initial_step, Self::BaseField::ZERO),
+            Assertion::single(8, initial_step, Self::BaseField::ZERO),
+            Assertion::single(9, initial_step, Self::BaseField::ZERO),
+            Assertion::single(10, initial_step, Self::BaseField::ZERO),
+            Assertion::single(11, initial_step, Self::BaseField::ZERO),
+            // Assert that the public key is the correct one
+            Assertion::single(4, last_step, self.pub_key[0]),
+            Assertion::single(5, last_step, self.pub_key[1]),
+            Assertion::single(6, last_step, self.pub_key[2]),
+            Assertion::single(7, last_step, self.pub_key[3]),
+        ]
+    }
+
+    fn get_periodic_column_values(&self) -> Vec<Vec<Self::BaseField>> {
+        get_round_constants()
+    }
+}
+
+pub struct PublicInputs {
+    pub(crate) pub_key: Word,
+    pub(crate) msg: Word,
+}
+
+impl PublicInputs {
+    pub fn new(pub_key: Word, msg: Word) -> Self {
+        Self { pub_key, msg }
+    }
+}
+
+impl ToElements<BaseElement> for PublicInputs {
+    fn to_elements(&self) -> Vec<BaseElement> {
+        let mut res = self.pub_key.to_vec();
+        res.extend_from_slice(self.msg.as_ref());
+        res
+    }
+}
+
+// HELPER EVALUATORS
+// ------------------------------------------------------------------------------------------------
+
+/// Enforces constraints for a single round of the Rescue Prime Optimized hash functions.
+pub fn enforce_rpo_round<E: FieldElement + From<BaseElement>>(
+    frame: &EvaluationFrame<E>,
+    result: &mut [E],
+    ark: &[E],
+) {
+    // compute the state that should result from applying the first 5 operations of the RPO round to
+    // the current hash state.
+    let mut step1 = [E::ZERO; STATE_WIDTH];
+    step1.copy_from_slice(frame.current());
+
+    apply_mds(&mut step1);
+    // add constants
+    for i in 0..STATE_WIDTH {
+        step1[i] += ark[i];
+    }
+    apply_sbox(&mut step1);
+    apply_mds(&mut step1);
+    // add constants
+    for i in 0..STATE_WIDTH {
+        step1[i] += ark[STATE_WIDTH + i];
+    }
+
+    // compute the state that should result from applying the inverse of the last operation of the
+    // RPO round to the next step of the computation.
+    let mut step2 = [E::ZERO; STATE_WIDTH];
+    step2.copy_from_slice(frame.next());
+    apply_sbox(&mut step2);
+
+    // make sure that the results are equal.
+    for i in 0..STATE_WIDTH {
+        result[i] = step2[i] - step1[i]
+    }
+}
+
+#[inline(always)]
+fn apply_sbox<E: FieldElement + From<BaseElement>>(state: &mut [E; STATE_WIDTH]) {
+    state.iter_mut().for_each(|v| {
+        let t2 = v.square();
+        let t4 = t2.square();
+        *v *= t2 * t4;
+    });
+}
+
+#[inline(always)]
+fn apply_mds<E: FieldElement + From<BaseElement>>(state: &mut [E; STATE_WIDTH]) {
+    let mut result = [E::ZERO; STATE_WIDTH];
+    result.iter_mut().zip(MDS).for_each(|(r, mds_row)| {
+        state.iter().zip(mds_row).for_each(|(&s, m)| {
+            *r += E::from(m) * s;
+        });
+    });
+    *state = result
+}
+
+/// Returns RPO round constants arranged in column-major form.
+pub fn get_round_constants() -> Vec<Vec<BaseElement>> {
+    let mut constants = Vec::new();
+    for _ in 0..(STATE_WIDTH * 2) {
+        constants.push(vec![ZERO; HASH_CYCLE_LEN]);
+    }
+
+    #[allow(clippy::needless_range_loop)]
+    for i in 0..HASH_CYCLE_LEN - 1 {
+        for j in 0..STATE_WIDTH {
+            constants[j][i] = ARK1[i][j];
+            constants[j + STATE_WIDTH][i] = ARK2[i][j];
+        }
+    }
+
+    constants
+}

src/dsa/rpo_stark/stark/mod.rs

@@ -0,0 +1,69 @@
+use core::marker::PhantomData;
+
+use prover::RpoSignatureProver;
+use rand::{distributions::Standard, prelude::Distribution, thread_rng, RngCore, SeedableRng};
+use rand_chacha::ChaCha20Rng;
+use winter_crypto::{ElementHasher, Hasher, SaltedMerkleTree};
+use winter_math::fields::f64::BaseElement;
+use winter_prover::{Proof, ProofOptions, Prover};
+use winter_verifier::{verify, AcceptableOptions, VerifierError};
+
+use crate::{
+    hash::{rpo::Rpo256, DIGEST_SIZE},
+    rand::RpoRandomCoin,
+};
+
+mod air;
+pub use air::{PublicInputs, RescueAir};
+mod prover;
+
+/// Represents an abstract STARK-based signature scheme with knowledge of RPO pre-image as
+/// the hard relation.
+pub struct RpoSignatureScheme<H: ElementHasher> {
+    options: ProofOptions,
+    _h: PhantomData<H>,
+}
+
+impl<H: ElementHasher<BaseField = BaseElement> + Sync> RpoSignatureScheme<H>
+where
+    Standard: Distribution<<H as Hasher>::Digest>,
+{
+    pub fn new(options: ProofOptions) -> Self {
+        RpoSignatureScheme { options, _h: PhantomData }
+    }
+
+    pub fn sign(&self, sk: [BaseElement; DIGEST_SIZE], msg: [BaseElement; DIGEST_SIZE]) -> Proof {
+        // create a prover
+        let prover = RpoSignatureProver::<H>::new(msg, self.options.clone());
+
+        // generate execution trace
+        let trace = prover.build_trace(sk);
+
+        // generate the initial seed for the PRNG used for zero-knowledge
+        let mut seed = <ChaCha20Rng as SeedableRng>::Seed::default();
+        let mut rng = thread_rng();
+        rng.fill_bytes(&mut seed);
+
+        // generate the proof
+        prover.prove(trace, Some(seed)).expect("failed to generate the signature")
+    }
+
+    pub fn verify(
+        &self,
+        pub_key: [BaseElement; DIGEST_SIZE],
+        msg: [BaseElement; DIGEST_SIZE],
+        proof: Proof,
+    ) -> Result<(), VerifierError> {
+        // we make sure that the parameters used in generating the proof match the expected ones
+        if *proof.options() != self.options {
+            return Err(VerifierError::UnacceptableProofOptions);
+        }
+        let pub_inputs = PublicInputs { pub_key, msg };
+        let acceptable_options = AcceptableOptions::OptionSet(vec![proof.options().clone()]);
+        verify::<RescueAir, Rpo256, RpoRandomCoin, SaltedMerkleTree<Rpo256, ChaCha20Rng>>(
+            proof,
+            pub_inputs,
+            &acceptable_options,
+        )
+    }
+}

src/dsa/rpo_stark/stark/prover.rs

@@ -0,0 +1,148 @@
+use core::marker::PhantomData;
+
+use rand_chacha::ChaCha20Rng;
+use winter_air::{
+    AuxRandElements, ConstraintCompositionCoefficients, PartitionOptions, ZkParameters,
+};
+use winter_crypto::{ElementHasher, SaltedMerkleTree};
+use winter_math::{fields::f64::BaseElement, FieldElement};
+use winter_prover::{
+    matrix::ColMatrix, CompositionPoly, CompositionPolyTrace, DefaultConstraintCommitment,
+    DefaultConstraintEvaluator, DefaultTraceLde, ProofOptions, Prover, StarkDomain, Trace,
+    TraceInfo, TracePolyTable, TraceTable,
+};
+
+use super::air::{PublicInputs, RescueAir, HASH_CYCLE_LEN};
+use crate::{
+    hash::{rpo::Rpo256, STATE_WIDTH},
+    rand::RpoRandomCoin,
+    Word, ZERO,
+};
+
+// PROVER
+// ================================================================================================
+
+/// A prover for the RPO STARK-based signature scheme.
+///
+/// The signature is based on the the one-wayness of the RPO hash function but it is generic over
+/// the hash function used for instantiating the random oracle for the BCS transform.
+pub(crate) struct RpoSignatureProver<H: ElementHasher + Sync> {
+    message: Word,
+    options: ProofOptions,
+    _hasher: PhantomData<H>,
+}
+
+impl<H: ElementHasher + Sync> RpoSignatureProver<H> {
+    pub(crate) fn new(message: Word, options: ProofOptions) -> Self {
+        Self { message, options, _hasher: PhantomData }
+    }
+
+    pub(crate) fn build_trace(&self, sk: Word) -> TraceTable<BaseElement> {
+        let mut trace = TraceTable::new(STATE_WIDTH, HASH_CYCLE_LEN);
+
+        trace.fill(
+            |state| {
+                // initialize first half of the rate portion of the state with the secret key
+                state[0] = ZERO;
+                state[1] = ZERO;
+                state[2] = ZERO;
+                state[3] = ZERO;
+                state[4] = sk[0];
+                state[5] = sk[1];
+                state[6] = sk[2];
+                state[7] = sk[3];
+                state[8] = ZERO;
+                state[9] = ZERO;
+                state[10] = ZERO;
+                state[11] = ZERO;
+            },
+            |step, state| {
+                Rpo256::apply_round(
+                    state.try_into().expect("should not fail given the size of the array"),
+                    step,
+                );
+            },
+        );
+        trace
+    }
+}
+
+impl<H: ElementHasher> Prover for RpoSignatureProver<H>
+where
+    H: ElementHasher<BaseField = BaseElement> + Sync,
+{
+    type BaseField = BaseElement;
+    type Air = RescueAir;
+    type Trace = TraceTable<BaseElement>;
+    type HashFn = Rpo256;
+    type VC = SaltedMerkleTree<Self::HashFn, Self::ZkPrng>;
+    type RandomCoin = RpoRandomCoin;
+    type TraceLde<E: FieldElement<BaseField = Self::BaseField>> =
+        DefaultTraceLde<E, Self::HashFn, Self::VC>;
+    type ConstraintCommitment<E: FieldElement<BaseField = Self::BaseField>> =
+        DefaultConstraintCommitment<E, Self::HashFn, Self::ZkPrng, Self::VC>;
+    type ConstraintEvaluator<'a, E: FieldElement<BaseField = Self::BaseField>> =
+        DefaultConstraintEvaluator<'a, Self::Air, E>;
+    type ZkPrng = ChaCha20Rng;
+
+    fn get_pub_inputs(&self, trace: &Self::Trace) -> PublicInputs {
+        let last_step = trace.length() - 1;
+        // Note that the message is not part of the execution trace but is part of the public
+        // inputs. This is explained in the reference description of the DSA and intuitively
+        // it is done in order to make sure that the message is part of the Fiat-Shamir
+        // transcript and hence binds the proof/signature to the message
+        PublicInputs {
+            pub_key: [
+                trace.get(4, last_step),
+                trace.get(5, last_step),
+                trace.get(6, last_step),
+                trace.get(7, last_step),
+            ],
+            msg: self.message,
+        }
+    }
+
+    fn options(&self) -> &ProofOptions {
+        &self.options
+    }
+
+    fn new_trace_lde<E: FieldElement<BaseField = Self::BaseField>>(
+        &self,
+        trace_info: &TraceInfo,
+        main_trace: &ColMatrix<Self::BaseField>,
+        domain: &StarkDomain<Self::BaseField>,
+        partition_option: PartitionOptions,
+        zk_parameters: Option<ZkParameters>,
+        prng: &mut Option<Self::ZkPrng>,
+    ) -> (Self::TraceLde<E>, TracePolyTable<E>) {
+        DefaultTraceLde::new(trace_info, main_trace, domain, partition_option, zk_parameters, prng)
+    }
+
+    fn new_evaluator<'a, E: FieldElement<BaseField = Self::BaseField>>(
+        &self,
+        air: &'a Self::Air,
+        aux_rand_elements: Option<AuxRandElements<E>>,
+        composition_coefficients: ConstraintCompositionCoefficients<E>,
+    ) -> Self::ConstraintEvaluator<'a, E> {
+        DefaultConstraintEvaluator::new(air, aux_rand_elements, composition_coefficients)
+    }
+
+    fn build_constraint_commitment<E: FieldElement<BaseField = Self::BaseField>>(
+        &self,
+        composition_poly_trace: CompositionPolyTrace<E>,
+        num_constraint_composition_columns: usize,
+        domain: &StarkDomain<Self::BaseField>,
+        partition_options: PartitionOptions,
+        zk_parameters: Option<ZkParameters>,
+        prng: &mut Option<Self::ZkPrng>,
+    ) -> (Self::ConstraintCommitment<E>, CompositionPoly<E>) {
+        DefaultConstraintCommitment::new(
+            composition_poly_trace,
+            num_constraint_composition_columns,
+            domain,
+            partition_options,
+            zk_parameters,
+            prng,
+        )
+    }
+}