bc_components/mlkem/
mlkem_public_key.rs

1use dcbor::prelude::*;
2use pqcrypto_mlkem::*;
3use pqcrypto_traits::kem::{PublicKey, SharedSecret};
4
5use super::{MLKEM, MLKEMCiphertext};
6use crate::{Error, Result, SymmetricKey, tags};
7
8/// A public key for the ML-KEM post-quantum key encapsulation mechanism.
9///
10/// `MLKEMPublicKey` represents a public key that can be used to encapsulate
11/// shared secrets using the ML-KEM (Module Lattice-based Key Encapsulation
12/// Mechanism) post-quantum algorithm. It supports multiple security levels
13/// through the variants:
14///
15/// - `MLKEM512`: NIST security level 1 (roughly equivalent to AES-128), 800
16///   bytes
17/// - `MLKEM768`: NIST security level 3 (roughly equivalent to AES-192), 1184
18///   bytes
19/// - `MLKEM1024`: NIST security level 5 (roughly equivalent to AES-256), 1568
20///   bytes
21///
22/// # Examples
23///
24/// ```
25/// use bc_components::MLKEM;
26///
27/// // Generate a keypair
28/// let (private_key, public_key) = MLKEM::MLKEM512.keypair();
29///
30/// // Encapsulate a shared secret using the public key
31/// let (shared_secret, ciphertext) =
32///     public_key.encapsulate_new_shared_secret();
33/// ```
34#[derive(Clone)]
35pub enum MLKEMPublicKey {
36    /// An ML-KEM-512 public key (NIST security level 1)
37    MLKEM512(Box<mlkem512::PublicKey>),
38    /// An ML-KEM-768 public key (NIST security level 3)
39    MLKEM768(Box<mlkem768::PublicKey>),
40    /// An ML-KEM-1024 public key (NIST security level 5)
41    MLKEM1024(Box<mlkem1024::PublicKey>),
42}
43
44/// Implements equality comparison for ML-KEM public keys.
45impl PartialEq for MLKEMPublicKey {
46    /// Compares two ML-KEM public keys for equality.
47    ///
48    /// Two ML-KEM public keys are equal if they have the same security level
49    /// and the same raw byte representation.
50    fn eq(&self, other: &Self) -> bool {
51        self.level() == other.level() && self.as_bytes() == other.as_bytes()
52    }
53}
54
55impl Eq for MLKEMPublicKey {}
56
57/// Implements hashing for ML-KEM public keys.
58impl std::hash::Hash for MLKEMPublicKey {
59    /// Hashes both the security level and the raw bytes of the public key.
60    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
61        self.level().hash(state);
62        self.as_bytes().hash(state);
63    }
64}
65
66impl MLKEMPublicKey {
67    /// Returns the security level of this ML-KEM public key.
68    pub fn level(&self) -> MLKEM {
69        match self {
70            MLKEMPublicKey::MLKEM512(_) => MLKEM::MLKEM512,
71            MLKEMPublicKey::MLKEM768(_) => MLKEM::MLKEM768,
72            MLKEMPublicKey::MLKEM1024(_) => MLKEM::MLKEM1024,
73        }
74    }
75
76    /// Returns the size of this ML-KEM public key in bytes.
77    pub fn size(&self) -> usize { self.level().public_key_size() }
78
79    /// Returns the raw bytes of this ML-KEM public key.
80    pub fn as_bytes(&self) -> &[u8] { self.as_ref() }
81
82    /// Creates an ML-KEM public key from raw bytes and a security level.
83    ///
84    /// # Parameters
85    ///
86    /// * `level` - The security level of the key.
87    /// * `bytes` - The raw bytes of the key.
88    ///
89    /// # Returns
90    ///
91    /// An `MLKEMPublicKey` if the bytes represent a valid key for the given
92    /// level, or an error otherwise.
93    ///
94    /// # Errors
95    ///
96    /// Returns an error if the bytes do not represent a valid ML-KEM public key
97    /// for the specified security level.
98    pub fn from_bytes(level: MLKEM, bytes: &[u8]) -> Result<Self> {
99        match level {
100            MLKEM::MLKEM512 => Ok(MLKEMPublicKey::MLKEM512(Box::new(
101                mlkem512::PublicKey::from_bytes(bytes)
102                    .map_err(|e| Error::post_quantum(e.to_string()))?,
103            ))),
104            MLKEM::MLKEM768 => Ok(MLKEMPublicKey::MLKEM768(Box::new(
105                mlkem768::PublicKey::from_bytes(bytes)
106                    .map_err(|e| Error::post_quantum(e.to_string()))?,
107            ))),
108            MLKEM::MLKEM1024 => Ok(MLKEMPublicKey::MLKEM1024(Box::new(
109                mlkem1024::PublicKey::from_bytes(bytes)
110                    .map_err(|e| Error::post_quantum(e.to_string()))?,
111            ))),
112        }
113    }
114
115    /// Encapsulates a new shared secret using this public key.
116    ///
117    /// This method generates a random shared secret and encapsulates it using
118    /// this public key, producing a ciphertext that can only be decapsulated
119    /// by the corresponding private key.
120    ///
121    /// # Returns
122    ///
123    /// A tuple containing:
124    /// - A `SymmetricKey` with the shared secret (32 bytes)
125    /// - An `MLKEMCiphertext` with the encapsulated shared secret
126    ///
127    /// # Examples
128    ///
129    /// ```
130    /// use bc_components::MLKEM;
131    ///
132    /// // Generate a keypair
133    /// let (private_key, public_key) = MLKEM::MLKEM512.keypair();
134    ///
135    /// // Encapsulate a shared secret
136    /// let (shared_secret, ciphertext) =
137    ///     public_key.encapsulate_new_shared_secret();
138    ///
139    /// // The private key holder can decapsulate the same shared secret
140    /// let decapsulated_secret =
141    ///     private_key.decapsulate_shared_secret(&ciphertext).unwrap();
142    /// assert_eq!(shared_secret, decapsulated_secret);
143    /// ```
144    pub fn encapsulate_new_shared_secret(
145        &self,
146    ) -> (SymmetricKey, MLKEMCiphertext) {
147        match self {
148            MLKEMPublicKey::MLKEM512(pk) => {
149                let (ss, ct) = mlkem512::encapsulate(pk.as_ref());
150                (
151                    SymmetricKey::from_data_ref(ss.as_bytes()).unwrap(),
152                    MLKEMCiphertext::MLKEM512(ct.into()),
153                )
154            }
155            MLKEMPublicKey::MLKEM768(pk) => {
156                let (ss, ct) = mlkem768::encapsulate(pk.as_ref());
157                (
158                    SymmetricKey::from_data_ref(ss.as_bytes()).unwrap(),
159                    MLKEMCiphertext::MLKEM768(ct.into()),
160                )
161            }
162            MLKEMPublicKey::MLKEM1024(pk) => {
163                let (ss, ct) = mlkem1024::encapsulate(pk.as_ref());
164                (
165                    SymmetricKey::from_data_ref(ss.as_bytes()).unwrap(),
166                    MLKEMCiphertext::MLKEM1024(ct.into()),
167                )
168            }
169        }
170    }
171}
172
173impl AsRef<[u8]> for MLKEMPublicKey {
174    /// Returns the raw bytes of the public key.
175    fn as_ref(&self) -> &[u8] {
176        match self {
177            MLKEMPublicKey::MLKEM512(pk) => pk.as_ref().as_bytes(),
178            MLKEMPublicKey::MLKEM768(pk) => pk.as_ref().as_bytes(),
179            MLKEMPublicKey::MLKEM1024(pk) => pk.as_ref().as_bytes(),
180        }
181    }
182}
183
184/// Provides debug formatting for ML-KEM public keys.
185impl std::fmt::Debug for MLKEMPublicKey {
186    /// Formats the public key as a string for debugging purposes.
187    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
188        match self {
189            MLKEMPublicKey::MLKEM512(_) => f.write_str("MLKEM512PublicKey"),
190            MLKEMPublicKey::MLKEM768(_) => f.write_str("MLKEM768PublicKey"),
191            MLKEMPublicKey::MLKEM1024(_) => f.write_str("MLKEM1024PublicKey"),
192        }
193    }
194}
195
196/// Defines CBOR tags for ML-KEM public keys.
197impl CBORTagged for MLKEMPublicKey {
198    /// Returns the CBOR tag for ML-KEM public keys.
199    fn cbor_tags() -> Vec<Tag> {
200        tags_for_values(&[tags::TAG_MLKEM_PUBLIC_KEY])
201    }
202}
203
204/// Converts an `MLKEMPublicKey` to CBOR.
205impl From<MLKEMPublicKey> for CBOR {
206    /// Converts to tagged CBOR.
207    fn from(value: MLKEMPublicKey) -> Self { value.tagged_cbor() }
208}
209
210/// Implements CBOR encoding for ML-KEM public keys.
211impl CBORTaggedEncodable for MLKEMPublicKey {
212    /// Creates the untagged CBOR representation as an array with level and key
213    /// bytes.
214    fn untagged_cbor(&self) -> CBOR {
215        vec![self.level().into(), CBOR::to_byte_string(self.as_bytes())].into()
216    }
217}
218
219/// Attempts to convert CBOR to an `MLKEMPublicKey`.
220impl TryFrom<CBOR> for MLKEMPublicKey {
221    type Error = dcbor::Error;
222
223    /// Converts from tagged CBOR.
224    fn try_from(cbor: CBOR) -> dcbor::Result<Self> {
225        Self::from_tagged_cbor(cbor)
226    }
227}
228
229/// Implements CBOR decoding for ML-KEM public keys.
230impl CBORTaggedDecodable for MLKEMPublicKey {
231    /// Creates an `MLKEMPublicKey` from untagged CBOR.
232    ///
233    /// # Errors
234    /// Returns an error if the CBOR value doesn't represent a valid ML-KEM
235    /// public key.
236    fn from_untagged_cbor(untagged_cbor: CBOR) -> dcbor::Result<Self> {
237        match untagged_cbor.as_case() {
238            CBORCase::Array(elements) => {
239                if elements.len() != 2 {
240                    return Err("MLKEMPublicKey must have two elements".into());
241                }
242
243                let level = MLKEM::try_from(elements[0].clone())?;
244                let data = CBOR::try_into_byte_string(elements[1].clone())?;
245                Ok(MLKEMPublicKey::from_bytes(level, &data)?)
246            }
247            _ => Err("MLKEMPublicKey must be an array".into()),
248        }
249    }
250}