1use rand::{distributions::Uniform, prelude::Distribution, Rng};
  2use winter_air::{FieldExtension, ProofOptions};
  3use winter_math::{fields::f64::BaseElement, FieldElement};
  4use winter_prover::Proof;
  5use winter_utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable};
  6
  7use crate::{
  8    dsa::rpo_stark::stark::RpoSignatureScheme,
  9    hash::{rpo::Rpo256, DIGEST_SIZE},
 10    StarkField, Word, ZERO,
 11};
 12
 13// CONSTANTS
 14// ================================================================================================
 15
 16/// Specifies the parameters of the STARK underlying the signature scheme. These parameters provide
 17/// at least 102 bits of security under the conjectured security of the toy protocol in
 18/// the ethSTARK paper [1].
 19///
 20/// [1]: https://eprint.iacr.org/2021/582
 21pub const PROOF_OPTIONS: ProofOptions =
 22    ProofOptions::new(30, 8, 12, FieldExtension::Quadratic, 4, 7, true);
 23
 24// PUBLIC KEY
 25// ================================================================================================
 26
 27/// A public key for verifying signatures.
 28///
 29/// The public key is a [Word] (i.e., 4 field elements) that is the hash of the secret key.
 30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 31pub struct PublicKey(Word);
 32
 33impl PublicKey {
 34    /// Returns the [Word] defining the public key.
 35    pub fn inner(&self) -> Word {
 36        self.0
 37    }
 38}
 39
 40impl PublicKey {
 41    /// Verifies the provided signature against provided message and this public key.
 42    pub fn verify(&self, message: Word, signature: &Signature) -> bool {
 43        signature.verify(message, *self)
 44    }
 45}
 46
 47impl Serializable for PublicKey {
 48    fn write_into<W: ByteWriter>(&self, target: &mut W) {
 49        self.0.write_into(target);
 50    }
 51}
 52
 53impl Deserializable for PublicKey {
 54    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
 55        let pk = <Word>::read_from(source)?;
 56        Ok(Self(pk))
 57    }
 58}
 59
 60// SECRET KEY
 61// ================================================================================================
 62
 63/// A secret key for generating signatures.
 64///
 65/// The secret key is a [Word] (i.e., 4 field elements).
 66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
 67pub struct SecretKey(Word);
 68
 69impl SecretKey {
 70    /// Generates a secret key from OS-provided randomness.
 71    pub fn new(word: Word) -> Self {
 72        Self(word)
 73    }
 74
 75    /// Generates a secret key from a [Word].
 76    #[cfg(feature = "std")]
 77    pub fn random() -> Self {
 78        use rand::{rngs::StdRng, SeedableRng};
 79
 80        let mut rng = StdRng::from_entropy();
 81        Self::with_rng(&mut rng)
 82    }
 83
 84    /// Generates a secret_key using the provided random number generator `Rng`.
 85    pub fn with_rng<R: Rng>(rng: &mut R) -> Self {
 86        let mut sk = [ZERO; 4];
 87        let uni_dist = Uniform::from(0..BaseElement::MODULUS);
 88
 89        for s in sk.iter_mut() {
 90            let sampled_integer = uni_dist.sample(rng);
 91            *s = BaseElement::new(sampled_integer);
 92        }
 93
 94        Self(sk)
 95    }
 96
 97    /// Computes the public key corresponding to this secret key.
 98    pub fn public_key(&self) -> PublicKey {
 99        let mut elements = [BaseElement::ZERO; 8];
100        elements[..DIGEST_SIZE].copy_from_slice(&self.0);
101        let pk = Rpo256::hash_elements(&elements);
102        PublicKey(pk.into())
103    }
104
105    /// Signs a message with this secret key.
106    pub fn sign(&self, message: Word) -> Signature {
107        let signature: RpoSignatureScheme<Rpo256> = RpoSignatureScheme::new(PROOF_OPTIONS);
108        let proof = signature.sign(self.0, message);
109        Signature { proof }
110    }
111}
112
113impl Serializable for SecretKey {
114    fn write_into<W: ByteWriter>(&self, target: &mut W) {
115        self.0.write_into(target);
116    }
117}
118
119impl Deserializable for SecretKey {
120    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
121        let sk = <Word>::read_from(source)?;
122        Ok(Self(sk))
123    }
124}
125
126// SIGNATURE
127// ================================================================================================
128
129/// An RPO STARK-based signature over a message.
130///
131/// The signature is a STARK proof of knowledge of a pre-image given an image where the map is
132/// the RPO permutation, the pre-image is the secret key and the image is the public key.
133/// The current implementation follows the description in [1] but relies on the conjectured security
134/// of the toy protocol in the ethSTARK paper [2], which gives us using the parameter set
135/// given in `PROOF_OPTIONS` a signature with $102$ bits of average-case existential unforgeability
136/// security against $2^{113}$-query bound adversaries that can obtain up to $2^{64}$ signatures
137/// under the same public key.
138///
139/// [1]: https://eprint.iacr.org/2024/1553
140/// [2]: https://eprint.iacr.org/2021/582
141#[derive(Debug, Clone, PartialEq, Eq)]
142pub struct Signature {
143    proof: Proof,
144}
145
146impl Signature {
147    /// Returns the STARK proof constituting the signature.
148    pub fn inner(&self) -> Proof {
149        self.proof.clone()
150    }
151
152    /// Returns true if this signature is a valid signature for the specified message generated
153    /// against the secret key matching the specified public key.
154    pub fn verify(&self, message: Word, pk: PublicKey) -> bool {
155        let signature: RpoSignatureScheme<Rpo256> = RpoSignatureScheme::new(PROOF_OPTIONS);
156
157        let res = signature.verify(pk.inner(), message, self.proof.clone());
158        res.is_ok()
159    }
160}
161
162impl Serializable for Signature {
163    fn write_into<W: ByteWriter>(&self, target: &mut W) {
164        self.proof.write_into(target);
165    }
166}
167
168impl Deserializable for Signature {
169    fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
170        let proof = Proof::read_from(source)?;
171        Ok(Self { proof })
172    }
173}