1use crate::{signature::Signature, Error};
  2use ark_crypto_primitives::sponge::{
  3    poseidon::{PoseidonConfig, PoseidonSponge},
  4    Absorb, CryptographicSponge,
  5};
  6use ark_ec::{
  7    twisted_edwards::{Affine, TECurveConfig},
  8    AffineRepr,
  9};
 10use ark_ff::PrimeField;
 11use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
 12use digest::Digest;
 13use digest::OutputSizeUser;
 14use rand_core::CryptoRngCore;
 15
 16fn prune_buffer<F: PrimeField>(mut bytes: [u8; 32]) -> F {
 17    bytes[0] &= 0b1111_1000;
 18    bytes[31] &= 0b0111_1111;
 19    bytes[31] |= 0b0100_0000;
 20    F::from_le_bytes_mod_order(&bytes[..])
 21}
 22
 23#[derive(Copy, Clone, Debug)]
 24/// EdDSA secret key is 32 byte data
 25pub struct SecretKey([u8; 32]);
 26
 27impl SecretKey {
 28    fn expand<F: PrimeField, D: Digest>(&self) -> (F, [u8; 32]) {
 29        let hash = D::new().chain_update(self.0).finalize();
 30        let (buffer, hash_prefix) = hash.split_at(32);
 31        let buffer: [u8; 32] = buffer.try_into().unwrap();
 32        let hash_prefix: [u8; 32] = hash_prefix.try_into().unwrap();
 33        let x = prune_buffer(buffer);
 34        (x, hash_prefix)
 35    }
 36
 37    pub fn to_bytes(&self) -> [u8; 32] {
 38        self.0
 39    }
 40
 41    pub fn from_bytes(bytes: [u8; 32]) -> Self {
 42        Self(bytes)
 43    }
 44}
 45
 46/// `PublicKey` is EdDSA signature verification key
 47#[derive(Copy, Clone, Debug, CanonicalSerialize, CanonicalDeserialize)]
 48pub struct PublicKey<TE: TECurveConfig>(Affine<TE>);
 49
 50impl<TE: TECurveConfig> PublicKey<TE> {
 51    pub fn xy(&self) -> (&TE::BaseField, &TE::BaseField) {
 52        self.as_ref().xy().unwrap()
 53    }
 54}
 55
 56impl<TE: TECurveConfig> From<Affine<TE>> for PublicKey<TE> {
 57    fn from(affine: Affine<TE>) -> Self {
 58        Self(affine)
 59    }
 60}
 61
 62impl<TE: TECurveConfig> AsRef<Affine<TE>> for PublicKey<TE> {
 63    fn as_ref(&self) -> &Affine<TE> {
 64        &self.0
 65    }
 66}
 67
 68#[derive(Copy, Clone, Debug)]
 69/// `SigningKey` produces EdDSA signatures for given message
 70pub struct SigningKey<TE: TECurveConfig> {
 71    secret_key: SecretKey,
 72    public_key: PublicKey<TE>,
 73}
 74
 75impl<TE: TECurveConfig + Clone> SigningKey<TE>
 76where
 77    TE::BaseField: PrimeField + Absorb,
 78{
 79    pub fn new<D: Digest>(secret_key: &SecretKey) -> Result<Self, Error> {
 80        (<D as OutputSizeUser>::output_size() == 64)
 81            .then_some(())
 82            .ok_or(Error::BadDigestOutput)?;
 83
 84        let (x, _) = secret_key.expand::<TE::ScalarField, D>();
 85        let public_key: Affine<TE> = (Affine::<TE>::generator() * x).into();
 86        let signing_key = Self {
 87            secret_key: *secret_key,
 88            public_key: PublicKey(public_key),
 89        };
 90
 91        Ok(signing_key)
 92    }
 93
 94    pub fn from_bytes<D: Digest>(bytes: [u8; 32]) -> Result<Self, Error> {
 95        let secret_key = SecretKey::from_bytes(bytes);
 96        Self::new::<D>(&secret_key)
 97    }
 98
 99    pub fn generate<D: Digest>(rng: &mut impl CryptoRngCore) -> Result<Self, Error> {
100        let mut secret_key = SecretKey([0; 32]);
101        rng.fill_bytes(&mut secret_key.0);
102        Self::new::<D>(&secret_key)
103    }
104
105    pub fn public_key(&self) -> &PublicKey<TE> {
106        &self.public_key
107    }
108
109    pub fn sign<D: Digest, E: Absorb>(
110        &self,
111        poseidon: &PoseidonConfig<TE::BaseField>,
112        message: &[E],
113    ) -> Signature<TE> {
114        let (x, prefix) = self.secret_key.expand::<TE::ScalarField, D>();
115
116        let mut h = D::new();
117        h.update(prefix);
118        message
119            .iter()
120            .for_each(|m| h.update(m.to_sponge_bytes_as_vec()));
121        let r: TE::ScalarField = crate::from_digest(h);
122        let sig_r: Affine<TE> = (Affine::<TE>::generator() * r).into();
123
124        let mut poseidon = PoseidonSponge::new(poseidon);
125
126        let (sig_r_x, sig_r_y) = sig_r.xy().unwrap();
127        poseidon.absorb(sig_r_x);
128        poseidon.absorb(sig_r_y);
129        let (pk_x, pk_y) = self.public_key.0.xy().unwrap();
130        poseidon.absorb(pk_x);
131        poseidon.absorb(pk_y);
132        message.iter().for_each(|m| poseidon.absorb(m));
133
134        let k = poseidon.squeeze_field_elements::<TE::ScalarField>(1);
135        let k = k.first().unwrap();
136
137        let sig_s = (x * k) + r;
138
139        Signature::new(sig_r, sig_s)
140    }
141}
142
143impl<TE: TECurveConfig> SigningKey<TE> {
144    pub fn shared_key<D: Digest>(&self, recipient: &PublicKey<TE>) -> [u8; 32] {
145        let (x, _) = self.secret_key.expand::<TE::ScalarField, D>();
146        let shared_key: Affine<TE> = (*recipient.as_ref() * x).into();
147        let mut data = Vec::new();
148        shared_key.serialize_compressed(&mut data).unwrap();
149        data[00..32].try_into().unwrap()
150    }
151}
152
153impl<TE: TECurveConfig + Clone> PublicKey<TE>
154where
155    TE::BaseField: PrimeField + Absorb,
156{
157    pub fn verify<E: Absorb>(
158        &self,
159        poseidon: &PoseidonConfig<TE::BaseField>,
160        message: &[E],
161        signature: &Signature<TE>,
162    ) -> Result<(), Error> {
163        let mut poseidon = PoseidonSponge::new(poseidon);
164
165        let (sig_r_x, sig_r_y) = signature.r().xy().unwrap();
166        poseidon.absorb(sig_r_x);
167        poseidon.absorb(sig_r_y);
168        let (pk_x, pk_y) = self.0.xy().unwrap();
169        poseidon.absorb(pk_x);
170        poseidon.absorb(pk_y);
171        message.iter().for_each(|m| poseidon.absorb(m));
172
173        let k = poseidon.squeeze_field_elements::<TE::ScalarField>(1);
174        let k = k.first().unwrap();
175
176        let kx_b = self.0 * k;
177        let s_b = Affine::<TE>::generator() * signature.s();
178        let r_rec: Affine<TE> = (s_b - kx_b).into();
179
180        (signature.r() == &r_rec).then_some(()).ok_or(Error::Verify)
181    }
182}