1use alloc::vec::Vec;
2use core::fmt;
3
4use rand::{CryptoRng, RngCore};
5
6use super::{IesError, IesScheme, crypto_box::CryptoBox, message::SealedMessage};
7use crate::{
8 Felt,
9 aead::{aead_poseidon2::AeadPoseidon2, xchacha::XChaCha},
10 dsa::{
11 ecdsa_k256_keccak::PUBLIC_KEY_BYTES as K256_PUBLIC_KEY_BYTES,
12 eddsa_25519_sha512::PUBLIC_KEY_BYTES as X25519_PUBLIC_KEY_BYTES,
13 },
14 ecdh::{KeyAgreementScheme, k256::K256, x25519::X25519},
15 utils::{ByteReader, ByteWriter, Deserializable, DeserializationError, Serializable},
16};
17
18type K256XChaCha20Poly1305 = CryptoBox<K256, XChaCha>;
23type X25519XChaCha20Poly1305 = CryptoBox<X25519, XChaCha>;
25type K256AeadPoseidon2 = CryptoBox<K256, AeadPoseidon2>;
27type X25519AeadPoseidon2 = CryptoBox<X25519, AeadPoseidon2>;
29
30macro_rules! impl_seal_bytes_with_associated_data {
35 ($($variant:path => $crypto_box:ty, $ephemeral_variant:path;)*) => {
36 pub fn seal_bytes_with_associated_data<R: CryptoRng + RngCore>(
42 &self,
43 rng: &mut R,
44 plaintext: &[u8],
45 associated_data: &[u8],
46 ) -> Result<SealedMessage, IesError> {
47 match self {
48 $(
49 $variant(key) => {
50 let scheme = self.scheme();
51 let (ciphertext, ephemeral) = <$crypto_box>::seal_bytes_with_associated_data(
52 rng,
53 key,
54 scheme,
55 plaintext,
56 associated_data,
57 )?;
58
59 Ok(SealedMessage {
60 ephemeral_key: $ephemeral_variant(ephemeral),
61 ciphertext,
62 })
63 }
64 )*
65 }
66 }
67 };
68}
69
70macro_rules! impl_seal_elements_with_associated_data {
72 ($($variant:path => $crypto_box:ty, $ephemeral_variant:path;)*) => {
73 pub fn seal_elements_with_associated_data<R: CryptoRng + RngCore>(
79 &self,
80 rng: &mut R,
81 plaintext: &[Felt],
82 associated_data: &[Felt],
83 ) -> Result<SealedMessage, IesError> {
84 match self {
85 $(
86 $variant(key) => {
87 let scheme = self.scheme();
88 let (ciphertext, ephemeral) = <$crypto_box>::seal_elements_with_associated_data(
89 rng,
90 key,
91 scheme,
92 plaintext,
93 associated_data,
94 )?;
95
96 Ok(SealedMessage {
97 ephemeral_key: $ephemeral_variant(ephemeral),
98 ciphertext,
99 })
100 }
101 )*
102 }
103 }
104 };
105}
106
107macro_rules! impl_unseal_bytes_with_associated_data {
109 ($($variant:path => $crypto_box:ty, $ephemeral_variant:path;)*) => {
110 pub fn unseal_bytes_with_associated_data(
119 &self,
120 sealed_message: SealedMessage,
121 associated_data: &[u8],
122 ) -> Result<Vec<u8>, IesError> {
123 let self_algo = self.scheme() as u8;
125 let msg_algo = sealed_message.ephemeral_key.scheme() as u8;
126
127 let compatible = self_algo == msg_algo;
128 if !compatible {
129 return Err(IesError::SchemeMismatch);
130 }
131
132 let SealedMessage { ephemeral_key, ciphertext } = sealed_message;
133
134 match (self, ephemeral_key) {
135 $(
136 ($variant(key), $ephemeral_variant(ephemeral)) => {
137 <$crypto_box>::unseal_bytes_with_associated_data(
138 key,
139 &ephemeral,
140 self.scheme(),
141 &ciphertext,
142 associated_data,
143 )
144 }
145 )*
146 _ => Err(IesError::SchemeMismatch),
147 }
148 }
149 };
150}
151
152macro_rules! impl_unseal_elements_with_associated_data {
154 ($($variant:path => $crypto_box:ty, $ephemeral_variant:path;)*) => {
155 pub fn unseal_elements_with_associated_data(
164 &self,
165 sealed_message: SealedMessage,
166 associated_data: &[Felt],
167 ) -> Result<Vec<Felt>, IesError> {
168 let self_algo = self.scheme() as u8;
170 let msg_algo = sealed_message.ephemeral_key.scheme() as u8;
171
172 let compatible = self_algo == msg_algo;
173 if !compatible {
174 return Err(IesError::SchemeMismatch);
175 }
176
177 let SealedMessage { ephemeral_key, ciphertext } = sealed_message;
178
179 match (self, ephemeral_key) {
180 $(
181 ($variant(key), $ephemeral_variant(ephemeral)) => {
182 <$crypto_box>::unseal_elements_with_associated_data(
183 key,
184 &ephemeral,
185 self.scheme(),
186 &ciphertext,
187 associated_data,
188 )
189 }
190 )*
191 _ => Err(IesError::SchemeMismatch),
192 }
193 }
194 };
195}
196
197#[derive(Debug, Clone, PartialEq, Eq)]
202pub enum SealingKey {
203 K256XChaCha20Poly1305(crate::dsa::ecdsa_k256_keccak::PublicKey),
204 X25519XChaCha20Poly1305(crate::dsa::eddsa_25519_sha512::PublicKey),
205 K256AeadPoseidon2(crate::dsa::ecdsa_k256_keccak::PublicKey),
206 X25519AeadPoseidon2(crate::dsa::eddsa_25519_sha512::PublicKey),
207}
208
209impl SealingKey {
210 pub fn scheme(&self) -> IesScheme {
212 match self {
213 SealingKey::K256XChaCha20Poly1305(_) => IesScheme::K256XChaCha20Poly1305,
214 SealingKey::X25519XChaCha20Poly1305(_) => IesScheme::X25519XChaCha20Poly1305,
215 SealingKey::K256AeadPoseidon2(_) => IesScheme::K256AeadPoseidon2,
216 SealingKey::X25519AeadPoseidon2(_) => IesScheme::X25519AeadPoseidon2,
217 }
218 }
219
220 pub fn seal_bytes<R: CryptoRng + RngCore>(
225 &self,
226 rng: &mut R,
227 plaintext: &[u8],
228 ) -> Result<SealedMessage, IesError> {
229 self.seal_bytes_with_associated_data(rng, plaintext, &[])
230 }
231
232 impl_seal_bytes_with_associated_data! {
233 SealingKey::K256XChaCha20Poly1305 => K256XChaCha20Poly1305, EphemeralPublicKey::K256XChaCha20Poly1305;
234 SealingKey::X25519XChaCha20Poly1305 => X25519XChaCha20Poly1305, EphemeralPublicKey::X25519XChaCha20Poly1305;
235 SealingKey::K256AeadPoseidon2 => K256AeadPoseidon2, EphemeralPublicKey::K256AeadPoseidon2;
236 SealingKey::X25519AeadPoseidon2 => X25519AeadPoseidon2, EphemeralPublicKey::X25519AeadPoseidon2;
237 }
238
239 pub fn seal_elements<R: CryptoRng + RngCore>(
244 &self,
245 rng: &mut R,
246 plaintext: &[Felt],
247 ) -> Result<SealedMessage, IesError> {
248 self.seal_elements_with_associated_data(rng, plaintext, &[])
249 }
250
251 impl_seal_elements_with_associated_data! {
252 SealingKey::K256XChaCha20Poly1305 => K256XChaCha20Poly1305, EphemeralPublicKey::K256XChaCha20Poly1305;
253 SealingKey::X25519XChaCha20Poly1305 => X25519XChaCha20Poly1305, EphemeralPublicKey::X25519XChaCha20Poly1305;
254 SealingKey::K256AeadPoseidon2 => K256AeadPoseidon2, EphemeralPublicKey::K256AeadPoseidon2;
255 SealingKey::X25519AeadPoseidon2 => X25519AeadPoseidon2, EphemeralPublicKey::X25519AeadPoseidon2;
256 }
257}
258
259impl fmt::Display for SealingKey {
260 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
261 write!(f, "{} sealing key", self.scheme())
262 }
263}
264
265impl Serializable for SealingKey {
266 fn write_into<W: ByteWriter>(&self, target: &mut W) {
267 target.write_u8(self.scheme().into());
268
269 match self {
270 SealingKey::K256XChaCha20Poly1305(key) => key.write_into(target),
271 SealingKey::X25519XChaCha20Poly1305(key) => key.write_into(target),
272 SealingKey::K256AeadPoseidon2(key) => key.write_into(target),
273 SealingKey::X25519AeadPoseidon2(key) => key.write_into(target),
274 }
275 }
276}
277
278impl Deserializable for SealingKey {
279 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
280 let scheme = IesScheme::try_from(source.read_u8()?)
281 .map_err(|_| DeserializationError::InvalidValue("Unsupported IES scheme".into()))?;
282
283 match scheme {
284 IesScheme::K256XChaCha20Poly1305 => {
285 let key = crate::dsa::ecdsa_k256_keccak::PublicKey::read_from(source)?;
286 Ok(SealingKey::K256XChaCha20Poly1305(key))
287 },
288 IesScheme::X25519XChaCha20Poly1305 => {
289 let key = crate::dsa::eddsa_25519_sha512::PublicKey::read_from(source)?;
290 Ok(SealingKey::X25519XChaCha20Poly1305(key))
291 },
292 IesScheme::K256AeadPoseidon2 => {
293 let key = crate::dsa::ecdsa_k256_keccak::PublicKey::read_from(source)?;
294 Ok(SealingKey::K256AeadPoseidon2(key))
295 },
296 IesScheme::X25519AeadPoseidon2 => {
297 let key = crate::dsa::eddsa_25519_sha512::PublicKey::read_from(source)?;
298 Ok(SealingKey::X25519AeadPoseidon2(key))
299 },
300 }
301 }
302}
303
304pub enum UnsealingKey {
309 K256XChaCha20Poly1305(crate::dsa::ecdsa_k256_keccak::SecretKey),
310 X25519XChaCha20Poly1305(crate::dsa::eddsa_25519_sha512::SecretKey),
311 K256AeadPoseidon2(crate::dsa::ecdsa_k256_keccak::SecretKey),
312 X25519AeadPoseidon2(crate::dsa::eddsa_25519_sha512::SecretKey),
313}
314
315impl UnsealingKey {
316 pub fn scheme(&self) -> IesScheme {
318 match self {
319 UnsealingKey::K256XChaCha20Poly1305(_) => IesScheme::K256XChaCha20Poly1305,
320 UnsealingKey::X25519XChaCha20Poly1305(_) => IesScheme::X25519XChaCha20Poly1305,
321 UnsealingKey::K256AeadPoseidon2(_) => IesScheme::K256AeadPoseidon2,
322 UnsealingKey::X25519AeadPoseidon2(_) => IesScheme::X25519AeadPoseidon2,
323 }
324 }
325
326 pub fn scheme_name(&self) -> &'static str {
328 self.scheme().name()
329 }
330
331 pub fn unseal_bytes(&self, sealed_message: SealedMessage) -> Result<Vec<u8>, IesError> {
336 self.unseal_bytes_with_associated_data(sealed_message, &[])
337 }
338
339 impl_unseal_bytes_with_associated_data! {
340 UnsealingKey::K256XChaCha20Poly1305 => K256XChaCha20Poly1305, EphemeralPublicKey::K256XChaCha20Poly1305;
341 UnsealingKey::X25519XChaCha20Poly1305 => X25519XChaCha20Poly1305, EphemeralPublicKey::X25519XChaCha20Poly1305;
342 UnsealingKey::K256AeadPoseidon2 => K256AeadPoseidon2, EphemeralPublicKey::K256AeadPoseidon2;
343 UnsealingKey::X25519AeadPoseidon2 => X25519AeadPoseidon2, EphemeralPublicKey::X25519AeadPoseidon2;
344 }
345
346 pub fn unseal_elements(&self, sealed_message: SealedMessage) -> Result<Vec<Felt>, IesError> {
351 self.unseal_elements_with_associated_data(sealed_message, &[])
352 }
353
354 impl_unseal_elements_with_associated_data! {
355 UnsealingKey::K256XChaCha20Poly1305 => K256XChaCha20Poly1305, EphemeralPublicKey::K256XChaCha20Poly1305;
356 UnsealingKey::X25519XChaCha20Poly1305 => X25519XChaCha20Poly1305, EphemeralPublicKey::X25519XChaCha20Poly1305;
357 UnsealingKey::K256AeadPoseidon2 => K256AeadPoseidon2, EphemeralPublicKey::K256AeadPoseidon2;
358 UnsealingKey::X25519AeadPoseidon2 => X25519AeadPoseidon2, EphemeralPublicKey::X25519AeadPoseidon2;
359 }
360}
361
362impl fmt::Display for UnsealingKey {
363 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
364 write!(f, "{} unsealing key", self.scheme())
365 }
366}
367
368impl Serializable for UnsealingKey {
369 fn write_into<W: ByteWriter>(&self, target: &mut W) {
370 target.write_u8(self.scheme().into());
371
372 match self {
373 UnsealingKey::K256XChaCha20Poly1305(key) => key.write_into(target),
374 UnsealingKey::X25519XChaCha20Poly1305(key) => key.write_into(target),
375 UnsealingKey::K256AeadPoseidon2(key) => key.write_into(target),
376 UnsealingKey::X25519AeadPoseidon2(key) => key.write_into(target),
377 }
378 }
379}
380
381impl Deserializable for UnsealingKey {
382 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
383 let scheme = IesScheme::try_from(source.read_u8()?)
384 .map_err(|_| DeserializationError::InvalidValue("Unsupported IES scheme".into()))?;
385
386 match scheme {
387 IesScheme::K256XChaCha20Poly1305 => {
388 let key = crate::dsa::ecdsa_k256_keccak::SecretKey::read_from(source)?;
389 Ok(UnsealingKey::K256XChaCha20Poly1305(key))
390 },
391 IesScheme::X25519XChaCha20Poly1305 => {
392 let key = crate::dsa::eddsa_25519_sha512::SecretKey::read_from(source)?;
393 Ok(UnsealingKey::X25519XChaCha20Poly1305(key))
394 },
395 IesScheme::K256AeadPoseidon2 => {
396 let key = crate::dsa::ecdsa_k256_keccak::SecretKey::read_from(source)?;
397 Ok(UnsealingKey::K256AeadPoseidon2(key))
398 },
399 IesScheme::X25519AeadPoseidon2 => {
400 let key = crate::dsa::eddsa_25519_sha512::SecretKey::read_from(source)?;
401 Ok(UnsealingKey::X25519AeadPoseidon2(key))
402 },
403 }
404 }
405}
406
407#[derive(Debug, Clone, PartialEq, Eq)]
412pub(super) enum EphemeralPublicKey {
413 K256XChaCha20Poly1305(crate::ecdh::k256::EphemeralPublicKey),
414 X25519XChaCha20Poly1305(crate::ecdh::x25519::EphemeralPublicKey),
415 K256AeadPoseidon2(crate::ecdh::k256::EphemeralPublicKey),
416 X25519AeadPoseidon2(crate::ecdh::x25519::EphemeralPublicKey),
417}
418
419impl EphemeralPublicKey {
420 pub fn scheme(&self) -> IesScheme {
422 match self {
423 EphemeralPublicKey::K256XChaCha20Poly1305(_) => IesScheme::K256XChaCha20Poly1305,
424 EphemeralPublicKey::X25519XChaCha20Poly1305(_) => IesScheme::X25519XChaCha20Poly1305,
425 EphemeralPublicKey::K256AeadPoseidon2(_) => IesScheme::K256AeadPoseidon2,
426 EphemeralPublicKey::X25519AeadPoseidon2(_) => IesScheme::X25519AeadPoseidon2,
427 }
428 }
429
430 pub fn to_bytes(&self) -> Vec<u8> {
432 match self {
433 EphemeralPublicKey::K256XChaCha20Poly1305(key) => key.to_bytes(),
434 EphemeralPublicKey::X25519XChaCha20Poly1305(key) => key.to_bytes(),
435 EphemeralPublicKey::K256AeadPoseidon2(key) => key.to_bytes(),
436 EphemeralPublicKey::X25519AeadPoseidon2(key) => key.to_bytes(),
437 }
438 }
439
440 pub fn from_bytes(scheme: IesScheme, bytes: &[u8]) -> Result<Self, IesError> {
442 let expected_len = match scheme {
443 IesScheme::K256XChaCha20Poly1305 | IesScheme::K256AeadPoseidon2 => {
444 K256_PUBLIC_KEY_BYTES
445 },
446 IesScheme::X25519XChaCha20Poly1305 | IesScheme::X25519AeadPoseidon2 => {
447 X25519_PUBLIC_KEY_BYTES
448 },
449 };
450
451 if bytes.len() != expected_len {
452 return Err(IesError::EphemeralPublicKeyDeserializationFailed);
453 }
454
455 match scheme {
456 IesScheme::K256XChaCha20Poly1305 => {
457 let key =
458 <K256 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
459 bytes,
460 expected_len,
461 )
462 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
463 Ok(EphemeralPublicKey::K256XChaCha20Poly1305(key))
464 },
465 IesScheme::K256AeadPoseidon2 => {
466 let key =
467 <K256 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
468 bytes,
469 expected_len,
470 )
471 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
472 Ok(EphemeralPublicKey::K256AeadPoseidon2(key))
473 },
474 IesScheme::X25519XChaCha20Poly1305 => {
475 let key =
476 <X25519 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
477 bytes,
478 expected_len,
479 )
480 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
481 Ok(EphemeralPublicKey::X25519XChaCha20Poly1305(key))
482 },
483 IesScheme::X25519AeadPoseidon2 => {
484 let key =
485 <X25519 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
486 bytes,
487 expected_len,
488 )
489 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
490 Ok(EphemeralPublicKey::X25519AeadPoseidon2(key))
491 },
492 }
493 }
494}