1use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
4use ml_dsa::{KeyGen, MlDsa65};
5use rand::rngs::OsRng;
6use signature::{Signer as SigSigner, Verifier as SigVerifier};
7use thiserror::Error;
8use x25519_dalek::{PublicKey as X25519PublicKey, StaticSecret};
9
10#[derive(Error, Debug)]
11pub enum KeyError {
12 #[error("Invalid key length: expected {expected}, got {actual}")]
13 InvalidKeyLength { expected: usize, actual: usize },
14
15 #[error("Signature verification failed")]
16 SignatureVerificationFailed,
17
18 #[error("Invalid signature format")]
19 InvalidSignatureFormat,
20
21 #[error("Key deserialization failed: {0}")]
22 KeyDeserializationFailed(String),
23
24 #[error("Key serialization failed: {0}")]
25 KeySerializationFailed(String),
26}
27
28pub struct Ed25519KeyPair {
30 signing_key: SigningKey,
31}
32
33impl Ed25519KeyPair {
34 pub fn generate() -> Self {
36 let signing_key = SigningKey::generate(&mut OsRng);
37 Self { signing_key }
38 }
39
40 pub fn from_secret(secret: &[u8; 32]) -> Self {
42 let signing_key = SigningKey::from_bytes(secret);
43 Self { signing_key }
44 }
45
46 pub fn public_key(&self) -> [u8; 32] {
48 self.signing_key.verifying_key().to_bytes()
49 }
50
51 pub fn secret_key(&self) -> [u8; 32] {
53 self.signing_key.to_bytes()
54 }
55
56 pub fn sign(&self, message: &[u8]) -> [u8; 64] {
58 self.signing_key.sign(message).to_bytes()
59 }
60
61 pub fn verify_with_pk(
63 pk: &[u8; 32],
64 message: &[u8],
65 signature: &[u8; 64],
66 ) -> Result<(), KeyError> {
67 let verifying_key =
68 VerifyingKey::from_bytes(pk).map_err(|_| KeyError::InvalidKeyLength {
69 expected: 32,
70 actual: pk.len(),
71 })?;
72
73 let sig = Signature::from_bytes(signature);
74
75 verifying_key
76 .verify(message, &sig)
77 .map_err(|_| KeyError::SignatureVerificationFailed)
78 }
79}
80
81pub struct MlDsa65KeyPair {
83 keypair: ml_dsa::KeyPair<MlDsa65>,
84}
85
86impl MlDsa65KeyPair {
87 pub fn generate() -> Self {
89 use rand::RngCore;
90 let mut seed = [0u8; 32];
91 OsRng.fill_bytes(&mut seed);
92 let keypair = MlDsa65::from_seed(&seed.into());
93 Self { keypair }
94 }
95
96 pub fn from_secret(secret_bytes: &[u8]) -> Result<Self, KeyError> {
98 if secret_bytes.len() != 32 {
99 return Err(KeyError::InvalidKeyLength {
100 expected: 32,
101 actual: secret_bytes.len(),
102 });
103 }
104 let seed: [u8; 32] = secret_bytes.try_into().unwrap();
105 let keypair = MlDsa65::from_seed(&seed.into());
106 Ok(Self { keypair })
107 }
108
109 pub fn public_key(&self) -> Vec<u8> {
111 self.keypair.verifying_key().encode().to_vec()
112 }
113
114 pub fn secret_key(&self) -> Vec<u8> {
116 self.keypair.to_seed().to_vec()
117 }
118
119 pub fn sign(&self, message: &[u8]) -> Vec<u8> {
121 let sig = self.keypair.signing_key().sign(message);
122 let encoded = sig.encode();
123 <[u8]>::to_vec(encoded.as_ref())
124 }
125
126 pub fn verify_with_pk(
128 pk_bytes: &[u8],
129 message: &[u8],
130 signature: &[u8],
131 ) -> Result<(), KeyError> {
132 use ml_dsa::{Signature, VerifyingKey};
133
134 let pk_array = pk_bytes
136 .try_into()
137 .map_err(|_| KeyError::InvalidKeyLength {
138 expected: 1952, actual: pk_bytes.len(),
140 })?;
141 let verifying_key = VerifyingKey::<MlDsa65>::decode(pk_array);
142
143 let sig_array = signature
145 .try_into()
146 .map_err(|_| KeyError::InvalidKeyLength {
147 expected: 3309, actual: signature.len(),
149 })?;
150 let sig =
151 Signature::<MlDsa65>::decode(sig_array).ok_or(KeyError::InvalidSignatureFormat)?;
152
153 verifying_key
154 .verify(message, &sig)
155 .map_err(|_| KeyError::SignatureVerificationFailed)
156 }
157}
158
159pub struct X25519KeyPair {
161 secret: StaticSecret,
162 public: X25519PublicKey,
163}
164
165impl X25519KeyPair {
166 pub fn generate() -> Self {
168 let secret = StaticSecret::random_from_rng(OsRng);
169 let public = X25519PublicKey::from(&secret);
170 Self { secret, public }
171 }
172
173 pub fn from_secret(secret_bytes: &[u8; 32]) -> Self {
175 let secret = StaticSecret::from(*secret_bytes);
176 let public = X25519PublicKey::from(&secret);
177 Self { secret, public }
178 }
179
180 pub fn public_key(&self) -> [u8; 32] {
182 self.public.to_bytes()
183 }
184
185 pub fn diffie_hellman(&self, their_public: &[u8; 32]) -> [u8; 32] {
187 let their_pk = X25519PublicKey::from(*their_public);
188 self.secret.diffie_hellman(&their_pk).to_bytes()
189 }
190}
191
192pub struct MlKem768KeyPair {
194 secret_key: Vec<u8>,
195 public_key: Vec<u8>,
196}
197
198impl MlKem768KeyPair {
199 pub fn generate() -> Self {
201 use ml_kem::{EncodedSizeUser, KemCore, MlKem768};
202
203 let mut rng = OsRng;
204 let (decapsulation_key, encapsulation_key) = MlKem768::generate(&mut rng);
205
206 Self {
207 secret_key: decapsulation_key.as_bytes().to_vec(),
208 public_key: encapsulation_key.as_bytes().to_vec(),
209 }
210 }
211
212 pub fn public_key(&self) -> &[u8] {
214 &self.public_key
215 }
216
217 pub fn secret_key(&self) -> &[u8] {
219 &self.secret_key
220 }
221
222 pub fn encapsulate(their_public_key: &[u8]) -> Result<(Vec<u8>, Vec<u8>), KeyError> {
224 use ::kem::Encapsulate;
225 use ml_kem::{EncodedSizeUser, MlKem768Params, kem::EncapsulationKey};
226
227 let encoded_key = their_public_key
229 .try_into()
230 .map_err(|_| KeyError::InvalidKeyLength {
231 expected: 1184,
232 actual: their_public_key.len(),
233 })?;
234 let encaps_key = EncapsulationKey::<MlKem768Params>::from_bytes(encoded_key);
235
236 let mut rng = OsRng;
238 let (ct, shared_secret) = encaps_key
239 .encapsulate(&mut rng)
240 .map_err(|e| KeyError::KeySerializationFailed(format!("{:?}", e)))?;
241
242 Ok((shared_secret.to_vec(), ct.to_vec()))
244 }
245
246 pub fn decapsulate(&self, ciphertext: &[u8]) -> Result<Vec<u8>, KeyError> {
248 use ::kem::Decapsulate;
249 use ml_kem::{
250 Ciphertext, EncodedSizeUser, MlKem768, MlKem768Params, kem::DecapsulationKey,
251 };
252
253 let encoded_key =
255 self.secret_key
256 .as_slice()
257 .try_into()
258 .map_err(|_| KeyError::InvalidKeyLength {
259 expected: 2400,
260 actual: self.secret_key.len(),
261 })?;
262 let decaps_key = DecapsulationKey::<MlKem768Params>::from_bytes(encoded_key);
263
264 let ct: &Ciphertext<MlKem768> =
266 ciphertext
267 .try_into()
268 .map_err(|_| KeyError::InvalidKeyLength {
269 expected: 1088,
270 actual: ciphertext.len(),
271 })?;
272
273 let shared_secret = decaps_key
275 .decapsulate(ct)
276 .map_err(|e| KeyError::KeyDeserializationFailed(format!("{:?}", e)))?;
277
278 Ok(shared_secret.to_vec())
280 }
281}
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286
287 #[test]
288 fn test_ed25519_sign_verify() {
289 let keypair = Ed25519KeyPair::generate();
290 let message = b"Hello, APFSDS!";
291
292 let signature = keypair.sign(message);
293 let pk = keypair.public_key();
294
295 assert!(Ed25519KeyPair::verify_with_pk(&pk, message, &signature).is_ok());
296 }
297
298 #[test]
299 fn test_ed25519_invalid_signature() {
300 let keypair = Ed25519KeyPair::generate();
301 let message = b"Hello, APFSDS!";
302 let wrong_message = b"Wrong message";
303
304 let signature = keypair.sign(message);
305 let pk = keypair.public_key();
306
307 assert!(Ed25519KeyPair::verify_with_pk(&pk, wrong_message, &signature).is_err());
308 }
309
310 #[test]
311 fn test_x25519_key_exchange() {
312 let alice = X25519KeyPair::generate();
313 let bob = X25519KeyPair::generate();
314
315 let alice_shared = alice.diffie_hellman(&bob.public_key());
316 let bob_shared = bob.diffie_hellman(&alice.public_key());
317
318 assert_eq!(alice_shared, bob_shared);
319 }
320
321 #[test]
322 fn test_key_serialization() {
323 let keypair = Ed25519KeyPair::generate();
324 let secret = keypair.secret_key();
325 let restored = Ed25519KeyPair::from_secret(&secret);
326
327 assert_eq!(keypair.public_key(), restored.public_key());
328 }
329}