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