tpe/
lib.rs

1use std::collections::BTreeMap;
2use std::io::Write;
3use std::ops::Mul;
4
5use bitcoin_hashes::{sha256, Hash};
6use bls12_381::{pairing, G1Projective, G2Projective, Scalar};
7pub use bls12_381::{G1Affine, G2Affine};
8use fedimint_core::bls12_381_serde;
9use fedimint_core::encoding::{Decodable, Encodable};
10use group::ff::Field;
11use group::{Curve, Group};
12use rand_chacha::rand_core::SeedableRng;
13use rand_chacha::ChaChaRng;
14use serde::{Deserialize, Serialize};
15
16#[derive(Copy, Clone, Debug, Eq, PartialEq, Encodable, Decodable, Serialize, Deserialize)]
17pub struct SecretKeyShare(#[serde(with = "bls12_381_serde::scalar")] pub Scalar);
18
19#[derive(Copy, Clone, Debug, Eq, PartialEq, Encodable, Decodable, Serialize, Deserialize)]
20pub struct PublicKeyShare(#[serde(with = "bls12_381_serde::g1")] pub G1Affine);
21
22#[derive(Copy, Clone, Debug, Eq, PartialEq, Encodable, Decodable, Serialize, Deserialize)]
23pub struct AggregatePublicKey(#[serde(with = "bls12_381_serde::g1")] pub G1Affine);
24
25#[derive(Copy, Clone, Debug, Eq, PartialEq, Encodable, Decodable, Serialize, Deserialize)]
26pub struct DecryptionKeyShare(#[serde(with = "bls12_381_serde::g1")] pub G1Affine);
27
28#[derive(Copy, Clone, Debug, Eq, PartialEq, Encodable, Decodable, Serialize, Deserialize)]
29pub struct AggregateDecryptionKey(#[serde(with = "bls12_381_serde::g1")] pub G1Affine);
30
31#[derive(Copy, Clone, Debug, Eq, PartialEq, Encodable, Decodable, Serialize, Deserialize)]
32pub struct EphemeralPublicKey(#[serde(with = "bls12_381_serde::g1")] pub G1Affine);
33
34#[derive(Copy, Clone, Debug, Eq, PartialEq, Encodable, Decodable, Serialize, Deserialize)]
35pub struct EphemeralSignature(#[serde(with = "bls12_381_serde::g2")] pub G2Affine);
36
37#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Encodable, Decodable, Serialize, Deserialize)]
38pub struct CipherText {
39    #[serde(with = "serde_big_array::BigArray")]
40    pub encrypted_preimage: [u8; 32],
41    pub pk: EphemeralPublicKey,
42    pub signature: EphemeralSignature,
43}
44
45pub fn derive_public_key_share(sk: &SecretKeyShare) -> PublicKeyShare {
46    PublicKeyShare(G1Projective::generator().mul(sk.0).to_affine())
47}
48
49pub fn encrypt_preimage(
50    agg_pk: &AggregatePublicKey,
51    encryption_seed: &[u8; 32],
52    preimage: &[u8; 32],
53    commitment: &sha256::Hash,
54) -> CipherText {
55    let agg_dk = derive_agg_decryption_key(agg_pk, encryption_seed);
56    let encrypted_preimage = xor_with_hash(*preimage, &agg_dk);
57
58    let ephemeral_sk = derive_ephemeral_sk(encryption_seed);
59    let ephemeral_pk = G1Projective::generator().mul(ephemeral_sk).to_affine();
60    let ephemeral_signature = hash_to_message(&encrypted_preimage, &ephemeral_pk, commitment)
61        .mul(ephemeral_sk)
62        .to_affine();
63
64    CipherText {
65        encrypted_preimage,
66        pk: EphemeralPublicKey(ephemeral_pk),
67        signature: EphemeralSignature(ephemeral_signature),
68    }
69}
70
71pub fn derive_agg_decryption_key(
72    agg_pk: &AggregatePublicKey,
73    encryption_seed: &[u8; 32],
74) -> AggregateDecryptionKey {
75    AggregateDecryptionKey(
76        agg_pk
77            .0
78            .mul(derive_ephemeral_sk(encryption_seed))
79            .to_affine(),
80    )
81}
82
83fn derive_ephemeral_sk(encryption_seed: &[u8; 32]) -> Scalar {
84    Scalar::random(&mut ChaChaRng::from_seed(*encryption_seed))
85}
86
87fn xor_with_hash(mut bytes: [u8; 32], agg_dk: &AggregateDecryptionKey) -> [u8; 32] {
88    let hash = sha256::Hash::hash(&agg_dk.0.to_compressed());
89
90    for i in 0..32 {
91        bytes[i] ^= hash[i];
92    }
93
94    bytes
95}
96
97fn hash_to_message(
98    encrypted_point: &[u8; 32],
99    ephemeral_pk: &G1Affine,
100    commitment: &sha256::Hash,
101) -> G2Affine {
102    let mut engine = sha256::HashEngine::default();
103
104    engine
105        .write_all("FEDIMINT_TPE_BLS12_381_MESSAGE".as_bytes())
106        .expect("Writing to a hash engine cannot fail");
107
108    engine
109        .write_all(encrypted_point)
110        .expect("Writing to a hash engine cannot fail");
111
112    engine
113        .write_all(&ephemeral_pk.to_compressed())
114        .expect("Writing to a hash engine cannot fail");
115
116    engine
117        .write_all(commitment.as_byte_array())
118        .expect("Writing to a hash engine cannot fail");
119
120    let seed = sha256::Hash::from_engine(engine).to_byte_array();
121
122    G2Projective::random(&mut ChaChaRng::from_seed(seed)).to_affine()
123}
124
125/// Verifying a ciphertext guarantees that it has not been malleated.
126pub fn verify_ciphertext(ct: &CipherText, commitment: &sha256::Hash) -> bool {
127    let message = hash_to_message(&ct.encrypted_preimage, &ct.pk.0, commitment);
128
129    pairing(&G1Affine::generator(), &ct.signature.0) == pairing(&ct.pk.0, &message)
130}
131
132pub fn decrypt_preimage(ct: &CipherText, agg_dk: &AggregateDecryptionKey) -> [u8; 32] {
133    xor_with_hash(ct.encrypted_preimage, agg_dk)
134}
135
136/// The function asserts that the ciphertext is valid.
137pub fn verify_agg_decryption_key(
138    agg_pk: &AggregatePublicKey,
139    agg_dk: &AggregateDecryptionKey,
140    ct: &CipherText,
141    commitment: &sha256::Hash,
142) -> bool {
143    let message = hash_to_message(&ct.encrypted_preimage, &ct.pk.0, commitment);
144
145    assert_eq!(
146        pairing(&G1Affine::generator(), &ct.signature.0),
147        pairing(&ct.pk.0, &message)
148    );
149
150    // Since the ciphertext is valid its signature is the ecdh point of the message
151    // and the ephemeral public key. Hence, the following equation holds if and only
152    // if the aggregate decryption key is the ecdh point of the ephemeral public key
153    // and the aggregate public key.
154
155    pairing(&agg_dk.0, &message) == pairing(&agg_pk.0, &ct.signature.0)
156}
157
158pub fn create_decryption_key_share(sks: &SecretKeyShare, ct: &CipherText) -> DecryptionKeyShare {
159    DecryptionKeyShare(ct.pk.0.mul(sks.0).to_affine())
160}
161
162/// The function asserts that the ciphertext is valid.
163pub fn verify_decryption_key_share(
164    pks: &PublicKeyShare,
165    dks: &DecryptionKeyShare,
166    ct: &CipherText,
167    commitment: &sha256::Hash,
168) -> bool {
169    let message = hash_to_message(&ct.encrypted_preimage, &ct.pk.0, commitment);
170
171    assert_eq!(
172        pairing(&G1Affine::generator(), &ct.signature.0),
173        pairing(&ct.pk.0, &message)
174    );
175
176    // Since the ciphertext is valid its signature is the ecdh point of the message
177    // and the ephemeral public key. Hence, the following equation holds if and only
178    // if the decryption key share is the ecdh point of the ephemeral public key and
179    // the public key share.
180
181    pairing(&dks.0, &message) == pairing(&pks.0, &ct.signature.0)
182}
183
184pub fn aggregate_decryption_shares(
185    shares: &BTreeMap<u64, DecryptionKeyShare>,
186) -> AggregateDecryptionKey {
187    AggregateDecryptionKey(
188        lagrange_multipliers(shares.keys().cloned().map(Scalar::from).collect())
189            .into_iter()
190            .zip(shares.values())
191            .map(|(lagrange_multiplier, share)| lagrange_multiplier * share.0)
192            .reduce(|a, b| a + b)
193            .expect("We have at least one share")
194            .to_affine(),
195    )
196}
197
198fn lagrange_multipliers(scalars: Vec<Scalar>) -> Vec<Scalar> {
199    scalars
200        .iter()
201        .map(|i| {
202            scalars
203                .iter()
204                .filter(|j| *j != i)
205                .map(|j| j * (j - i).invert().expect("We filtered the case j == i"))
206                .reduce(|a, b| a * b)
207                .expect("We have at least one share")
208        })
209        .collect()
210}
211
212macro_rules! impl_hash_with_serialized_compressed {
213    ($type:ty) => {
214        impl std::hash::Hash for $type {
215            fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
216                state.write(&self.0.to_compressed());
217            }
218        }
219    };
220}
221
222impl_hash_with_serialized_compressed!(AggregatePublicKey);
223impl_hash_with_serialized_compressed!(DecryptionKeyShare);
224impl_hash_with_serialized_compressed!(AggregateDecryptionKey);
225impl_hash_with_serialized_compressed!(EphemeralPublicKey);
226impl_hash_with_serialized_compressed!(EphemeralSignature);
227impl_hash_with_serialized_compressed!(PublicKeyShare);
228
229#[cfg(test)]
230mod tests {
231    use std::collections::BTreeMap;
232
233    use bitcoin_hashes::{sha256, Hash};
234    use bls12_381::{G1Projective, Scalar};
235    use group::ff::Field;
236    use group::Curve;
237    use rand::rngs::OsRng;
238
239    use crate::{
240        aggregate_decryption_shares, create_decryption_key_share, decrypt_preimage,
241        derive_agg_decryption_key, encrypt_preimage, verify_agg_decryption_key, verify_ciphertext,
242        verify_decryption_key_share, AggregatePublicKey, DecryptionKeyShare, PublicKeyShare,
243        SecretKeyShare,
244    };
245
246    fn dealer_keygen(
247        threshold: usize,
248        keys: usize,
249    ) -> (AggregatePublicKey, Vec<PublicKeyShare>, Vec<SecretKeyShare>) {
250        let poly: Vec<Scalar> = (0..threshold).map(|_| Scalar::random(&mut OsRng)).collect();
251
252        let apk = (G1Projective::generator() * eval_polynomial(&poly, &Scalar::zero())).to_affine();
253
254        let sks: Vec<SecretKeyShare> = (0..keys)
255            .map(|idx| SecretKeyShare(eval_polynomial(&poly, &Scalar::from(idx as u64 + 1))))
256            .collect();
257
258        let pks = sks
259            .iter()
260            .map(|sk| PublicKeyShare((G1Projective::generator() * sk.0).to_affine()))
261            .collect();
262
263        (AggregatePublicKey(apk), pks, sks)
264    }
265
266    fn eval_polynomial(coefficients: &[Scalar], x: &Scalar) -> Scalar {
267        coefficients
268            .iter()
269            .cloned()
270            .rev()
271            .reduce(|acc, coefficient| acc * x + coefficient)
272            .expect("We have at least one coefficient")
273    }
274
275    #[test]
276    fn test_roundtrip() {
277        let (agg_pk, pks, sks) = dealer_keygen(3, 4);
278
279        let encryption_seed = [7_u8; 32];
280        let preimage = [42_u8; 32];
281        let commitment = sha256::Hash::hash(&[0_u8; 32]);
282        let ciphertext = encrypt_preimage(&agg_pk, &encryption_seed, &preimage, &commitment);
283
284        assert!(verify_ciphertext(&ciphertext, &commitment));
285
286        let shares: Vec<DecryptionKeyShare> = sks
287            .iter()
288            .map(|sk| create_decryption_key_share(sk, &ciphertext))
289            .collect();
290
291        for (pk, share) in pks.iter().zip(shares.iter()) {
292            assert!(verify_decryption_key_share(
293                pk,
294                share,
295                &ciphertext,
296                &commitment
297            ));
298        }
299
300        let selected_shares: BTreeMap<u64, DecryptionKeyShare> = (1_u64..4).zip(shares).collect();
301
302        assert_eq!(selected_shares.len(), 3);
303
304        let agg_dk = aggregate_decryption_shares(&selected_shares);
305
306        assert_eq!(agg_dk, derive_agg_decryption_key(&agg_pk, &encryption_seed));
307
308        assert!(verify_agg_decryption_key(
309            &agg_pk,
310            &agg_dk,
311            &ciphertext,
312            &commitment
313        ));
314
315        assert_eq!(preimage, decrypt_preimage(&ciphertext, &agg_dk));
316    }
317}