Skip to main content

ml_kem/
encapsulation_key.rs

1use crate::{
2    B32, SharedKey,
3    crypto::{G, H},
4    kem::{InvalidKey, Kem, Key, KeyExport, KeySizeUser, TryKeyInit},
5    param::{EncapsulationKeySize, KemParams},
6    pke::EncryptionKey,
7};
8use array::sizes::U32;
9use kem::{Ciphertext, Encapsulate, Generate};
10use rand_core::CryptoRng;
11
12/// An `EncapsulationKey` provides the ability to encapsulate a shared key so that it can only be
13/// decapsulated by the holder of the corresponding decapsulation key.
14#[derive(Clone, Debug)]
15pub struct EncapsulationKey<P>
16where
17    P: KemParams,
18{
19    ek_pke: EncryptionKey<P>,
20    h: B32,
21}
22
23impl<P> EncapsulationKey<P>
24where
25    P: Kem<SharedKeySize = U32> + KemParams,
26{
27    /// Create a new [`EncapsulationKey`] from its serialized form.
28    ///
29    /// # Errors
30    /// If the key failed validation during decoding.
31    pub fn new(encapsulation_key: &Key<Self>) -> Result<Self, InvalidKey> {
32        EncryptionKey::from_bytes(encapsulation_key)
33            .map(Self::from_encryption_key)
34            .map_err(|_| InvalidKey)
35    }
36
37    /// Encapsulates with the given randomness. This is useful for testing against known vectors.
38    ///
39    /// # Warning
40    /// Do NOT use this function unless you know what you're doing. If you fail to use all uniform
41    /// random bytes even once, you can have catastrophic security failure.
42    #[cfg_attr(not(feature = "hazmat"), doc(hidden))]
43    pub fn encapsulate_deterministic(&self, m: &B32) -> (Ciphertext<P>, SharedKey) {
44        let (K, r) = G(&[m, &self.h]);
45        let c = self.ek_pke.encrypt(m, &r);
46        (c, K)
47    }
48
49    /// Convert from an `EncryptionKey`.
50    pub(crate) fn from_encryption_key(ek_pke: EncryptionKey<P>) -> Self {
51        let h = H(ek_pke.to_bytes());
52        Self { ek_pke, h }
53    }
54
55    /// Borrow the encryption key.
56    pub(crate) fn ek_pke(&self) -> &EncryptionKey<P> {
57        &self.ek_pke
58    }
59
60    /// Retrieve the hash of the encryption key.
61    pub(crate) fn h(&self) -> B32 {
62        self.h
63    }
64}
65
66impl<P> Encapsulate for EncapsulationKey<P>
67where
68    P: Kem + KemParams,
69{
70    type Kem = P;
71
72    fn encapsulate_with_rng<R>(&self, rng: &mut R) -> (Ciphertext<P>, SharedKey)
73    where
74        R: CryptoRng + ?Sized,
75    {
76        let m = B32::generate_from_rng(rng);
77        self.encapsulate_deterministic(&m)
78    }
79}
80
81impl<P> KeyExport for EncapsulationKey<P>
82where
83    P: KemParams,
84{
85    fn to_bytes(&self) -> Key<Self> {
86        self.ek_pke.to_bytes()
87    }
88}
89
90impl<P> KeySizeUser for EncapsulationKey<P>
91where
92    P: KemParams,
93{
94    type KeySize = EncapsulationKeySize<P>;
95}
96
97impl<P> TryKeyInit for EncapsulationKey<P>
98where
99    P: KemParams,
100{
101    fn new(encapsulation_key: &Key<Self>) -> Result<Self, InvalidKey> {
102        Self::new(encapsulation_key)
103    }
104}
105
106impl<P> Eq for EncapsulationKey<P> where P: KemParams {}
107impl<P> PartialEq for EncapsulationKey<P>
108where
109    P: KemParams,
110{
111    fn eq(&self, other: &Self) -> bool {
112        // Handwritten to avoid derive putting `Eq` bounds on `KemParams`
113        self.ek_pke == other.ek_pke && self.h == other.h
114    }
115}