penumbra_sdk_custody/threshold/dkg/
encryption.rs1use anyhow::{anyhow, Result};
2use chacha20poly1305::{aead::AeadInPlace, aead::NewAead, ChaCha20Poly1305, Key as SymmetricKey};
3use rand_core::CryptoRngCore;
4
5const PK_SIZE: usize = 32;
7const TAG_SIZE: usize = 16;
9const NONCE_SIZE: usize = 12;
11
12fn 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
23fn 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
38fn 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#[derive(Clone, Copy)]
66pub struct EncryptionKey(decaf377_ka::Public);
67
68impl EncryptionKey {
69 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 let ciphertext = {
78 let mut ciphertext = Vec::new();
79 ciphertext.extend_from_slice(&epk.0);
80 ciphertext.extend_from_slice(&[0u8; TAG_SIZE]);
82 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 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#[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 pub fn public(&self) -> EncryptionKey {
132 self.clone().into()
133 }
134
135 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 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}