iris_crypto/
cheetah.rs

1use ibig::UBig;
2use iris_ztd::{
3    crypto::cheetah::{
4        ch_add, ch_neg, ch_scal_big, trunc_g_order, CheetahPoint, F6lt, A_GEN, G_ORDER,
5    },
6    tip5::hash::hash_varlen,
7    Belt, Digest, Hashable, Noun, NounDecode, NounEncode,
8};
9use iris_ztd_derive::{NounDecode, NounEncode};
10extern crate alloc;
11
12#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, NounEncode, NounDecode)]
13pub struct PublicKey(pub CheetahPoint);
14
15impl PublicKey {
16    pub fn verify(&self, m: &Digest, sig: &Signature) -> bool {
17        if sig.c == UBig::from(0u64)
18            || sig.c >= *G_ORDER
19            || sig.s == UBig::from(0u64)
20            || sig.s >= *G_ORDER
21        {
22            return false;
23        }
24
25        // Compute scalar = s*G - c*pubkey
26        // This is equivalent to: scalar = s*G + (-c)*pubkey
27        let sg = match ch_scal_big(&sig.s, &A_GEN) {
28            Ok(pt) => pt,
29            Err(_) => return false,
30        };
31        let c_pk = match ch_scal_big(&sig.c, &self.0) {
32            Ok(pt) => pt,
33            Err(_) => return false,
34        };
35        let scalar = match ch_add(&sg, &ch_neg(&c_pk)) {
36            Ok(pt) => pt,
37            Err(_) => return false,
38        };
39        let chal = {
40            let mut transcript: Vec<Belt> = Vec::new();
41            transcript.extend_from_slice(&scalar.x.0);
42            transcript.extend_from_slice(&scalar.y.0);
43            transcript.extend_from_slice(&self.0.x.0);
44            transcript.extend_from_slice(&self.0.y.0);
45            transcript.extend_from_slice(&m.0);
46            trunc_g_order(&hash_varlen(&mut transcript))
47        };
48
49        chal == sig.c
50    }
51
52    pub fn to_be_bytes(&self) -> [u8; 97] {
53        let mut data = [0u8; 97];
54        data[0] = 0x01; // prefix byte
55        let mut offset = 1;
56        // y-coordinate: 6 belts × 8 bytes = 48 bytes
57        for belt in self.0.y.0.iter().rev() {
58            data[offset..offset + 8].copy_from_slice(&belt.0.to_be_bytes());
59            offset += 8;
60        }
61        // x-coordinate: 6 belts × 8 bytes = 48 bytes
62        for belt in self.0.x.0.iter().rev() {
63            data[offset..offset + 8].copy_from_slice(&belt.0.to_be_bytes());
64            offset += 8;
65        }
66        data
67    }
68
69    pub fn from_be_bytes(bytes: &[u8]) -> PublicKey {
70        let mut x = [Belt(0); 6];
71        let mut y = [Belt(0); 6];
72
73        // y-coordinate: bytes 1-48
74        for i in 0..6 {
75            let offset = 1 + i * 8;
76            let mut buf = [0u8; 8];
77            buf.copy_from_slice(&bytes[offset..offset + 8]);
78            y[5 - i] = Belt(u64::from_be_bytes(buf));
79        }
80
81        // x-coordinate: bytes 49-96
82        for i in 0..6 {
83            let offset = 49 + i * 8;
84            let mut buf = [0u8; 8];
85            buf.copy_from_slice(&bytes[offset..offset + 8]);
86            x[5 - i] = Belt(u64::from_be_bytes(buf));
87        }
88
89        PublicKey(CheetahPoint {
90            x: F6lt(x),
91            y: F6lt(y),
92            inf: false,
93        })
94    }
95
96    /// SLIP-10 compatible serialization (legacy 65-byte format for compatibility)
97    pub(crate) fn to_slip10_bytes(&self) -> Vec<u8> {
98        let mut data = Vec::new();
99        for belt in self.0.y.0.iter().rev().chain(self.0.x.0.iter().rev()) {
100            data.extend_from_slice(&belt.0.to_be_bytes());
101        }
102        data
103    }
104}
105
106impl core::ops::Add for &PublicKey {
107    type Output = PublicKey;
108
109    fn add(self, other: &PublicKey) -> PublicKey {
110        PublicKey(ch_add(&self.0, &other.0).unwrap())
111    }
112}
113
114impl core::ops::Add for PublicKey {
115    type Output = PublicKey;
116
117    fn add(self, other: PublicKey) -> PublicKey {
118        &self + &other
119    }
120}
121
122impl core::ops::AddAssign for PublicKey {
123    fn add_assign(&mut self, other: PublicKey) {
124        *self = &*self + &other;
125    }
126}
127
128impl core::ops::Sub for &PublicKey {
129    type Output = PublicKey;
130
131    fn sub(self, other: &PublicKey) -> PublicKey {
132        PublicKey(ch_add(&self.0, &ch_neg(&other.0)).unwrap())
133    }
134}
135
136impl core::ops::SubAssign for PublicKey {
137    fn sub_assign(&mut self, other: PublicKey) {
138        *self = &*self - &other;
139    }
140}
141
142impl core::iter::Sum<PublicKey> for PublicKey {
143    fn sum<I: Iterator<Item = PublicKey>>(iter: I) -> Self {
144        iter.fold(PublicKey(CheetahPoint::identity()), |acc, x| &acc + &x)
145    }
146}
147
148impl<'a> core::iter::Sum<&'a PublicKey> for PublicKey {
149    fn sum<I: Iterator<Item = &'a PublicKey>>(iter: I) -> Self {
150        iter.fold(PublicKey(CheetahPoint::identity()), |acc, x| &acc + x)
151    }
152}
153
154impl Hashable for PublicKey {
155    fn hash(&self) -> Digest {
156        self.to_noun().hash()
157    }
158}
159
160#[derive(Debug, Clone)]
161pub struct Signature {
162    pub c: UBig, // challenge
163    pub s: UBig, // signature scalar
164}
165
166// Aggregate signature of the same challenge
167impl core::iter::Sum<Signature> for Option<Signature> {
168    fn sum<I: Iterator<Item = Signature>>(mut iter: I) -> Self {
169        let mut c = None;
170        let s = iter.try_fold(UBig::from(0u64), |acc, x| {
171            if c.is_some() && c.as_ref() != Some(&x.c) {
172                return None;
173            }
174            c = Some(x.c);
175            Some((acc + x.s) % &*G_ORDER)
176        });
177        Some(Signature { c: c?, s: s? })
178    }
179}
180
181impl NounEncode for Signature {
182    fn to_noun(&self) -> Noun {
183        (
184            Belt::from_bytes(&self.c.to_le_bytes()).as_slice(),
185            Belt::from_bytes(&self.s.to_le_bytes()).as_slice(),
186        )
187            .to_noun()
188    }
189}
190
191impl NounDecode for Signature {
192    fn from_noun(noun: &Noun) -> Option<Self> {
193        let (c, s): ([Belt; 8], [Belt; 8]) = NounDecode::from_noun(noun)?;
194
195        let c = Belt::to_bytes(&c);
196        let s = Belt::to_bytes(&s);
197
198        Some(Signature {
199            c: UBig::from_le_bytes(&c),
200            s: UBig::from_le_bytes(&s),
201        })
202    }
203}
204
205impl Hashable for Signature {
206    fn hash(&self) -> Digest {
207        self.to_noun().hash()
208    }
209}
210
211#[derive(Debug, Clone)]
212pub struct PrivateKey(pub UBig);
213
214impl PrivateKey {
215    pub fn public_key(&self) -> PublicKey {
216        PublicKey(ch_scal_big(&self.0, &A_GEN).unwrap())
217    }
218
219    pub fn sign(&self, m: &Digest) -> Signature {
220        self.sign_multi(m, &self.nonce_for(m), &self.public_key())
221    }
222
223    pub fn nonce_for(&self, m: &Digest) -> UBig {
224        let pubkey = self.public_key().0;
225        let nonce = {
226            let mut transcript = Vec::new();
227            transcript.extend_from_slice(&pubkey.x.0);
228            transcript.extend_from_slice(&pubkey.y.0);
229            transcript.extend_from_slice(&m.0);
230            self.0.to_le_bytes().chunks(4).for_each(|chunk| {
231                let mut buf = [0u8; 4];
232                buf[..chunk.len()].copy_from_slice(chunk);
233                transcript.push(Belt(u32::from_le_bytes(buf) as u64));
234            });
235            trunc_g_order(&hash_varlen(&mut transcript))
236        };
237        nonce
238    }
239
240    pub fn combine_nonces(nonces: &[UBig]) -> UBig {
241        nonces.iter().fold(UBig::from(0u64), |acc, x| &acc + x) % &*G_ORDER
242    }
243
244    /// Perform a multiparty sign
245    ///
246    /// # Arguments
247    /// * `m` - The digest of message to sign
248    /// * `shared_nonce` - The challenge nonce. This is after taking `nonce_for(m)` on all private keys, and combining them with [`PrivateKey::combine_nonces`].
249    /// * `combined_pubkey` - The combined public key to sign against.
250    ///
251    /// # Returns
252    /// * `Signature` - The partial signature. This will be invalid until combined with other partial signatures.
253    ///
254    /// # Example
255    ///
256    /// ```
257    /// # use iris_ztd::{Digest, Belt};
258    /// # use iris_crypto::cheetah::*;
259    /// # use ibig::UBig;
260    /// let pk1 = PrivateKey(UBig::from(123u64));
261    /// let pk2 = PrivateKey(UBig::from(456u64));
262    /// let m = Digest([Belt(8), Belt(9), Belt(10), Belt(11), Belt(12)]);
263    /// let nonce1 = pk1.nonce_for(&m);
264    /// let nonce2 = pk2.nonce_for(&m);
265    /// let combined_nonce = PrivateKey::combine_nonces(&[nonce1, nonce2]);
266    /// let combined_pubkey = pk1.public_key() + pk2.public_key();
267    /// let sig1 = pk1.sign_multi(&m, &combined_nonce, &combined_pubkey);
268    /// let sig2 = pk2.sign_multi(&m, &combined_nonce, &combined_pubkey);
269    /// let sig = [sig1, sig2].into_iter().sum::<Option<Signature>>().unwrap();
270    /// assert!(combined_pubkey.verify(&m, &sig));
271    /// ```
272    pub fn sign_multi(
273        &self,
274        m: &Digest,
275        shared_nonce: &UBig,
276        combined_pubkey: &PublicKey,
277    ) -> Signature {
278        let chal = {
279            // scalar = nonce * G
280            let scalar = ch_scal_big(shared_nonce, &A_GEN).unwrap();
281            let mut transcript = Vec::new();
282            transcript.extend_from_slice(&scalar.x.0);
283            transcript.extend_from_slice(&scalar.y.0);
284            transcript.extend_from_slice(&combined_pubkey.0.x.0);
285            transcript.extend_from_slice(&combined_pubkey.0.y.0);
286            transcript.extend_from_slice(&m.0);
287            trunc_g_order(&hash_varlen(&mut transcript))
288        };
289        let nonce = self.nonce_for(m);
290        let sig = (&nonce + &chal * &self.0) % &*G_ORDER;
291        Signature { c: chal, s: sig }
292    }
293
294    pub fn to_be_bytes(&self) -> [u8; 32] {
295        let bytes = self.0.to_be_bytes();
296        let mut arr = [0u8; 32];
297        arr[32 - bytes.len()..].copy_from_slice(&bytes);
298        arr
299    }
300}
301
302impl core::ops::Add for &PrivateKey {
303    type Output = PrivateKey;
304
305    fn add(self, other: &PrivateKey) -> PrivateKey {
306        PrivateKey((&self.0 + &other.0) % &*G_ORDER)
307    }
308}
309
310impl core::ops::Add for PrivateKey {
311    type Output = PrivateKey;
312
313    fn add(self, other: PrivateKey) -> PrivateKey {
314        PrivateKey((&self.0 + &other.0) % &*G_ORDER)
315    }
316}
317
318impl core::ops::AddAssign for PrivateKey {
319    fn add_assign(&mut self, other: PrivateKey) {
320        *self = &*self + &other;
321    }
322}
323
324impl core::ops::Sub for &PrivateKey {
325    type Output = PrivateKey;
326
327    fn sub(self, other: &PrivateKey) -> PrivateKey {
328        PrivateKey((&self.0 - &other.0) % &*G_ORDER)
329    }
330}
331
332impl core::ops::SubAssign for PrivateKey {
333    fn sub_assign(&mut self, other: PrivateKey) {
334        *self = &*self - &other;
335    }
336}
337
338impl core::iter::Sum<PrivateKey> for PrivateKey {
339    fn sum<I: Iterator<Item = PrivateKey>>(iter: I) -> Self {
340        iter.fold(PrivateKey(UBig::from(0u64)), |acc, x| &acc + &x)
341    }
342}
343
344impl<'a> core::iter::Sum<&'a PrivateKey> for PrivateKey {
345    fn sum<I: Iterator<Item = &'a PrivateKey>>(iter: I) -> Self {
346        iter.fold(PrivateKey(UBig::from(0u64)), |acc, x| &acc + x)
347    }
348}
349
350#[cfg(test)]
351mod tests {
352    use super::*;
353
354    #[test]
355    fn mupk_test() {
356        let privs = [
357            UBig::from(123u64),
358            UBig::from(124u64),
359            &*G_ORDER - &UBig::from(1u64),
360        ]
361        .map(PrivateKey);
362        let pubs = privs.clone().map(|p| p.public_key());
363        let pub_key: PublicKey = pubs.iter().sum();
364        let priv_key: PrivateKey = privs.iter().sum();
365        let pub_key_from_priv = priv_key.public_key();
366        assert_eq!(pub_key, pub_key_from_priv);
367    }
368
369    #[test]
370    fn musig_test() {
371        let privs = [
372            UBig::from(123u64),
373            UBig::from(124u64),
374            &*G_ORDER - &UBig::from(1u64),
375        ]
376        .map(PrivateKey);
377        let pubs = privs.clone().map(|p| p.public_key());
378        let pub_key: PublicKey = pubs.iter().sum();
379        let priv_key: PrivateKey = privs.iter().sum();
380
381        let digest = Digest([Belt(1), Belt(2), Belt(3), Belt(4), Belt(5)]);
382        let signature_all = priv_key.sign(&digest);
383        // Just testing regular signing
384        assert!(pub_key.verify(&digest, &signature_all));
385
386        // Now do split signing
387        let nonces = privs
388            .iter()
389            .map(|p| p.nonce_for(&digest))
390            .collect::<Vec<_>>();
391        let nonce = PrivateKey::combine_nonces(&nonces);
392        let mut sigs = vec![];
393        for priv_key in &privs {
394            sigs.push(priv_key.sign_multi(&digest, &nonce, &pub_key));
395        }
396        // Combine all signatures
397        let sig = sigs.into_iter().sum::<Option<Signature>>().unwrap();
398        // Verify combined signature
399        assert!(pub_key.verify(&digest, &sig));
400    }
401
402    #[test]
403    fn test_sign_and_verify() {
404        let priv_key = PrivateKey(UBig::from(123u64));
405        let digest = Digest([Belt(1), Belt(2), Belt(3), Belt(4), Belt(5)]);
406        let signature = priv_key.sign(&digest);
407        let pubkey = priv_key.public_key();
408        assert!(
409            pubkey.verify(&digest, &signature),
410            "Signature verification failed!"
411        );
412
413        // Corrupting digest, signature, or pubkey should all cause failure
414        let mut wrong_digest = digest;
415        wrong_digest.0[0] = Belt(0);
416        assert!(
417            !pubkey.verify(&wrong_digest, &signature),
418            "Should reject wrong digest"
419        );
420        let mut wrong_sig = signature.clone();
421        wrong_sig.s += UBig::from(1u64);
422        assert!(
423            !pubkey.verify(&digest, &wrong_sig),
424            "Should reject wrong signature"
425        );
426        let mut wrong_pubkey = pubkey.clone();
427        wrong_pubkey.0.x.0[0].0 += 1;
428        assert!(
429            !wrong_pubkey.verify(&digest, &signature),
430            "Should reject wrong public key"
431        );
432    }
433
434    #[test]
435    fn test_vector() {
436        // from nockchain zkvm-jetpack cheetah_jets.rs test_batch_verify_affine
437        let digest = Digest([Belt(8), Belt(9), Belt(10), Belt(11), Belt(12)]);
438        let pubkey = PublicKey(CheetahPoint {
439            x: F6lt([
440                Belt(2754611494552410273),
441                Belt(8599518745794843693),
442                Belt(10526511002404673680),
443                Belt(4830863958577994148),
444                Belt(375185138577093320),
445                Belt(12938930721685970739),
446            ]),
447            y: F6lt([
448                Belt(3062714866612034253),
449                Belt(15671931273416742386),
450                Belt(4071440668668521568),
451                Belt(7738250649524482367),
452                Belt(5259065445844042557),
453                Belt(8456011930642078370),
454            ]),
455            inf: false,
456        });
457        let c_hex = "6f3cd43cd8709f4368aed04cd84292ab1c380cb645aaa7d010669d70375cbe88";
458        let s_hex = "5197ab182e307a350b5cf3606d6e99a6f35b0d382c8330dde6e51fb6ef8ebb8c";
459        let signature = Signature {
460            c: UBig::from_str_radix(c_hex, 16).unwrap(),
461            s: UBig::from_str_radix(s_hex, 16).unwrap(),
462        };
463        assert!(pubkey.verify(&digest, &signature));
464    }
465}