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
125pub 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
136pub 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 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
162pub 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 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}