1use pqcrypto_kyber::kyber1024;
6use pqcrypto_traits::kem::{PublicKey, SecretKey, Ciphertext, SharedSecret};
7use zeroize::Zeroize;
8use crate::{CryptoError, Result};
9
10#[derive(Clone)]
12pub struct KyberPublicKey(kyber1024::PublicKey);
13
14impl KyberPublicKey {
15 pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
17 if bytes.len() != kyber1024::public_key_bytes() {
18 return Err(CryptoError::InvalidKeyLength {
19 expected: kyber1024::public_key_bytes(),
20 actual: bytes.len(),
21 });
22 }
23 let key = kyber1024::PublicKey::from_bytes(bytes)
24 .map_err(|_| CryptoError::InvalidKeyLength {
25 expected: kyber1024::public_key_bytes(),
26 actual: bytes.len(),
27 })?;
28 Ok(Self(key))
29 }
30
31 pub fn as_bytes(&self) -> &[u8] {
33 self.0.as_bytes()
34 }
35
36 pub fn size() -> usize {
38 kyber1024::public_key_bytes()
39 }
40}
41
42pub struct KyberSecretKey(kyber1024::SecretKey);
44
45impl KyberSecretKey {
46 pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
48 if bytes.len() != kyber1024::secret_key_bytes() {
49 return Err(CryptoError::InvalidKeyLength {
50 expected: kyber1024::secret_key_bytes(),
51 actual: bytes.len(),
52 });
53 }
54 let key = kyber1024::SecretKey::from_bytes(bytes)
55 .map_err(|_| CryptoError::InvalidKeyLength {
56 expected: kyber1024::secret_key_bytes(),
57 actual: bytes.len(),
58 })?;
59 Ok(Self(key))
60 }
61
62 pub fn as_bytes(&self) -> &[u8] {
64 self.0.as_bytes()
65 }
66
67 pub fn size() -> usize {
69 kyber1024::secret_key_bytes()
70 }
71}
72
73impl Drop for KyberSecretKey {
75 fn drop(&mut self) {
76 }
80}
81
82pub struct KyberKeyPair {
84 pub public_key: KyberPublicKey,
85 pub secret_key: KyberSecretKey,
86}
87
88#[derive(Clone, Zeroize)]
90#[zeroize(drop)]
91pub struct KyberSharedSecret([u8; 32]);
92
93impl KyberSharedSecret {
94 pub fn as_bytes(&self) -> &[u8] {
96 &self.0
97 }
98
99 pub fn size() -> usize {
101 kyber1024::shared_secret_bytes()
102 }
103}
104
105pub struct KyberCiphertext(kyber1024::Ciphertext);
107
108impl KyberCiphertext {
109 pub fn from_bytes(bytes: &[u8]) -> Result<Self> {
111 if bytes.len() != kyber1024::ciphertext_bytes() {
112 return Err(CryptoError::InvalidKeyLength {
113 expected: kyber1024::ciphertext_bytes(),
114 actual: bytes.len(),
115 });
116 }
117 let ct = kyber1024::Ciphertext::from_bytes(bytes)
118 .map_err(|_| CryptoError::InvalidKeyLength {
119 expected: kyber1024::ciphertext_bytes(),
120 actual: bytes.len(),
121 })?;
122 Ok(Self(ct))
123 }
124
125 pub fn as_bytes(&self) -> &[u8] {
127 self.0.as_bytes()
128 }
129
130 pub fn size() -> usize {
132 kyber1024::ciphertext_bytes()
133 }
134}
135
136pub struct KyberCrypto;
138
139impl KyberCrypto {
140 pub fn generate_keypair() -> KyberKeyPair {
142 let (pk, sk) = kyber1024::keypair();
143 KyberKeyPair {
144 public_key: KyberPublicKey(pk),
145 secret_key: KyberSecretKey(sk),
146 }
147 }
148
149 pub fn encapsulate(public_key: &KyberPublicKey) -> (KyberSharedSecret, KyberCiphertext) {
157 let (ss, ct) = kyber1024::encapsulate(&public_key.0);
158 let mut shared_secret = [0u8; 32];
159 shared_secret.copy_from_slice(ss.as_bytes());
160 (KyberSharedSecret(shared_secret), KyberCiphertext(ct))
161 }
162
163 pub fn decapsulate(
172 ciphertext: &KyberCiphertext,
173 secret_key: &KyberSecretKey,
174 ) -> KyberSharedSecret {
175 let ss = kyber1024::decapsulate(&ciphertext.0, &secret_key.0);
176 let mut shared_secret = [0u8; 32];
177 shared_secret.copy_from_slice(ss.as_bytes());
178 KyberSharedSecret(shared_secret)
179 }
180
181 pub fn sizes() -> KyberSizes {
183 KyberSizes {
184 public_key: kyber1024::public_key_bytes(),
185 secret_key: kyber1024::secret_key_bytes(),
186 ciphertext: kyber1024::ciphertext_bytes(),
187 shared_secret: kyber1024::shared_secret_bytes(),
188 }
189 }
190}
191
192#[derive(Debug, Clone, Copy)]
194pub struct KyberSizes {
195 pub public_key: usize,
196 pub secret_key: usize,
197 pub ciphertext: usize,
198 pub shared_secret: usize,
199}
200
201#[cfg(test)]
202mod tests {
203 use super::*;
204
205 #[test]
206 fn test_keypair_generation() {
207 let keypair = KyberCrypto::generate_keypair();
208 assert_eq!(keypair.public_key.as_bytes().len(), KyberPublicKey::size());
209 assert_eq!(keypair.secret_key.as_bytes().len(), KyberSecretKey::size());
210 }
211
212 #[test]
213 fn test_encapsulate_decapsulate() {
214 let alice_keypair = KyberCrypto::generate_keypair();
216
217 let (bob_shared_secret, ciphertext) = KyberCrypto::encapsulate(&alice_keypair.public_key);
219
220 let alice_shared_secret = KyberCrypto::decapsulate(&ciphertext, &alice_keypair.secret_key);
222
223 assert_eq!(bob_shared_secret.as_bytes(), alice_shared_secret.as_bytes());
225 }
226
227 #[test]
228 fn test_serialization() {
229 let keypair = KyberCrypto::generate_keypair();
230
231 let pk_bytes = keypair.public_key.as_bytes();
233 let pk_restored = KyberPublicKey::from_bytes(pk_bytes).unwrap();
234 assert_eq!(pk_bytes, pk_restored.as_bytes());
235
236 let sk_bytes = keypair.secret_key.as_bytes();
238 let sk_restored = KyberSecretKey::from_bytes(sk_bytes).unwrap();
239 assert_eq!(sk_bytes, sk_restored.as_bytes());
240
241 let (ss1, ct) = KyberCrypto::encapsulate(&pk_restored);
243 let ss2 = KyberCrypto::decapsulate(&ct, &sk_restored);
244 assert_eq!(ss1.as_bytes(), ss2.as_bytes());
245 }
246
247 #[test]
248 fn test_sizes() {
249 let sizes = KyberCrypto::sizes();
250 assert_eq!(sizes.public_key, 1568); assert_eq!(sizes.secret_key, 3168); assert_eq!(sizes.ciphertext, 1568); assert_eq!(sizes.shared_secret, 32); }
255}