Skip to main content

apfsds_crypto/
keys.rs

1//! Ed25519, X25519, and ML-DSA-65 key management
2
3use 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
28/// Ed25519 key pair for signing
29pub struct Ed25519KeyPair {
30    signing_key: SigningKey,
31}
32
33impl Ed25519KeyPair {
34    /// Generate a new random key pair
35    pub fn generate() -> Self {
36        let signing_key = SigningKey::generate(&mut OsRng);
37        Self { signing_key }
38    }
39
40    /// Create from secret key bytes
41    pub fn from_secret(secret: &[u8; 32]) -> Self {
42        let signing_key = SigningKey::from_bytes(secret);
43        Self { signing_key }
44    }
45
46    /// Get the public key
47    pub fn public_key(&self) -> [u8; 32] {
48        self.signing_key.verifying_key().to_bytes()
49    }
50
51    /// Get the secret key
52    pub fn secret_key(&self) -> [u8; 32] {
53        self.signing_key.to_bytes()
54    }
55
56    /// Sign a message
57    pub fn sign(&self, message: &[u8]) -> [u8; 64] {
58        self.signing_key.sign(message).to_bytes()
59    }
60
61    /// Verify a signature (requires only the public key)
62    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
81/// ML-DSA-65 (Dilithium3) key pair for post-quantum signatures
82pub struct MlDsa65KeyPair {
83    keypair: ml_dsa::KeyPair<MlDsa65>,
84}
85
86impl MlDsa65KeyPair {
87    /// Generate a new random key pair
88    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    /// Create from secret key bytes (32-byte seed)
97    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    /// Get the public key bytes
110    pub fn public_key(&self) -> Vec<u8> {
111        self.keypair.verifying_key().encode().to_vec()
112    }
113
114    /// Get the secret key bytes (32-byte seed)
115    pub fn secret_key(&self) -> Vec<u8> {
116        self.keypair.to_seed().to_vec()
117    }
118
119    /// Sign a message (returns detached signature)
120    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    /// Verify a signature with public key
127    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        // Convert pk_bytes to EncodedVerifyingKey
135        let pk_array = pk_bytes
136            .try_into()
137            .map_err(|_| KeyError::InvalidKeyLength {
138                expected: 1952, // ML-DSA-65 public key size
139                actual: pk_bytes.len(),
140            })?;
141        let verifying_key = VerifyingKey::<MlDsa65>::decode(pk_array);
142
143        // Convert signature bytes to EncodedSignature and decode
144        let sig_array = signature
145            .try_into()
146            .map_err(|_| KeyError::InvalidKeyLength {
147                expected: 3309, // ML-DSA-65 signature size
148                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
159/// X25519 key pair for ECDH key exchange
160pub struct X25519KeyPair {
161    secret: StaticSecret,
162    public: X25519PublicKey,
163}
164
165impl X25519KeyPair {
166    /// Generate a new random key pair
167    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    /// Create from secret key bytes
174    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    /// Get the public key
181    pub fn public_key(&self) -> [u8; 32] {
182        self.public.to_bytes()
183    }
184
185    /// Perform ECDH to derive a shared secret
186    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
192/// ML-KEM-768 (Kyber) key pair for post-quantum key exchange
193pub struct MlKem768KeyPair {
194    secret_key: Vec<u8>,
195    public_key: Vec<u8>,
196}
197
198impl MlKem768KeyPair {
199    /// Generate a new random key pair
200    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    /// Get the public key bytes
213    pub fn public_key(&self) -> &[u8] {
214        &self.public_key
215    }
216
217    /// Get the secret key bytes
218    pub fn secret_key(&self) -> &[u8] {
219        &self.secret_key
220    }
221
222    /// Encapsulate: generate shared secret and ciphertext (for sender)
223    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        // Deserialize the public key
228        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        // Encapsulate to generate shared secret and ciphertext
237        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        // ct and shared_secret are Array types, convert directly to Vec
243        Ok((shared_secret.to_vec(), ct.to_vec()))
244    }
245
246    /// Decapsulate: recover shared secret from ciphertext (for receiver)
247    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        // Deserialize the secret key
254        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        // Deserialize the ciphertext (Ciphertext is parameterized by MlKem768, not MlKem768Params)
265        let ct: &Ciphertext<MlKem768> =
266            ciphertext
267                .try_into()
268                .map_err(|_| KeyError::InvalidKeyLength {
269                    expected: 1088,
270                    actual: ciphertext.len(),
271                })?;
272
273        // Decapsulate to recover shared secret
274        let shared_secret = decaps_key
275            .decapsulate(ct)
276            .map_err(|e| KeyError::KeyDeserializationFailed(format!("{:?}", e)))?;
277
278        // shared_secret is an Array type, convert directly to Vec
279        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}