Skip to main content

lib_q_core/contexts/
kem.rs

1//! KEM context implementation for lib-Q Core
2//!
3//! This module provides the KEM context that handles key encapsulation
4//! mechanism operations with proper security validation.
5
6#[cfg(feature = "alloc")]
7use alloc::{
8    boxed::Box,
9    string::String,
10    vec::Vec,
11};
12
13use super::BaseContext;
14#[cfg(test)]
15use crate::api::KemOperations;
16use crate::api::{
17    Algorithm,
18    AlgorithmCategory,
19    CryptoProvider,
20};
21use crate::error::Result;
22use crate::traits::{
23    KemKeypair,
24    KemPublicKey,
25    KemSecretKey,
26};
27
28/// KEM context for key encapsulation operations
29#[cfg(feature = "alloc")]
30pub struct KemContext {
31    inner: BaseContext<Self>,
32}
33
34#[cfg(feature = "alloc")]
35impl KemContext {
36    /// Create a new KEM context with no provider
37    pub fn new() -> Self {
38        Self {
39            inner: BaseContext::new(),
40        }
41    }
42
43    /// Create a new KEM context with a provider
44    pub fn with_provider(provider: Box<dyn CryptoProvider>) -> Self {
45        Self {
46            inner: BaseContext::with_provider(provider),
47        }
48    }
49
50    /// Create a new KEM context with the default provider
51    #[cfg(feature = "alloc")]
52    pub fn with_default_provider() -> Self {
53        Self {
54            inner: match crate::providers::LibQCryptoProvider::new() {
55                Ok(provider) => BaseContext::with_provider(Box::new(provider)),
56                Err(_) => BaseContext::new(),
57            },
58        }
59    }
60
61    /// Set the cryptographic provider
62    pub fn set_provider(&mut self, provider: Box<dyn CryptoProvider>) {
63        self.inner.set_provider(provider);
64    }
65
66    /// Get the current provider
67    pub fn provider(&self) -> Option<&dyn CryptoProvider> {
68        self.inner.provider()
69    }
70
71    /// Generate a keypair for the specified algorithm
72    pub fn generate_keypair(
73        &mut self,
74        algorithm: Algorithm,
75        randomness: Option<&[u8]>,
76    ) -> Result<KemKeypair> {
77        self.inner.ensure_initialized()?;
78
79        // Validate algorithm category
80        if algorithm.category() != AlgorithmCategory::Kem {
81            return Err(crate::error::Error::InvalidAlgorithm {
82                algorithm: "Algorithm is not a KEM algorithm",
83            });
84        }
85
86        // Use provider if available
87        match self.inner.provider().and_then(|p| p.kem()) {
88            Some(kem_ops) => kem_ops.generate_keypair(algorithm, randomness),
89            None => Err(crate::error::Error::ProviderNotConfigured {
90                operation: String::from("KEM"),
91            }),
92        }
93    }
94
95    /// Encapsulate a key using the given public key
96    pub fn encapsulate(
97        &self,
98        algorithm: Algorithm,
99        public_key: &KemPublicKey,
100        randomness: Option<&[u8]>,
101    ) -> Result<(Vec<u8>, Vec<u8>)> {
102        if !self.inner.is_initialized() {
103            return Err(crate::error::Error::InvalidState {
104                operation: String::from("encapsulate"),
105                reason: String::from("Context not initialized"),
106            });
107        }
108
109        // Validate algorithm category
110        if algorithm.category() != AlgorithmCategory::Kem {
111            return Err(crate::error::Error::InvalidAlgorithm {
112                algorithm: "Algorithm is not a KEM algorithm",
113            });
114        }
115
116        // Use provider if available
117        match self.inner.provider().and_then(|p| p.kem()) {
118            Some(kem_ops) => kem_ops.encapsulate(algorithm, public_key, randomness),
119            None => Err(crate::error::Error::ProviderNotConfigured {
120                operation: String::from("KEM"),
121            }),
122        }
123    }
124
125    /// Decapsulate a key using the given secret key and ciphertext
126    pub fn decapsulate(
127        &self,
128        algorithm: Algorithm,
129        secret_key: &KemSecretKey,
130        ciphertext: &[u8],
131    ) -> Result<Vec<u8>> {
132        if !self.inner.is_initialized() {
133            return Err(crate::error::Error::InvalidState {
134                operation: String::from("decapsulate"),
135                reason: String::from("Context not initialized"),
136            });
137        }
138
139        // Validate algorithm category
140        if algorithm.category() != AlgorithmCategory::Kem {
141            return Err(crate::error::Error::InvalidAlgorithm {
142                algorithm: "Algorithm is not a KEM algorithm",
143            });
144        }
145
146        // Use provider if available
147        match self.inner.provider().and_then(|p| p.kem()) {
148            Some(kem_ops) => kem_ops.decapsulate(algorithm, secret_key, ciphertext),
149            None => Err(crate::error::Error::ProviderNotConfigured {
150                operation: String::from("KEM"),
151            }),
152        }
153    }
154
155    /// Check if the context is initialized
156    pub fn is_initialized(&self) -> bool {
157        self.inner.is_initialized()
158    }
159}
160
161#[cfg(feature = "alloc")]
162impl Default for KemContext {
163    fn default() -> Self {
164        Self::new()
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171    use crate::api::CryptoProvider;
172
173    // Mock provider for testing
174    struct MockKemProvider;
175
176    impl CryptoProvider for MockKemProvider {
177        fn kem(&self) -> Option<&dyn KemOperations> {
178            Some(self)
179        }
180        fn signature(&self) -> Option<&dyn crate::api::SignatureOperations> {
181            None
182        }
183        fn hash(&self) -> Option<&dyn crate::api::HashOperations> {
184            None
185        }
186        fn aead(&self) -> Option<&dyn crate::api::AeadOperations> {
187            None
188        }
189    }
190
191    impl KemOperations for MockKemProvider {
192        fn generate_keypair(
193            &self,
194            _algorithm: Algorithm,
195            _randomness: Option<&[u8]>,
196        ) -> Result<KemKeypair> {
197            Err(crate::error::Error::NotImplemented {
198                feature: "Mock KEM operations not implemented".to_string(),
199            })
200        }
201
202        fn encapsulate(
203            &self,
204            _algorithm: Algorithm,
205            _public_key: &KemPublicKey,
206            _randomness: Option<&[u8]>,
207        ) -> Result<(Vec<u8>, Vec<u8>)> {
208            Err(crate::error::Error::NotImplemented {
209                feature: "Mock KEM operations not implemented".to_string(),
210            })
211        }
212
213        fn decapsulate(
214            &self,
215            _algorithm: Algorithm,
216            _secret_key: &KemSecretKey,
217            _ciphertext: &[u8],
218        ) -> Result<Vec<u8>> {
219            Err(crate::error::Error::NotImplemented {
220                feature: "Mock KEM operations not implemented".to_string(),
221            })
222        }
223
224        fn derive_public_key(
225            &self,
226            _algorithm: Algorithm,
227            _secret_key: &KemSecretKey,
228        ) -> Result<KemPublicKey> {
229            Err(crate::error::Error::NotImplemented {
230                feature: "Mock KEM operations not implemented".to_string(),
231            })
232        }
233    }
234
235    #[test]
236    fn test_kem_context_creation() {
237        let context = KemContext::new();
238        assert!(!context.is_initialized());
239        assert!(context.provider().is_none());
240    }
241
242    #[test]
243    fn test_kem_context_with_provider() {
244        let provider = Box::new(MockKemProvider);
245        let context = KemContext::with_provider(provider);
246        assert!(!context.is_initialized());
247        assert!(context.provider().is_some());
248    }
249
250    #[test]
251    fn test_kem_context_provider_management() {
252        let mut context = KemContext::new();
253        assert!(context.provider().is_none());
254
255        let provider = Box::new(MockKemProvider);
256        context.set_provider(provider);
257        assert!(context.provider().is_some());
258    }
259
260    #[test]
261    fn test_kem_context_initialization() {
262        let mut context = KemContext::new();
263        assert!(!context.is_initialized());
264
265        // Should initialize automatically on first operation
266        let result = context.generate_keypair(Algorithm::MlKem512, None);
267        assert!(result.is_err()); // Will fail due to no provider, but context should be initialized
268        assert!(context.is_initialized());
269    }
270
271    #[test]
272    fn test_kem_context_algorithm_validation() {
273        let mut context = KemContext::new();
274
275        // Try to use a non-KEM algorithm
276        let result = context.generate_keypair(Algorithm::Sha3_256, None);
277        assert!(result.is_err());
278        if let Err(crate::error::Error::InvalidAlgorithm { .. }) = result {
279            // Expected error
280        } else {
281            panic!("Expected InvalidAlgorithm error");
282        }
283    }
284}