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.clone())
 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 to_bytes(&self) -> [u8; 32] {
100        self.secret_key.to_bytes()
101    }
102
103    pub fn generate<D: Digest>(rng: &mut impl CryptoRngCore) -> Result<Self, Error> {
104        let mut secret_key = SecretKey([0; 32]);
105        rng.fill_bytes(&mut secret_key.0);
106        Self::new::<D>(&secret_key)
107    }
108
109    pub fn public_key(&self) -> &PublicKey<TE> {
110        &self.public_key
111    }
112
113    pub fn sign<D: Digest, E: Absorb>(
114        &self,
115        poseidon: &PoseidonConfig<TE::BaseField>,
116        message: &[E],
117    ) -> Signature<TE> {
118        let (x, prefix) = self.secret_key.expand::<TE::ScalarField, D>();
119
120        let mut h = D::new();
121        h.update(prefix);
122        message
123            .iter()
124            .for_each(|m| h.update(m.to_sponge_bytes_as_vec()));
125        let r: TE::ScalarField = crate::from_digest(h);
126        let sig_r: Affine<TE> = (Affine::<TE>::generator() * r).into();
127
128        let mut poseidon = PoseidonSponge::new(poseidon);
129
130        let (sig_r_x, sig_r_y) = sig_r.xy().unwrap();
131        poseidon.absorb(sig_r_x);
132        poseidon.absorb(sig_r_y);
133        let (pk_x, pk_y) = self.public_key.0.xy().unwrap();
134        poseidon.absorb(pk_x);
135        poseidon.absorb(pk_y);
136        message.iter().for_each(|m| poseidon.absorb(m));
137
138        let k = poseidon.squeeze_field_elements::<TE::ScalarField>(1);
139        let k = k.first().unwrap();
140
141        let sig_s = (x * k) + r;
142
143        Signature::new(sig_r, sig_s)
144    }
145}
146
147impl<TE: TECurveConfig> SigningKey<TE> {
148    pub fn shared_key<D: Digest>(&self, recipient: &PublicKey<TE>) -> [u8; 32] {
149        let (x, _) = self.secret_key.expand::<TE::ScalarField, D>();
150        let shared_key: Affine<TE> = (*recipient.as_ref() * x).into();
151        let mut data = Vec::new();
152        shared_key.serialize_compressed(&mut data).unwrap();
153        data[00..32].try_into().unwrap()
154    }
155}
156
157impl<TE: TECurveConfig + Clone> PublicKey<TE>
158where
159    TE::BaseField: PrimeField + Absorb,
160{
161    pub fn verify<E: Absorb>(
162        &self,
163        poseidon: &PoseidonConfig<TE::BaseField>,
164        message: &[E],
165        signature: &Signature<TE>,
166    ) -> Result<(), Error> {
167        let mut poseidon = PoseidonSponge::new(poseidon);
168
169        let (sig_r_x, sig_r_y) = signature.r().xy().unwrap();
170        poseidon.absorb(sig_r_x);
171        poseidon.absorb(sig_r_y);
172        let (pk_x, pk_y) = self.0.xy().unwrap();
173        poseidon.absorb(pk_x);
174        poseidon.absorb(pk_y);
175        message.iter().for_each(|m| poseidon.absorb(m));
176
177        let k = poseidon.squeeze_field_elements::<TE::ScalarField>(1);
178        let k = k.first().unwrap();
179
180        let kx_b = self.0 * k;
181        let s_b = Affine::<TE>::generator() * signature.s();
182        let r_rec: Affine<TE> = (s_b - kx_b).into();
183
184        (signature.r() == &r_rec).then_some(()).ok_or(Error::Verify)
185    }
186}