1use alloc::vec::Vec;
2use core::fmt;
3
4use rand::CryptoRng;
5
6use super::{IesError, IesScheme, crypto_box::CryptoBox, message::SealedMessage};
7use crate::{
8 Felt,
9 aead::{aead_eidos::AeadEidos, 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>;
29type K256AeadEidos = CryptoBox<K256, AeadEidos>;
31type X25519AeadEidos = CryptoBox<X25519, AeadEidos>;
33
34macro_rules! impl_seal_bytes_with_associated_data {
39 ($($variant:path => $crypto_box:ty, $ephemeral_variant:path;)*) => {
40 pub fn seal_bytes_with_associated_data<R: CryptoRng>(
46 &self,
47 rng: &mut R,
48 plaintext: &[u8],
49 associated_data: &[u8],
50 ) -> Result<SealedMessage, IesError> {
51 match self {
52 $(
53 $variant(key) => {
54 let scheme = self.scheme();
55 let (ciphertext, ephemeral) = <$crypto_box>::seal_bytes_with_associated_data(
56 rng,
57 key,
58 scheme,
59 plaintext,
60 associated_data,
61 )?;
62
63 Ok(SealedMessage {
64 ephemeral_key: $ephemeral_variant(ephemeral),
65 ciphertext,
66 })
67 }
68 )*
69 }
70 }
71 };
72}
73
74macro_rules! impl_seal_elements_with_associated_data {
76 ($($variant:path => $crypto_box:ty, $ephemeral_variant:path;)*) => {
77 pub fn seal_elements_with_associated_data<R: CryptoRng>(
83 &self,
84 rng: &mut R,
85 plaintext: &[Felt],
86 associated_data: &[Felt],
87 ) -> Result<SealedMessage, IesError> {
88 match self {
89 $(
90 $variant(key) => {
91 let scheme = self.scheme();
92 let (ciphertext, ephemeral) = <$crypto_box>::seal_elements_with_associated_data(
93 rng,
94 key,
95 scheme,
96 plaintext,
97 associated_data,
98 )?;
99
100 Ok(SealedMessage {
101 ephemeral_key: $ephemeral_variant(ephemeral),
102 ciphertext,
103 })
104 }
105 )*
106 }
107 }
108 };
109}
110
111macro_rules! impl_unseal_bytes_with_associated_data {
113 ($($variant:path => $crypto_box:ty, $ephemeral_variant:path;)*) => {
114 pub fn unseal_bytes_with_associated_data(
123 &self,
124 sealed_message: SealedMessage,
125 associated_data: &[u8],
126 ) -> Result<Vec<u8>, IesError> {
127 let self_algo = self.scheme() as u8;
128 let msg_algo = sealed_message.ephemeral_key.scheme() as u8;
129
130 let compatible = self_algo == msg_algo;
131 if !compatible {
132 return Err(IesError::SchemeMismatch);
133 }
134
135 let SealedMessage { ephemeral_key, ciphertext } = sealed_message;
136
137 match (self, ephemeral_key) {
138 $(
139 ($variant(key), $ephemeral_variant(ephemeral)) => {
140 <$crypto_box>::unseal_bytes_with_associated_data(
141 key,
142 &ephemeral,
143 self.scheme(),
144 &ciphertext,
145 associated_data,
146 )
147 }
148 )*
149 _ => Err(IesError::SchemeMismatch),
150 }
151 }
152 };
153}
154
155macro_rules! impl_unseal_elements_with_associated_data {
157 ($($variant:path => $crypto_box:ty, $ephemeral_variant:path;)*) => {
158 pub fn unseal_elements_with_associated_data(
167 &self,
168 sealed_message: SealedMessage,
169 associated_data: &[Felt],
170 ) -> Result<Vec<Felt>, IesError> {
171 let self_algo = self.scheme() as u8;
172 let msg_algo = sealed_message.ephemeral_key.scheme() as u8;
173
174 let compatible = self_algo == msg_algo;
175 if !compatible {
176 return Err(IesError::SchemeMismatch);
177 }
178
179 let SealedMessage { ephemeral_key, ciphertext } = sealed_message;
180
181 match (self, ephemeral_key) {
182 $(
183 ($variant(key), $ephemeral_variant(ephemeral)) => {
184 <$crypto_box>::unseal_elements_with_associated_data(
185 key,
186 &ephemeral,
187 self.scheme(),
188 &ciphertext,
189 associated_data,
190 )
191 }
192 )*
193 _ => Err(IesError::SchemeMismatch),
194 }
195 }
196 };
197}
198
199#[derive(Debug, Clone, PartialEq, Eq)]
204pub enum SealingKey {
205 K256XChaCha20Poly1305(crate::dsa::ecdsa_k256_keccak::PublicKey),
206 X25519XChaCha20Poly1305(crate::dsa::eddsa_25519_sha512::PublicKey),
207 K256AeadPoseidon2(crate::dsa::ecdsa_k256_keccak::PublicKey),
208 X25519AeadPoseidon2(crate::dsa::eddsa_25519_sha512::PublicKey),
209 K256AeadEidos(crate::dsa::ecdsa_k256_keccak::PublicKey),
210 X25519AeadEidos(crate::dsa::eddsa_25519_sha512::PublicKey),
211}
212
213impl SealingKey {
214 pub fn scheme(&self) -> IesScheme {
216 match self {
217 SealingKey::K256XChaCha20Poly1305(_) => IesScheme::K256XChaCha20Poly1305,
218 SealingKey::X25519XChaCha20Poly1305(_) => IesScheme::X25519XChaCha20Poly1305,
219 SealingKey::K256AeadPoseidon2(_) => IesScheme::K256AeadPoseidon2,
220 SealingKey::X25519AeadPoseidon2(_) => IesScheme::X25519AeadPoseidon2,
221 SealingKey::K256AeadEidos(_) => IesScheme::K256AeadEidos,
222 SealingKey::X25519AeadEidos(_) => IesScheme::X25519AeadEidos,
223 }
224 }
225
226 pub fn seal_bytes<R: CryptoRng>(
231 &self,
232 rng: &mut R,
233 plaintext: &[u8],
234 ) -> Result<SealedMessage, IesError> {
235 self.seal_bytes_with_associated_data(rng, plaintext, &[])
236 }
237
238 impl_seal_bytes_with_associated_data! {
239 SealingKey::K256XChaCha20Poly1305 => K256XChaCha20Poly1305, EphemeralPublicKey::K256XChaCha20Poly1305;
240 SealingKey::X25519XChaCha20Poly1305 => X25519XChaCha20Poly1305, EphemeralPublicKey::X25519XChaCha20Poly1305;
241 SealingKey::K256AeadPoseidon2 => K256AeadPoseidon2, EphemeralPublicKey::K256AeadPoseidon2;
242 SealingKey::X25519AeadPoseidon2 => X25519AeadPoseidon2, EphemeralPublicKey::X25519AeadPoseidon2;
243 SealingKey::K256AeadEidos => K256AeadEidos, EphemeralPublicKey::K256AeadEidos;
244 SealingKey::X25519AeadEidos => X25519AeadEidos, EphemeralPublicKey::X25519AeadEidos;
245 }
246
247 pub fn seal_elements<R: CryptoRng>(
252 &self,
253 rng: &mut R,
254 plaintext: &[Felt],
255 ) -> Result<SealedMessage, IesError> {
256 self.seal_elements_with_associated_data(rng, plaintext, &[])
257 }
258
259 impl_seal_elements_with_associated_data! {
260 SealingKey::K256XChaCha20Poly1305 => K256XChaCha20Poly1305, EphemeralPublicKey::K256XChaCha20Poly1305;
261 SealingKey::X25519XChaCha20Poly1305 => X25519XChaCha20Poly1305, EphemeralPublicKey::X25519XChaCha20Poly1305;
262 SealingKey::K256AeadPoseidon2 => K256AeadPoseidon2, EphemeralPublicKey::K256AeadPoseidon2;
263 SealingKey::X25519AeadPoseidon2 => X25519AeadPoseidon2, EphemeralPublicKey::X25519AeadPoseidon2;
264 SealingKey::K256AeadEidos => K256AeadEidos, EphemeralPublicKey::K256AeadEidos;
265 SealingKey::X25519AeadEidos => X25519AeadEidos, EphemeralPublicKey::X25519AeadEidos;
266 }
267}
268
269impl fmt::Display for SealingKey {
270 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
271 write!(f, "{} sealing key", self.scheme())
272 }
273}
274
275impl Serializable for SealingKey {
276 fn write_into<W: ByteWriter>(&self, target: &mut W) {
277 target.write_u8(self.scheme().into());
278
279 match self {
280 SealingKey::K256XChaCha20Poly1305(key) => key.write_into(target),
281 SealingKey::X25519XChaCha20Poly1305(key) => key.write_into(target),
282 SealingKey::K256AeadPoseidon2(key) => key.write_into(target),
283 SealingKey::X25519AeadPoseidon2(key) => key.write_into(target),
284 SealingKey::K256AeadEidos(key) => key.write_into(target),
285 SealingKey::X25519AeadEidos(key) => key.write_into(target),
286 }
287 }
288}
289
290impl Deserializable for SealingKey {
291 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
292 let scheme = IesScheme::try_from(source.read_u8()?)
293 .map_err(|_| DeserializationError::InvalidValue("Unsupported IES scheme".into()))?;
294
295 match scheme {
296 IesScheme::K256XChaCha20Poly1305 => {
297 let key = crate::dsa::ecdsa_k256_keccak::PublicKey::read_from(source)?;
298 Ok(SealingKey::K256XChaCha20Poly1305(key))
299 },
300 IesScheme::X25519XChaCha20Poly1305 => {
301 let key = crate::dsa::eddsa_25519_sha512::PublicKey::read_from(source)?;
302 Ok(SealingKey::X25519XChaCha20Poly1305(key))
303 },
304 IesScheme::K256AeadPoseidon2 => {
305 let key = crate::dsa::ecdsa_k256_keccak::PublicKey::read_from(source)?;
306 Ok(SealingKey::K256AeadPoseidon2(key))
307 },
308 IesScheme::X25519AeadPoseidon2 => {
309 let key = crate::dsa::eddsa_25519_sha512::PublicKey::read_from(source)?;
310 Ok(SealingKey::X25519AeadPoseidon2(key))
311 },
312 IesScheme::K256AeadEidos => {
313 let key = crate::dsa::ecdsa_k256_keccak::PublicKey::read_from(source)?;
314 Ok(SealingKey::K256AeadEidos(key))
315 },
316 IesScheme::X25519AeadEidos => {
317 let key = crate::dsa::eddsa_25519_sha512::PublicKey::read_from(source)?;
318 Ok(SealingKey::X25519AeadEidos(key))
319 },
320 }
321 }
322}
323
324pub enum UnsealingKey {
329 K256XChaCha20Poly1305(crate::dsa::ecdsa_k256_keccak::KeyExchangeKey),
330 X25519XChaCha20Poly1305(crate::dsa::eddsa_25519_sha512::KeyExchangeKey),
331 K256AeadPoseidon2(crate::dsa::ecdsa_k256_keccak::KeyExchangeKey),
332 X25519AeadPoseidon2(crate::dsa::eddsa_25519_sha512::KeyExchangeKey),
333 K256AeadEidos(crate::dsa::ecdsa_k256_keccak::KeyExchangeKey),
334 X25519AeadEidos(crate::dsa::eddsa_25519_sha512::KeyExchangeKey),
335}
336
337impl UnsealingKey {
338 pub fn scheme(&self) -> IesScheme {
340 match self {
341 UnsealingKey::K256XChaCha20Poly1305(_) => IesScheme::K256XChaCha20Poly1305,
342 UnsealingKey::X25519XChaCha20Poly1305(_) => IesScheme::X25519XChaCha20Poly1305,
343 UnsealingKey::K256AeadPoseidon2(_) => IesScheme::K256AeadPoseidon2,
344 UnsealingKey::X25519AeadPoseidon2(_) => IesScheme::X25519AeadPoseidon2,
345 UnsealingKey::K256AeadEidos(_) => IesScheme::K256AeadEidos,
346 UnsealingKey::X25519AeadEidos(_) => IesScheme::X25519AeadEidos,
347 }
348 }
349
350 pub fn scheme_name(&self) -> &'static str {
352 self.scheme().name()
353 }
354
355 pub fn unseal_bytes(&self, sealed_message: SealedMessage) -> Result<Vec<u8>, IesError> {
360 self.unseal_bytes_with_associated_data(sealed_message, &[])
361 }
362
363 impl_unseal_bytes_with_associated_data! {
364 UnsealingKey::K256XChaCha20Poly1305 => K256XChaCha20Poly1305, EphemeralPublicKey::K256XChaCha20Poly1305;
365 UnsealingKey::X25519XChaCha20Poly1305 => X25519XChaCha20Poly1305, EphemeralPublicKey::X25519XChaCha20Poly1305;
366 UnsealingKey::K256AeadPoseidon2 => K256AeadPoseidon2, EphemeralPublicKey::K256AeadPoseidon2;
367 UnsealingKey::X25519AeadPoseidon2 => X25519AeadPoseidon2, EphemeralPublicKey::X25519AeadPoseidon2;
368 UnsealingKey::K256AeadEidos => K256AeadEidos, EphemeralPublicKey::K256AeadEidos;
369 UnsealingKey::X25519AeadEidos => X25519AeadEidos, EphemeralPublicKey::X25519AeadEidos;
370 }
371
372 pub fn unseal_elements(&self, sealed_message: SealedMessage) -> Result<Vec<Felt>, IesError> {
377 self.unseal_elements_with_associated_data(sealed_message, &[])
378 }
379
380 impl_unseal_elements_with_associated_data! {
381 UnsealingKey::K256XChaCha20Poly1305 => K256XChaCha20Poly1305, EphemeralPublicKey::K256XChaCha20Poly1305;
382 UnsealingKey::X25519XChaCha20Poly1305 => X25519XChaCha20Poly1305, EphemeralPublicKey::X25519XChaCha20Poly1305;
383 UnsealingKey::K256AeadPoseidon2 => K256AeadPoseidon2, EphemeralPublicKey::K256AeadPoseidon2;
384 UnsealingKey::X25519AeadPoseidon2 => X25519AeadPoseidon2, EphemeralPublicKey::X25519AeadPoseidon2;
385 UnsealingKey::K256AeadEidos => K256AeadEidos, EphemeralPublicKey::K256AeadEidos;
386 UnsealingKey::X25519AeadEidos => X25519AeadEidos, EphemeralPublicKey::X25519AeadEidos;
387 }
388}
389
390impl fmt::Display for UnsealingKey {
391 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
392 write!(f, "{} unsealing key", self.scheme())
393 }
394}
395
396impl Serializable for UnsealingKey {
397 fn write_into<W: ByteWriter>(&self, target: &mut W) {
398 target.write_u8(self.scheme().into());
399
400 match self {
401 UnsealingKey::K256XChaCha20Poly1305(key) => key.write_into(target),
402 UnsealingKey::X25519XChaCha20Poly1305(key) => key.write_into(target),
403 UnsealingKey::K256AeadPoseidon2(key) => key.write_into(target),
404 UnsealingKey::X25519AeadPoseidon2(key) => key.write_into(target),
405 UnsealingKey::K256AeadEidos(key) => key.write_into(target),
406 UnsealingKey::X25519AeadEidos(key) => key.write_into(target),
407 }
408 }
409}
410
411impl Deserializable for UnsealingKey {
412 fn read_from<R: ByteReader>(source: &mut R) -> Result<Self, DeserializationError> {
413 let scheme = IesScheme::try_from(source.read_u8()?)
414 .map_err(|_| DeserializationError::InvalidValue("Unsupported IES scheme".into()))?;
415
416 match scheme {
417 IesScheme::K256XChaCha20Poly1305 => {
418 let key = crate::dsa::ecdsa_k256_keccak::KeyExchangeKey::read_from(source)?;
419 Ok(UnsealingKey::K256XChaCha20Poly1305(key))
420 },
421 IesScheme::X25519XChaCha20Poly1305 => {
422 let key = crate::dsa::eddsa_25519_sha512::KeyExchangeKey::read_from(source)?;
423 Ok(UnsealingKey::X25519XChaCha20Poly1305(key))
424 },
425 IesScheme::K256AeadPoseidon2 => {
426 let key = crate::dsa::ecdsa_k256_keccak::KeyExchangeKey::read_from(source)?;
427 Ok(UnsealingKey::K256AeadPoseidon2(key))
428 },
429 IesScheme::X25519AeadPoseidon2 => {
430 let key = crate::dsa::eddsa_25519_sha512::KeyExchangeKey::read_from(source)?;
431 Ok(UnsealingKey::X25519AeadPoseidon2(key))
432 },
433 IesScheme::K256AeadEidos => {
434 let key = crate::dsa::ecdsa_k256_keccak::KeyExchangeKey::read_from(source)?;
435 Ok(UnsealingKey::K256AeadEidos(key))
436 },
437 IesScheme::X25519AeadEidos => {
438 let key = crate::dsa::eddsa_25519_sha512::KeyExchangeKey::read_from(source)?;
439 Ok(UnsealingKey::X25519AeadEidos(key))
440 },
441 }
442 }
443}
444
445#[derive(Debug, Clone, PartialEq, Eq)]
450pub(super) enum EphemeralPublicKey {
451 K256XChaCha20Poly1305(crate::ecdh::k256::EphemeralPublicKey),
452 X25519XChaCha20Poly1305(crate::ecdh::x25519::EphemeralPublicKey),
453 K256AeadPoseidon2(crate::ecdh::k256::EphemeralPublicKey),
454 X25519AeadPoseidon2(crate::ecdh::x25519::EphemeralPublicKey),
455 K256AeadEidos(crate::ecdh::k256::EphemeralPublicKey),
456 X25519AeadEidos(crate::ecdh::x25519::EphemeralPublicKey),
457}
458
459impl EphemeralPublicKey {
460 pub fn scheme(&self) -> IesScheme {
462 match self {
463 EphemeralPublicKey::K256XChaCha20Poly1305(_) => IesScheme::K256XChaCha20Poly1305,
464 EphemeralPublicKey::X25519XChaCha20Poly1305(_) => IesScheme::X25519XChaCha20Poly1305,
465 EphemeralPublicKey::K256AeadPoseidon2(_) => IesScheme::K256AeadPoseidon2,
466 EphemeralPublicKey::X25519AeadPoseidon2(_) => IesScheme::X25519AeadPoseidon2,
467 EphemeralPublicKey::K256AeadEidos(_) => IesScheme::K256AeadEidos,
468 EphemeralPublicKey::X25519AeadEidos(_) => IesScheme::X25519AeadEidos,
469 }
470 }
471
472 pub fn to_bytes(&self) -> Vec<u8> {
474 match self {
475 EphemeralPublicKey::K256XChaCha20Poly1305(key) => key.to_bytes(),
476 EphemeralPublicKey::X25519XChaCha20Poly1305(key) => key.to_bytes(),
477 EphemeralPublicKey::K256AeadPoseidon2(key) => key.to_bytes(),
478 EphemeralPublicKey::X25519AeadPoseidon2(key) => key.to_bytes(),
479 EphemeralPublicKey::K256AeadEidos(key) => key.to_bytes(),
480 EphemeralPublicKey::X25519AeadEidos(key) => key.to_bytes(),
481 }
482 }
483
484 pub fn from_bytes(scheme: IesScheme, bytes: &[u8]) -> Result<Self, IesError> {
486 let expected_len = match scheme {
487 IesScheme::K256XChaCha20Poly1305
488 | IesScheme::K256AeadPoseidon2
489 | IesScheme::K256AeadEidos => K256_PUBLIC_KEY_BYTES,
490 IesScheme::X25519XChaCha20Poly1305
491 | IesScheme::X25519AeadPoseidon2
492 | IesScheme::X25519AeadEidos => X25519_PUBLIC_KEY_BYTES,
493 };
494
495 if bytes.len() != expected_len {
496 return Err(IesError::EphemeralPublicKeyDeserializationFailed);
497 }
498
499 match scheme {
500 IesScheme::K256XChaCha20Poly1305 => {
501 let key =
502 <K256 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
503 bytes,
504 expected_len,
505 )
506 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
507 Ok(EphemeralPublicKey::K256XChaCha20Poly1305(key))
508 },
509 IesScheme::K256AeadPoseidon2 => {
510 let key =
511 <K256 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
512 bytes,
513 expected_len,
514 )
515 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
516 Ok(EphemeralPublicKey::K256AeadPoseidon2(key))
517 },
518 IesScheme::K256AeadEidos => {
519 let key =
520 <K256 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
521 bytes,
522 expected_len,
523 )
524 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
525 Ok(EphemeralPublicKey::K256AeadEidos(key))
526 },
527 IesScheme::X25519XChaCha20Poly1305 => {
528 let key =
529 <X25519 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
530 bytes,
531 expected_len,
532 )
533 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
534 Ok(EphemeralPublicKey::X25519XChaCha20Poly1305(key))
535 },
536 IesScheme::X25519AeadPoseidon2 => {
537 let key =
538 <X25519 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
539 bytes,
540 expected_len,
541 )
542 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
543 Ok(EphemeralPublicKey::X25519AeadPoseidon2(key))
544 },
545 IesScheme::X25519AeadEidos => {
546 let key =
547 <X25519 as KeyAgreementScheme>::EphemeralPublicKey::read_from_bytes_with_budget(
548 bytes,
549 expected_len,
550 )
551 .map_err(|_| IesError::EphemeralPublicKeyDeserializationFailed)?;
552 Ok(EphemeralPublicKey::X25519AeadEidos(key))
553 },
554 }
555 }
556}