Skip to main content

ri/protocol/
kyber.rs

1//! Copyright © 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18#![allow(non_snake_case)]
19
20//! # Kyber KEM
21//!
22//! This module implements the Kyber Key Encapsulation Mechanism (KEM)
23//! using liboqs. Kyber is IND-CCA2 secure and is based on the hardness
24//! of the Module-LWE (Learning With Errors over Modules) problem.
25//!
26//! ## Security Level
27//!
28//! - **Kyber512**: NIST Level 1 ≈ AES-128
29//! - **Kyber768**: NIST Level 3 ≈ AES-192
30//! - **Kyber1024**: NIST Level 5 ≈ AES-256
31//!
32//! ## Usage
33//!
34//! ```rust,ignore
35//! use ri::protocol::kyber::KyberKEM;
36//!
37//! let kem = KyberKEM::new();
38//! let (public_key, secret_key) = kem.keygen()?;
39//! let (ciphertext, shared_secret_1) = kem.encapsulate(&public_key)?;
40//! let shared_secret_2 = kem.decapsulate(&ciphertext, &secret_key)?;
41//! assert_eq!(shared_secret_1, shared_secret_2);
42//! ```
43
44use std::sync::Arc;
45use crate::core::{RiResult, RiError};
46
47#[derive(Debug, Clone)]
48#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
49pub struct KyberPublicKey(pub Vec<u8>);
50
51#[derive(Debug, Clone)]
52#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
53pub struct KyberSecretKey(pub Vec<u8>);
54
55#[derive(Debug, Clone)]
56#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
57pub struct KyberCiphertext(pub Vec<u8>);
58
59#[derive(Debug, Clone)]
60#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
61pub struct KyberKEM {
62    algorithm: Arc<std::sync::RwLock<KyberAlgorithm>>,
63}
64
65#[derive(Debug, Clone, Copy)]
66enum KyberAlgorithm {
67    Kyber512,
68    Kyber768,
69    Kyber1024,
70}
71
72impl KyberKEM {
73    pub fn new() -> Self {
74        Self {
75            algorithm: Arc::new(std::sync::RwLock::new(KyberAlgorithm::Kyber512)),
76        }
77    }
78
79    pub fn with_algorithm(algorithm: super::RiPostQuantumAlgorithm) -> Self {
80        let algo = match algorithm {
81            super::RiPostQuantumAlgorithm::Kyber512 => KyberAlgorithm::Kyber512,
82            super::RiPostQuantumAlgorithm::Kyber768 => KyberAlgorithm::Kyber768,
83            super::RiPostQuantumAlgorithm::Kyber1024 => KyberAlgorithm::Kyber1024,
84            _ => KyberAlgorithm::Kyber512,
85        };
86        Self {
87            algorithm: Arc::new(std::sync::RwLock::new(algo)),
88        }
89    }
90
91    #[cfg(feature = "protocol")]
92    pub fn keygen(&self) -> RiResult<(Vec<u8>, Vec<u8>)> {
93        use oqs::kem::Kem;
94
95        let algo = *self.algorithm.read().map_err(|e| 
96            RiError::InvalidState(format!("Lock error: {}", e))
97        )?;
98        let kem = match algo {
99            KyberAlgorithm::Kyber512 => Kem::new(oqs::kem::Algorithm::Kyber512),
100            KyberAlgorithm::Kyber768 => Kem::new(oqs::kem::Algorithm::Kyber768),
101            KyberAlgorithm::Kyber1024 => Kem::new(oqs::kem::Algorithm::Kyber1024),
102        }.map_err(|e| RiError::Other(format!("Failed to initialize Kyber: {:?}", e)))?;
103
104        let (pk, sk) = kem.keypair()
105            .map_err(|e| RiError::Other(format!("Kyber keygen failed: {:?}", e)))?;
106        Ok((pk.into_vec(), sk.into_vec()))
107    }
108
109    #[cfg(not(feature = "protocol"))]
110    pub fn keygen(&self) -> RiResult<(Vec<u8>, Vec<u8>)> {
111        Err(RiError::Other(
112            "Post-quantum cryptography requires the 'protocol' feature. \
113             Enable with: cargo build --features protocol".to_string()
114        ))
115    }
116
117    #[cfg(feature = "protocol")]
118    pub fn encapsulate(&self, public_key: &[u8]) -> RiResult<super::KEMResult> {
119        use oqs::kem::Kem;
120
121        let algo = *self.algorithm.read().map_err(|e| 
122            RiError::InvalidState(format!("Lock error: {}", e))
123        )?;
124        let kem = match algo {
125            KyberAlgorithm::Kyber512 => Kem::new(oqs::kem::Algorithm::Kyber512),
126            KyberAlgorithm::Kyber768 => Kem::new(oqs::kem::Algorithm::Kyber768),
127            KyberAlgorithm::Kyber1024 => Kem::new(oqs::kem::Algorithm::Kyber1024),
128        }.map_err(|e| RiError::Other(format!("Failed to initialize Kyber: {:?}", e)))?;
129
130        let pk = kem.public_key_from_bytes(public_key)
131            .ok_or_else(|| RiError::Other("Invalid public key".to_string()))?;
132        let (ct, ss) = kem.encapsulate(&pk)
133            .map_err(|e| RiError::Other(format!("Kyber encapsulate failed: {:?}", e)))?;
134        Ok(super::KEMResult {
135            ciphertext: ct.into_vec(),
136            shared_secret: ss.into_vec(),
137        })
138    }
139
140    #[cfg(not(feature = "protocol"))]
141    pub fn encapsulate(&self, _public_key: &[u8]) -> RiResult<super::KEMResult> {
142        Err(RiError::Other(
143            "Post-quantum cryptography requires the 'protocol' feature. \
144             Enable with: cargo build --features protocol".to_string()
145        ))
146    }
147
148    #[cfg(feature = "protocol")]
149    pub fn decapsulate(&self, ciphertext: &[u8], secret_key: &[u8]) -> RiResult<Vec<u8>> {
150        use oqs::kem::Kem;
151
152        let algo = *self.algorithm.read().map_err(|e| 
153            RiError::InvalidState(format!("Lock error: {}", e))
154        )?;
155        let kem = match algo {
156            KyberAlgorithm::Kyber512 => Kem::new(oqs::kem::Algorithm::Kyber512),
157            KyberAlgorithm::Kyber768 => Kem::new(oqs::kem::Algorithm::Kyber768),
158            KyberAlgorithm::Kyber1024 => Kem::new(oqs::kem::Algorithm::Kyber1024),
159        }.map_err(|e| RiError::Other(format!("Failed to initialize Kyber: {:?}", e)))?;
160
161        let ct = kem.ciphertext_from_bytes(ciphertext)
162            .ok_or_else(|| RiError::Other("Invalid ciphertext".to_string()))?;
163        let sk = kem.secret_key_from_bytes(secret_key)
164            .ok_or_else(|| RiError::Other("Invalid secret key".to_string()))?;
165        let ss = kem.decapsulate(&sk, &ct)
166            .map_err(|e| RiError::Other(format!("Kyber decapsulate failed: {:?}", e)))?;
167        Ok(ss.into_vec())
168    }
169
170    #[cfg(not(feature = "protocol"))]
171    pub fn decapsulate(&self, _ciphertext: &[u8], _secret_key: &[u8]) -> RiResult<Vec<u8>> {
172        Err(RiError::Other(
173            "Post-quantum cryptography requires the 'protocol' feature. \
174             Enable with: cargo build --features protocol".to_string()
175        ))
176    }
177}
178
179impl Default for KyberKEM {
180    fn default() -> Self {
181        Self::new()
182    }
183}
184
185#[cfg(feature = "pyo3")]
186#[pyo3::prelude::pymethods]
187impl KyberKEM {
188    #[new]
189    pub fn new_py() -> Self {
190        Self::new()
191    }
192
193    pub fn keygen_py(&self) -> Option<(Vec<u8>, Vec<u8>)> {
194        self.keygen().ok()
195    }
196
197    pub fn encapsulate_py(&self, public_key: &[u8]) -> Option<(Vec<u8>, Vec<u8>)> {
198        self.encapsulate(public_key).ok().map(|r| (r.ciphertext, r.shared_secret))
199    }
200
201    pub fn decapsulate_py(&self, ciphertext: &[u8], secret_key: &[u8]) -> Option<Vec<u8>> {
202        self.decapsulate(ciphertext, secret_key).ok()
203    }
204}