Skip to main content

ri/protocol/
dilithium.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//! # Dilithium Signature
21//!
22//! This module implements the Dilithium digital signature algorithm
23//! using liboqs. Dilithium is EUF-CMA secure and is based on the hardness
24//! of the Module-LWE (Learning With Errors over Modules) problem.
25//!
26//! ## Security Level
27//!
28//! - **Dilithium2**: NIST Level 2 ≈ AES-128
29//! - **Dilithium3**: NIST Level 3 ≈ AES-192
30//! - **Dilithium5**: NIST Level 5 ≈ AES-256
31//!
32//! ## Usage
33//!
34//! ```rust,ignore
35//! use ri::protocol::dilithium::DilithiumSigner;
36//!
37//! let signer = DilithiumSigner::new();
38//! let (public_key, secret_key) = signer.keygen()?;
39//! let message = b"Hello, Post-Quantum World!";
40//! let signature = signer.sign(&secret_key, message)?;
41//! assert!(signer.verify(&public_key, message, &signature)?);
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 DilithiumPublicKey(pub Vec<u8>);
50
51#[derive(Debug, Clone)]
52#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
53pub struct DilithiumSecretKey(pub Vec<u8>);
54
55#[derive(Debug, Clone)]
56#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
57pub struct DilithiumSignature(pub Vec<u8>);
58
59#[derive(Debug, Clone)]
60#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
61pub struct DilithiumSigner {
62    algorithm: Arc<std::sync::RwLock<DilithiumAlgorithm>>,
63}
64
65#[derive(Debug, Clone, Copy)]
66enum DilithiumAlgorithm {
67    Dilithium2,
68    Dilithium3,
69    Dilithium5,
70}
71
72impl DilithiumSigner {
73    pub fn new() -> Self {
74        Self {
75            algorithm: Arc::new(std::sync::RwLock::new(DilithiumAlgorithm::Dilithium2)),
76        }
77    }
78
79    pub fn with_algorithm(algorithm: super::RiPostQuantumAlgorithm) -> Self {
80        let algo = match algorithm {
81            super::RiPostQuantumAlgorithm::Dilithium2 => DilithiumAlgorithm::Dilithium2,
82            super::RiPostQuantumAlgorithm::Dilithium3 => DilithiumAlgorithm::Dilithium3,
83            super::RiPostQuantumAlgorithm::Dilithium5 => DilithiumAlgorithm::Dilithium5,
84            _ => DilithiumAlgorithm::Dilithium2,
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::sig::Sig;
94
95        let algo = *self.algorithm.read().map_err(|e| 
96            RiError::InvalidState(format!("Lock error: {}", e))
97        )?;
98        let sig = match algo {
99            DilithiumAlgorithm::Dilithium2 => Sig::new(oqs::sig::Algorithm::Dilithium2),
100            DilithiumAlgorithm::Dilithium3 => Sig::new(oqs::sig::Algorithm::Dilithium3),
101            DilithiumAlgorithm::Dilithium5 => Sig::new(oqs::sig::Algorithm::Dilithium5),
102        }.map_err(|e| RiError::Other(format!("Failed to initialize Dilithium: {:?}", e)))?;
103
104        let (pk, sk) = sig.keypair()
105            .map_err(|e| RiError::Other(format!("Dilithium 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 sign(&self, secret_key: &[u8], message: &[u8]) -> RiResult<Vec<u8>> {
119        use oqs::sig::Sig;
120
121        let algo = *self.algorithm.read().map_err(|e| 
122            RiError::InvalidState(format!("Lock error: {}", e))
123        )?;
124        let sig = match algo {
125            DilithiumAlgorithm::Dilithium2 => Sig::new(oqs::sig::Algorithm::Dilithium2),
126            DilithiumAlgorithm::Dilithium3 => Sig::new(oqs::sig::Algorithm::Dilithium3),
127            DilithiumAlgorithm::Dilithium5 => Sig::new(oqs::sig::Algorithm::Dilithium5),
128        }.map_err(|e| RiError::Other(format!("Failed to initialize Dilithium: {:?}", e)))?;
129
130        let sk = sig.secret_key_from_bytes(secret_key)
131            .ok_or_else(|| RiError::Other("Invalid secret key".to_string()))?;
132        let signature = sig.sign(message, &sk)
133            .map_err(|e| RiError::Other(format!("Dilithium sign failed: {:?}", e)))?;
134        Ok(signature.into_vec())
135    }
136
137    #[cfg(not(feature = "protocol"))]
138    pub fn sign(&self, _secret_key: &[u8], _message: &[u8]) -> RiResult<Vec<u8>> {
139        Err(RiError::Other(
140            "Post-quantum cryptography requires the 'protocol' feature. \
141             Enable with: cargo build --features protocol".to_string()
142        ))
143    }
144
145    #[cfg(feature = "protocol")]
146    pub fn verify(&self, public_key: &[u8], message: &[u8], signature: &[u8]) -> RiResult<bool> {
147        use oqs::sig::Sig;
148
149        let algo = *self.algorithm.read().map_err(|e| 
150            RiError::InvalidState(format!("Lock error: {}", e))
151        )?;
152        let sig = match algo {
153            DilithiumAlgorithm::Dilithium2 => Sig::new(oqs::sig::Algorithm::Dilithium2),
154            DilithiumAlgorithm::Dilithium3 => Sig::new(oqs::sig::Algorithm::Dilithium3),
155            DilithiumAlgorithm::Dilithium5 => Sig::new(oqs::sig::Algorithm::Dilithium5),
156        }.map_err(|e| RiError::Other(format!("Failed to initialize Dilithium: {:?}", e)))?;
157
158        let pk = sig.public_key_from_bytes(public_key)
159            .ok_or_else(|| RiError::Other("Invalid public key".to_string()))?;
160        let sig_bytes = sig.signature_from_bytes(signature)
161            .ok_or_else(|| RiError::Other("Invalid signature".to_string()))?;
162        let result = sig.verify(message, &sig_bytes, &pk);
163        Ok(result.is_ok())
164    }
165
166    #[cfg(not(feature = "protocol"))]
167    pub fn verify(&self, _public_key: &[u8], _message: &[u8], _signature: &[u8]) -> RiResult<bool> {
168        Err(RiError::Other(
169            "Post-quantum cryptography requires the 'protocol' feature. \
170             Enable with: cargo build --features protocol".to_string()
171        ))
172    }
173}
174
175impl Default for DilithiumSigner {
176    fn default() -> Self {
177        Self::new()
178    }
179}
180
181#[cfg(feature = "pyo3")]
182#[pyo3::prelude::pymethods]
183impl DilithiumSigner {
184    #[new]
185    pub fn new_py() -> Self {
186        Self::new()
187    }
188
189    pub fn keygen_py(&self) -> Option<(Vec<u8>, Vec<u8>)> {
190        self.keygen().ok()
191    }
192
193    pub fn sign_py(&self, secret_key: &[u8], message: &[u8]) -> Option<Vec<u8>> {
194        self.sign(secret_key, message).ok()
195    }
196
197    pub fn verify_py(&self, public_key: &[u8], message: &[u8], signature: &[u8]) -> bool {
198        self.verify(public_key, message, signature).unwrap_or(false)
199    }
200}