penumbra_sdk_custody/threshold/dkg/
encryption.rs

1use anyhow::{anyhow, Result};
2use chacha20poly1305::{aead::AeadInPlace, aead::NewAead, ChaCha20Poly1305, Key as SymmetricKey};
3use rand_core::CryptoRngCore;
4
5/// The number of bytes in a key agreement public key.
6const PK_SIZE: usize = 32;
7/// The number of bytes in our AEAD's authentication tag.
8const TAG_SIZE: usize = 16;
9/// The number of bytes in our nonce.
10const NONCE_SIZE: usize = 12;
11
12/// Generate a random shared secret.
13///
14/// I had to write this method because it didn't exist natively. That's probably
15/// a good idea. The reason I added it here is because we want to do the trick
16/// of using a fake shared secret to make key agreement infallible.
17fn random_shared_secret(rng: &mut impl CryptoRngCore) -> decaf377_ka::SharedSecret {
18    let mut data = [0u8; 32];
19    rng.fill_bytes(&mut data);
20    decaf377_ka::SharedSecret(data)
21}
22
23/// Perform key agreement in a way that cannot fail.
24///
25/// Whenever key agreement were to fail, instead we return a random shared secret.
26/// This means that if an invalid public key is announced, the ciphertexts encrypted
27/// to that key will not be decryptable. This is the same situation as announcing a
28/// public key whose corresponding secret key you do not know; this behavior thus seems fine.
29fn infallible_key_agreement(
30    rng: &mut impl CryptoRngCore,
31    sk: &decaf377_ka::Secret,
32    pk: &decaf377_ka::Public,
33) -> decaf377_ka::SharedSecret {
34    let fake_secret = random_shared_secret(rng);
35    sk.key_agreement_with(pk).unwrap_or(fake_secret)
36}
37
38/// Derive a symmetric key from the information used to produce a shared secret.
39///
40/// `pk` is the target public key we're encrypting data to.
41/// `epk` is data we add to the ciphertext so that the owner of that public key can decrypt it.
42/// `secret` is the result of a key exchange between those two keys.
43fn derive_symmetric_key(
44    pk: &decaf377_ka::Public,
45    epk: &decaf377_ka::Public,
46    secret: &decaf377_ka::SharedSecret,
47) -> SymmetricKey {
48    TryInto::<[u8; 32]>::try_into(
49        &blake2b_simd::Params::new()
50            .personal(b"dkg-encryption")
51            .to_state()
52            .update(&pk.0)
53            .update(&epk.0)
54            .update(&secret.0)
55            .finalize()
56            .as_array()[..32],
57    )
58    .expect("array conversion should not fail")
59    .into()
60}
61
62/// A key to which data can be encrypted.
63///
64/// This key has a corresponding decryption key which can decrypt the messages encrypted to it.
65#[derive(Clone, Copy)]
66pub struct EncryptionKey(decaf377_ka::Public);
67
68impl EncryptionKey {
69    /// Encrypt a message, producing a ciphertext.
70    pub fn encrypt(&self, rng: &mut impl CryptoRngCore, message: &[u8]) -> Vec<u8> {
71        let esk = decaf377_ka::Secret::new(rng);
72        let epk = esk.public();
73        let secret = infallible_key_agreement(rng, &esk, &self.0);
74        let key = derive_symmetric_key(&self.0, &epk, &secret);
75        // ciphertext = EPK || <aead tag> || <encrypted data>
76        // The tag will also include the EPK as associated data that gets authenticated.
77        let ciphertext = {
78            let mut ciphertext = Vec::new();
79            ciphertext.extend_from_slice(&epk.0);
80            // Reserve space for the tag
81            ciphertext.extend_from_slice(&[0u8; TAG_SIZE]);
82            // Include the message, which will be written over in place
83            ciphertext.extend_from_slice(message);
84            let tag = ChaCha20Poly1305::new(&key)
85                .encrypt_in_place_detached(
86                    &[0u8; NONCE_SIZE].into(),
87                    &epk.0,
88                    &mut ciphertext[PK_SIZE + TAG_SIZE..],
89                )
90                .expect("chacha20poly1305 encryption should not fail");
91            ciphertext[PK_SIZE..PK_SIZE + TAG_SIZE].copy_from_slice(&tag);
92            ciphertext
93        };
94
95        ciphertext
96    }
97
98    /// Return a view of this value's underlying bytes
99    pub fn as_bytes(&self) -> &[u8; 32] {
100        &self.0 .0
101    }
102}
103
104impl TryFrom<&[u8]> for EncryptionKey {
105    type Error = anyhow::Error;
106
107    fn try_from(value: &[u8]) -> std::result::Result<Self, Self::Error> {
108        let repr: [u8; 32] = value.try_into()?;
109        Ok(Self(decaf377_ka::Public(repr)))
110    }
111}
112
113/// A key that allows decrypting ciphertexts sent to the corresponding encryption key.
114#[derive(Clone)]
115pub struct DecryptionKey(decaf377_ka::Secret);
116
117impl From<DecryptionKey> for EncryptionKey {
118    fn from(value: DecryptionKey) -> Self {
119        EncryptionKey(value.0.public())
120    }
121}
122
123impl DecryptionKey {
124    pub fn new(rng: &mut impl CryptoRngCore) -> Self {
125        Self(decaf377_ka::Secret::new(rng))
126    }
127
128    /// Get the corresponding public encryption key for this decryption key.
129    ///
130    /// This is a synonym for the From impl.
131    pub fn public(&self) -> EncryptionKey {
132        self.clone().into()
133    }
134
135    /// Decrypt a ciphertext, extracitng out the corresponding message.
136    ///
137    /// This may potentially fail, if the ciphertext is malformed, or was tampered with.
138    pub fn decrypt(&self, rng: &mut impl CryptoRngCore, ciphertext: &[u8]) -> Result<Vec<u8>> {
139        if ciphertext.len() < PK_SIZE + TAG_SIZE {
140            anyhow::bail!("failed to decrypt ciphertext");
141        }
142        let (header, message) = ciphertext.split_at(PK_SIZE + TAG_SIZE);
143        let mut message = message.to_owned();
144        let epk = decaf377_ka::Public(
145            header[..PK_SIZE]
146                .try_into()
147                .expect("array conversion should not fail"),
148        );
149        // Not key committing, but shouldn't be a big concern.
150        //
151        // (By this I mean that decryption may still succeed even if the public key gets mangled in
152        // the ciphertext, but that's alright I guess).
153        let secret = infallible_key_agreement(rng, &self.0, &epk);
154        let key = derive_symmetric_key(&self.0.public(), &epk, &secret);
155        ChaCha20Poly1305::new(&key)
156            .decrypt_in_place_detached(
157                &[0u8; NONCE_SIZE].into(),
158                &header[..PK_SIZE],
159                &mut message,
160                header[PK_SIZE..PK_SIZE + TAG_SIZE].into(),
161            )
162            .map_err(|_| anyhow!("failed to decrypt ciphertext"))?;
163        Ok(message)
164    }
165}
166
167#[cfg(test)]
168mod test {
169    use rand_core::OsRng;
170
171    use super::*;
172
173    #[test]
174    fn test_encryption_roundtrip() -> Result<()> {
175        let mut rng = OsRng;
176
177        let dk = DecryptionKey::new(&mut rng);
178        let ek = dk.public();
179        let msg = "ペンブラが好きです".as_bytes();
180        let ciphertext = ek.encrypt(&mut rng, msg);
181        let msg2 = dk.decrypt(&mut rng, &ciphertext)?;
182        assert_eq!(msg, &msg2);
183
184        Ok(())
185    }
186}