Make DKG Encryption a bit more flexible

No longer requires the use of an EncryptionKeyMessage, and allows pre-defined keys for encryption.
2025-12-13 14:39:25 +00:00 · 2024-07-23 19:21:27 -04:00
parent 772d033bb2
commit b8472963c9
4 changed files with 123 additions and 109 deletions
--- a/crypto/dkg/src/encryption.rs
+++ b/crypto/dkg/src/encryption.rs
@@ -48,8 +48,8 @@ pub(crate) use sealed::*;
 /// Wraps a message with a key to use for encryption in the future.
 #[derive(Clone, PartialEq, Eq, Debug, Zeroize)]
 pub struct EncryptionKeyMessage<C: Ciphersuite, M: Message> {
-  msg: M,
-  enc_key: C::G,
+  pub(crate) msg: M,
+  pub(crate) enc_key: C::G,
 }

 // Doesn't impl ReadWrite so that doesn't need to be imported
@@ -98,11 +98,11 @@ fn ecdh<C: Ciphersuite>(private: &Zeroizing<C::F>, public: C::G) -> Zeroizing<C:

 // Each ecdh must be distinct. Reuse of an ecdh for multiple ciphers will cause the messages to be
 // leaked.
-fn cipher<C: Ciphersuite>(context: &str, ecdh: &Zeroizing<C::G>) -> ChaCha20 {
+fn cipher<C: Ciphersuite>(context: [u8; 32], ecdh: &Zeroizing<C::G>) -> ChaCha20 {
  // Ideally, we'd box this transcript with ZAlloc, yet that's only possible on nightly
  // TODO: https://github.com/serai-dex/serai/issues/151
  let mut transcript = RecommendedTranscript::new(b"DKG Encryption v0.2");
-  transcript.append_message(b"context", context.as_bytes());
+  transcript.append_message(b"context", context);

  transcript.domain_separate(b"encryption_key");

@@ -134,7 +134,7 @@ fn cipher<C: Ciphersuite>(context: &str, ecdh: &Zeroizing<C::G>) -> ChaCha20 {

 fn encrypt<R: RngCore + CryptoRng, C: Ciphersuite, E: Encryptable>(
  rng: &mut R,
-  context: &str,
+  context: [u8; 32],
  from: Participant,
  to: C::G,
  mut msg: Zeroizing<E>,
@@ -197,7 +197,7 @@ impl<C: Ciphersuite, E: Encryptable> EncryptedMessage<C, E> {
  pub(crate) fn invalidate_msg<R: RngCore + CryptoRng>(
    &mut self,
    rng: &mut R,
-    context: &str,
+    context: [u8; 32],
    from: Participant,
  ) {
    // Invalidate the message by specifying a new key/Schnorr PoP
@@ -219,7 +219,7 @@ impl<C: Ciphersuite, E: Encryptable> EncryptedMessage<C, E> {
  pub(crate) fn invalidate_share_serialization<R: RngCore + CryptoRng>(
    &mut self,
    rng: &mut R,
-    context: &str,
+    context: [u8; 32],
    from: Participant,
    to: C::G,
  ) {
@@ -243,7 +243,7 @@ impl<C: Ciphersuite, E: Encryptable> EncryptedMessage<C, E> {
  pub(crate) fn invalidate_share_value<R: RngCore + CryptoRng>(
    &mut self,
    rng: &mut R,
-    context: &str,
+    context: [u8; 32],
    from: Participant,
    to: C::G,
  ) {
@@ -300,14 +300,14 @@ impl<C: Ciphersuite> EncryptionKeyProof<C> {
 // This still doesn't mean the DKG offers an authenticated channel. The per-message keys have no
 // root of trust other than their existence in the assumed-to-exist external authenticated channel.
 fn pop_challenge<C: Ciphersuite>(
-  context: &str,
+  context: [u8; 32],
  nonce: C::G,
  key: C::G,
  sender: Participant,
  msg: &[u8],
 ) -> C::F {
  let mut transcript = RecommendedTranscript::new(b"DKG Encryption Key Proof of Possession v0.2");
-  transcript.append_message(b"context", context.as_bytes());
+  transcript.append_message(b"context", context);

  transcript.domain_separate(b"proof_of_possession");

@@ -323,9 +323,9 @@ fn pop_challenge<C: Ciphersuite>(
  C::hash_to_F(b"DKG-encryption-proof_of_possession", &transcript.challenge(b"schnorr"))
 }

-fn encryption_key_transcript(context: &str) -> RecommendedTranscript {
+fn encryption_key_transcript(context: [u8; 32]) -> RecommendedTranscript {
  let mut transcript = RecommendedTranscript::new(b"DKG Encryption Key Correctness Proof v0.2");
-  transcript.append_message(b"context", context.as_bytes());
+  transcript.append_message(b"context", context);
  transcript
 }

@@ -337,14 +337,71 @@ pub(crate) enum DecryptionError {
  InvalidProof,
 }

+// A simple box for managing decryption.
+#[derive(Clone, Debug)]
+pub(crate) struct Decryption<C: Ciphersuite> {
+  context: [u8; 32],
+  enc_keys: HashMap<Participant, C::G>,
+}
+
+impl<C: Ciphersuite> Decryption<C> {
+pub(crate) fn new(context: [u8; 32]) -> Self { Self { context, enc_keys: HashMap::new()} }
+pub(crate) fn register(
+  &mut self,
+  participant: Participant,
+  key: C::G,
+) {
+  assert!(
+    !self.enc_keys.contains_key(&participant),
+    "Re-registering encryption key for a participant"
+  );
+  self.enc_keys.insert(participant, key);
+}
+
+// Given a message, and the intended decryptor, and a proof for its key, decrypt the message.
+// Returns None if the key was wrong.
+pub(crate) fn decrypt_with_proof<E: Encryptable>(
+  &self,
+  from: Participant,
+  decryptor: Participant,
+  mut msg: EncryptedMessage<C, E>,
+  // There's no encryption key proof if the accusation is of an invalid signature
+  proof: Option<EncryptionKeyProof<C>>,
+) -> Result<Zeroizing<E>, DecryptionError> {
+  if !msg.pop.verify(
+    msg.key,
+    pop_challenge::<C>(self.context, msg.pop.R, msg.key, from, msg.msg.deref().as_ref()),
+  ) {
+    Err(DecryptionError::InvalidSignature)?;
+  }
+
+  if let Some(proof) = proof {
+    // Verify this is the decryption key for this message
+    proof
+      .dleq
+      .verify(
+        &mut encryption_key_transcript(self.context),
+        &[C::generator(), msg.key],
+        &[self.enc_keys[&decryptor], *proof.key],
+      )
+      .map_err(|_| DecryptionError::InvalidProof)?;
+
+    cipher::<C>(self.context, &proof.key).apply_keystream(msg.msg.as_mut().as_mut());
+    Ok(msg.msg)
+  } else {
+    Err(DecryptionError::InvalidProof)
+  }
+}
+}
+
 // A simple box for managing encryption.
 #[derive(Clone)]
 pub(crate) struct Encryption<C: Ciphersuite> {
-  context: String,
-  i: Option<Participant>,
+  context: [u8; 32],
+  i: Participant,
  enc_key: Zeroizing<C::F>,
  enc_pub_key: C::G,
-  enc_keys: HashMap<Participant, C::G>,
+  decryption: Decryption<C>,
 }

 impl<C: Ciphersuite> fmt::Debug for Encryption<C> {
@@ -354,7 +411,7 @@ impl<C: Ciphersuite> fmt::Debug for Encryption<C> {
      .field("context", &self.context)
      .field("i", &self.i)
      .field("enc_pub_key", &self.enc_pub_key)
-      .field("enc_keys", &self.enc_keys)
+      .field("decryption", &self.decryption)
      .finish_non_exhaustive()
  }
 }
@@ -363,25 +420,24 @@ impl<C: Ciphersuite> Zeroize for Encryption<C> {
  fn zeroize(&mut self) {
    self.enc_key.zeroize();
    self.enc_pub_key.zeroize();
-    for (_, mut value) in self.enc_keys.drain() {
+    for (_, mut value) in self.decryption.enc_keys.drain() {
      value.zeroize();
    }
  }
 }

 impl<C: Ciphersuite> Encryption<C> {
-  pub(crate) fn new<R: RngCore + CryptoRng>(
-    context: String,
-    i: Option<Participant>,
-    rng: &mut R,
+  pub(crate) fn new(
+    context: [u8; 32],
+    i: Participant,
+    enc_key: Zeroizing<C::F>,
  ) -> Self {
-    let enc_key = Zeroizing::new(C::random_nonzero_F(rng));
    Self {
      context,
      i,
      enc_pub_key: C::generator() * enc_key.deref(),
      enc_key,
-      enc_keys: HashMap::new(),
+      decryption: Decryption::new(context),
    }
  }

@@ -389,17 +445,12 @@ impl<C: Ciphersuite> Encryption<C> {
    EncryptionKeyMessage { msg, enc_key: self.enc_pub_key }
  }

-  pub(crate) fn register<M: Message>(
+  pub(crate) fn register(
    &mut self,
    participant: Participant,
-    msg: EncryptionKeyMessage<C, M>,
-  ) -> M {
-    assert!(
-      !self.enc_keys.contains_key(&participant),
-      "Re-registering encryption key for a participant"
-    );
-    self.enc_keys.insert(participant, msg.enc_key);
-    msg.msg
+    key: C::G,
+  ) {
+    self.decryption.register(participant, key)
  }

  pub(crate) fn encrypt<R: RngCore + CryptoRng, E: Encryptable>(
@@ -408,7 +459,7 @@ impl<C: Ciphersuite> Encryption<C> {
    participant: Participant,
    msg: Zeroizing<E>,
  ) -> EncryptedMessage<C, E> {
-    encrypt(rng, &self.context, self.i.unwrap(), self.enc_keys[&participant], msg)
+    encrypt(rng, self.context, self.i, self.decryption.enc_keys[&participant], msg)
  }

  pub(crate) fn decrypt<R: RngCore + CryptoRng, I: Copy + Zeroize, E: Encryptable>(
@@ -426,18 +477,18 @@ impl<C: Ciphersuite> Encryption<C> {
      batch,
      batch_id,
      msg.key,
-      pop_challenge::<C>(&self.context, msg.pop.R, msg.key, from, msg.msg.deref().as_ref()),
+      pop_challenge::<C>(self.context, msg.pop.R, msg.key, from, msg.msg.deref().as_ref()),
    );

    let key = ecdh::<C>(&self.enc_key, msg.key);
-    cipher::<C>(&self.context, &key).apply_keystream(msg.msg.as_mut().as_mut());
+    cipher::<C>(self.context, &key).apply_keystream(msg.msg.as_mut().as_mut());
    (
      msg.msg,
      EncryptionKeyProof {
        key,
        dleq: DLEqProof::prove(
          rng,
-          &mut encryption_key_transcript(&self.context),
+          &mut encryption_key_transcript(self.context),
          &[C::generator(), msg.key],
          &self.enc_key,
        ),
@@ -445,38 +496,5 @@ impl<C: Ciphersuite> Encryption<C> {
    )
  }

-  // Given a message, and the intended decryptor, and a proof for its key, decrypt the message.
-  // Returns None if the key was wrong.
-  pub(crate) fn decrypt_with_proof<E: Encryptable>(
-    &self,
-    from: Participant,
-    decryptor: Participant,
-    mut msg: EncryptedMessage<C, E>,
-    // There's no encryption key proof if the accusation is of an invalid signature
-    proof: Option<EncryptionKeyProof<C>>,
-  ) -> Result<Zeroizing<E>, DecryptionError> {
-    if !msg.pop.verify(
-      msg.key,
-      pop_challenge::<C>(&self.context, msg.pop.R, msg.key, from, msg.msg.deref().as_ref()),
-    ) {
-      Err(DecryptionError::InvalidSignature)?;
-    }
-
-    if let Some(proof) = proof {
-      // Verify this is the decryption key for this message
-      proof
-        .dleq
-        .verify(
-          &mut encryption_key_transcript(&self.context),
-          &[C::generator(), msg.key],
-          &[self.enc_keys[&decryptor], *proof.key],
-        )
-        .map_err(|_| DecryptionError::InvalidProof)?;
-
-      cipher::<C>(&self.context, &proof.key).apply_keystream(msg.msg.as_mut().as_mut());
-      Ok(msg.msg)
-    } else {
-      Err(DecryptionError::InvalidProof)
-    }
-  }
+  pub(crate) fn into_decryption(self) -> Decryption<C> { self.decryption }
 }
--- a/crypto/dkg/src/pedpop.rs
+++ b/crypto/dkg/src/pedpop.rs
@@ -24,7 +24,7 @@ use schnorr::SchnorrSignature;
 use crate::{
  Participant, DkgError, ThresholdParams, ThresholdCore, validate_map,
  encryption::{
-    ReadWrite, EncryptionKeyMessage, EncryptedMessage, Encryption, EncryptionKeyProof,
+    ReadWrite, EncryptionKeyMessage, EncryptedMessage, Encryption, Decryption, EncryptionKeyProof,
    DecryptionError,
  },
 };
@@ -32,10 +32,10 @@ use crate::{
 type FrostError<C> = DkgError<EncryptionKeyProof<C>>;

 #[allow(non_snake_case)]
-fn challenge<C: Ciphersuite>(context: &str, l: Participant, R: &[u8], Am: &[u8]) -> C::F {
+fn challenge<C: Ciphersuite>(context: [u8; 32], l: Participant, R: &[u8], Am: &[u8]) -> C::F {
  let mut transcript = RecommendedTranscript::new(b"DKG FROST v0.2");
  transcript.domain_separate(b"schnorr_proof_of_knowledge");
-  transcript.append_message(b"context", context.as_bytes());
+  transcript.append_message(b"context", context);
  transcript.append_message(b"participant", l.to_bytes());
  transcript.append_message(b"nonce", R);
  transcript.append_message(b"commitments", Am);
@@ -86,15 +86,15 @@ impl<C: Ciphersuite> ReadWrite for Commitments<C> {
 #[derive(Debug, Zeroize)]
 pub struct KeyGenMachine<C: Ciphersuite> {
  params: ThresholdParams,
-  context: String,
+  context: [u8; 32],
  _curve: PhantomData<C>,
 }

 impl<C: Ciphersuite> KeyGenMachine<C> {
  /// Create a new machine to generate a key.
  ///
-  /// The context string should be unique among multisigs.
-  pub fn new(params: ThresholdParams, context: String) -> KeyGenMachine<C> {
+  /// The context should be unique among multisigs.
+  pub fn new(params: ThresholdParams, context: [u8; 32]) -> KeyGenMachine<C> {
    KeyGenMachine { params, context, _curve: PhantomData }
  }

@@ -129,11 +129,12 @@ impl<C: Ciphersuite> KeyGenMachine<C> {
      // There's no reason to spend the time and effort to make this deterministic besides a
      // general obsession with canonicity and determinism though
      r,
-      challenge::<C>(&self.context, self.params.i(), nonce.to_bytes().as_ref(), &cached_msg),
+      challenge::<C>(self.context, self.params.i(), nonce.to_bytes().as_ref(), &cached_msg),
    );

    // Additionally create an encryption mechanism to protect the secret shares
-    let encryption = Encryption::new(self.context.clone(), Some(self.params.i), rng);
+    let encryption =
+      Encryption::new(self.context, self.params.i, Zeroizing::new(C::random_nonzero_F(rng)));

    // Step 4: Broadcast
    let msg =
@@ -177,7 +178,7 @@ fn polynomial<F: PrimeField + Zeroize>(
 // The encryption system also explicitly uses Zeroizing<M> so it can ensure anything being
 // encrypted is within Zeroizing. Accordingly, internally having Zeroizing would be redundant.
 #[derive(Clone, PartialEq, Eq)]
-pub struct SecretShare<F: PrimeField>(F::Repr);
+pub struct SecretShare<F: PrimeField>(pub(crate) F::Repr);
 impl<F: PrimeField> AsRef<[u8]> for SecretShare<F> {
  fn as_ref(&self) -> &[u8] {
    self.0.as_ref()
@@ -225,7 +226,7 @@ impl<F: PrimeField> ReadWrite for SecretShare<F> {
 #[derive(Zeroize)]
 pub struct SecretShareMachine<C: Ciphersuite> {
  params: ThresholdParams,
-  context: String,
+  context: [u8; 32],
  coefficients: Vec<Zeroizing<C::F>>,
  our_commitments: Vec<C::G>,
  encryption: Encryption<C>,
@@ -261,7 +262,8 @@ impl<C: Ciphersuite> SecretShareMachine<C> {
    let mut commitments = HashMap::new();
    for l in (1 ..= self.params.n()).map(Participant) {
      let Some(msg) = commitment_msgs.remove(&l) else { continue };
-      let mut msg = self.encryption.register(l, msg);
+      self.encryption.register(l, msg.enc_key);
+      let mut msg = msg.msg;

      if msg.commitments.len() != self.params.t().into() {
        Err(FrostError::InvalidCommitments(l))?;
@@ -274,7 +276,7 @@ impl<C: Ciphersuite> SecretShareMachine<C> {
        &mut batch,
        l,
        msg.commitments[0],
-        challenge::<C>(&self.context, l, msg.sig.R.to_bytes().as_ref(), &msg.cached_msg),
+        challenge::<C>(self.context, l, msg.sig.R.to_bytes().as_ref(), &msg.cached_msg),
      );

      commitments.insert(l, msg.commitments.drain(..).collect::<Vec<_>>());
@@ -472,7 +474,7 @@ impl<C: Ciphersuite> KeyMachine<C> {
    let KeyMachine { commitments, encryption, params, secret } = self;
    Ok(BlameMachine {
      commitments,
-      encryption,
+      encryption: encryption.into_decryption(),
      result: Some(ThresholdCore {
        params,
        secret_share: secret,
@@ -486,7 +488,7 @@ impl<C: Ciphersuite> KeyMachine<C> {
 /// A machine capable of handling blame proofs.
 pub struct BlameMachine<C: Ciphersuite> {
  commitments: HashMap<Participant, Vec<C::G>>,
-  encryption: Encryption<C>,
+  encryption: Decryption<C>,
  result: Option<ThresholdCore<C>>,
 }

@@ -505,7 +507,6 @@ impl<C: Ciphersuite> Zeroize for BlameMachine<C> {
    for commitments in self.commitments.values_mut() {
      commitments.zeroize();
    }
-    self.encryption.zeroize();
    self.result.zeroize();
  }
 }
@@ -598,18 +599,18 @@ impl<C: Ciphersuite> AdditionalBlameMachine<C> {
  /// authenticated as having come from the supposed party and verified as valid. Usage of invalid
  /// commitments is considered undefined behavior, and may cause everything from inaccurate blame
  /// to panics.
-  pub fn new<R: RngCore + CryptoRng>(
-    rng: &mut R,
-    context: String,
+  pub fn new(
+    context: [u8; 32],
    n: u16,
    mut commitment_msgs: HashMap<Participant, EncryptionKeyMessage<C, Commitments<C>>>,
  ) -> Result<Self, FrostError<C>> {
    let mut commitments = HashMap::new();
-    let mut encryption = Encryption::new(context, None, rng);
+    let mut encryption = Decryption::new(context);
    for i in 1 ..= n {
      let i = Participant::new(i).unwrap();
      let Some(msg) = commitment_msgs.remove(&i) else { Err(DkgError::MissingParticipant(i))? };
-      commitments.insert(i, encryption.register(i, msg).commitments);
+      encryption.register(i, msg.enc_key);
+      commitments.insert(i, msg.msg.commitments);
    }
    Ok(AdditionalBlameMachine(BlameMachine { commitments, encryption, result: None }))
  }
--- a/crypto/dkg/src/tests/pedpop.rs
+++ b/crypto/dkg/src/tests/pedpop.rs
@@ -14,7 +14,7 @@ use crate::{
 type PedPoPEncryptedMessage<C> = EncryptedMessage<C, SecretShare<<C as Ciphersuite>::F>>;
 type PedPoPSecretShares<C> = HashMap<Participant, PedPoPEncryptedMessage<C>>;

-const CONTEXT: &str = "DKG Test Key Generation";
+const CONTEXT: [u8; 32] = *b"DKG Test Key Generation         ";

 // Commit, then return commitment messages, enc keys, and shares
 #[allow(clippy::type_complexity)]
@@ -31,7 +31,7 @@ fn commit_enc_keys_and_shares<R: RngCore + CryptoRng, C: Ciphersuite>(
  let mut enc_keys = HashMap::new();
  for i in (1 ..= PARTICIPANTS).map(Participant) {
    let params = ThresholdParams::new(THRESHOLD, PARTICIPANTS, i).unwrap();
-    let machine = KeyGenMachine::<C>::new(params, CONTEXT.to_string());
+    let machine = KeyGenMachine::<C>::new(params, CONTEXT);
    let (machine, these_commitments) = machine.generate_coefficients(rng);
    machines.insert(i, machine);

@@ -147,14 +147,12 @@ mod literal {

      // Verify machines constructed with AdditionalBlameMachine::new work
      assert_eq!(
-        AdditionalBlameMachine::new(
-          &mut OsRng,
-          CONTEXT.to_string(),
-          PARTICIPANTS,
-          commitment_msgs.clone()
-        )
-        .unwrap()
-        .blame(ONE, TWO, msg.clone(), blame.clone()),
+        AdditionalBlameMachine::new(CONTEXT, PARTICIPANTS, commitment_msgs.clone()).unwrap().blame(
+          ONE,
+          TWO,
+          msg.clone(),
+          blame.clone()
+        ),
        ONE,
      );
    }
--- a/processor/src/key_gen.rs
+++ b/processor/src/key_gen.rs
@@ -184,13 +184,12 @@ impl<N: Network, D: Db> KeyGen<N, D> {
    const NETWORK_KEY_CONTEXT: &str = "network";
    let context = |id: &KeyGenId, key| {
      // TODO2: Also embed the chain ID/genesis block
-      format!(
-        "Serai Key Gen. Session: {:?}, Network: {:?}, Attempt: {}, Key: {}",
-        id.session,
-        N::NETWORK,
-        id.attempt,
-        key,
-      )
+      let mut transcript = RecommendedTranscript::new(b"Serai Key Gen");
+      transcript.append_message(b"session", id.session.0.to_le_bytes());
+      transcript.append_message(b"network", N::ID);
+      transcript.append_message(b"attempt", id.attempt.to_le_bytes());
+      transcript.append_message(b"key", key);
+      <[u8; 32]>::try_from(&(&transcript.challenge(b"context"))[.. 32]).unwrap()
    };

    let rng = |label, id: KeyGenId| {
@@ -557,7 +556,6 @@ impl<N: Network, D: Db> KeyGen<N, D> {
          blame.clone().and_then(|blame| EncryptionKeyProof::read(&mut blame.as_slice()).ok());

        let substrate_blame = AdditionalBlameMachine::new(
-          &mut rand_core::OsRng,
          context(&id, SUBSTRATE_KEY_CONTEXT),
          params.n(),
          substrate_commitment_msgs,
@@ -565,7 +563,6 @@ impl<N: Network, D: Db> KeyGen<N, D> {
        .unwrap()
        .blame(accuser, accused, substrate_share, substrate_blame);
        let network_blame = AdditionalBlameMachine::new(
-          &mut rand_core::OsRng,
          context(&id, NETWORK_KEY_CONTEXT),
          params.n(),
          network_commitment_msgs,