Skip to main content

rns_crypto/
identity.rs

1use alloc::vec::Vec;
2use core::fmt;
3
4use crate::ed25519::{Ed25519PrivateKey, Ed25519PublicKey};
5use crate::hkdf;
6use crate::sha256;
7use crate::token::{Token, TokenError};
8use crate::x25519::{X25519PrivateKey, X25519PublicKey};
9use crate::Rng;
10
11pub const KEYSIZE: usize = 512; // bits
12pub const DERIVED_KEY_LENGTH: usize = 64; // bytes
13pub const TRUNCATED_HASHLENGTH: usize = 128; // bits (16 bytes)
14
15#[derive(Debug)]
16pub enum CryptoError {
17    NoPrivateKey,
18    NoPublicKey,
19    TokenError(TokenError),
20    HkdfError(hkdf::HkdfError),
21    InvalidCiphertext,
22}
23
24impl fmt::Display for CryptoError {
25    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
26        match self {
27            CryptoError::NoPrivateKey => write!(f, "No private key"),
28            CryptoError::NoPublicKey => write!(f, "No public key"),
29            CryptoError::TokenError(e) => write!(f, "Token error: {}", e),
30            CryptoError::HkdfError(e) => write!(f, "HKDF error: {}", e),
31            CryptoError::InvalidCiphertext => write!(f, "Invalid ciphertext"),
32        }
33    }
34}
35
36pub struct Identity {
37    prv: Option<X25519PrivateKey>,
38    sig_prv: Option<Ed25519PrivateKey>,
39    pub_key: Option<X25519PublicKey>,
40    sig_pub: Option<Ed25519PublicKey>,
41    hash: [u8; 16],
42}
43
44impl Identity {
45    pub fn new(rng: &mut dyn Rng) -> Self {
46        let prv = X25519PrivateKey::generate(rng);
47        let sig_prv = Ed25519PrivateKey::generate(rng);
48
49        let pub_key = prv.public_key();
50        let sig_pub = sig_prv.public_key();
51
52        let mut pub_bytes = [0u8; 64];
53        pub_bytes[..32].copy_from_slice(&pub_key.public_bytes());
54        pub_bytes[32..].copy_from_slice(&sig_pub.public_bytes());
55
56        let hash = truncated_hash(&pub_bytes);
57
58        Identity {
59            prv: Some(prv),
60            sig_prv: Some(sig_prv),
61            pub_key: Some(pub_key),
62            sig_pub: Some(sig_pub),
63            hash,
64        }
65    }
66
67    pub fn from_private_key(prv_bytes: &[u8; 64]) -> Self {
68        let x_prv_bytes: [u8; 32] = prv_bytes[..32].try_into().unwrap();
69        let ed_seed: [u8; 32] = prv_bytes[32..].try_into().unwrap();
70
71        let prv = X25519PrivateKey::from_bytes(&x_prv_bytes);
72        let sig_prv = Ed25519PrivateKey::from_bytes(&ed_seed);
73
74        let pub_key = prv.public_key();
75        let sig_pub = sig_prv.public_key();
76
77        let mut pub_bytes = [0u8; 64];
78        pub_bytes[..32].copy_from_slice(&pub_key.public_bytes());
79        pub_bytes[32..].copy_from_slice(&sig_pub.public_bytes());
80
81        let hash = truncated_hash(&pub_bytes);
82
83        Identity {
84            prv: Some(prv),
85            sig_prv: Some(sig_prv),
86            pub_key: Some(pub_key),
87            sig_pub: Some(sig_pub),
88            hash,
89        }
90    }
91
92    pub fn from_public_key(pub_bytes: &[u8; 64]) -> Self {
93        let x_pub_bytes: [u8; 32] = pub_bytes[..32].try_into().unwrap();
94        let ed_pub_bytes: [u8; 32] = pub_bytes[32..].try_into().unwrap();
95
96        let pub_key = X25519PublicKey::from_bytes(&x_pub_bytes);
97        let sig_pub = Ed25519PublicKey::from_bytes(&ed_pub_bytes);
98
99        let hash = truncated_hash(pub_bytes);
100
101        Identity {
102            prv: None,
103            sig_prv: None,
104            pub_key: Some(pub_key),
105            sig_pub: Some(sig_pub),
106            hash,
107        }
108    }
109
110    pub fn get_private_key(&self) -> Option<[u8; 64]> {
111        match (&self.prv, &self.sig_prv) {
112            (Some(prv), Some(sig_prv)) => {
113                let mut result = [0u8; 64];
114                result[..32].copy_from_slice(&prv.private_bytes());
115                result[32..].copy_from_slice(&sig_prv.private_bytes());
116                Some(result)
117            }
118            _ => None,
119        }
120    }
121
122    pub fn get_public_key(&self) -> Option<[u8; 64]> {
123        match (&self.pub_key, &self.sig_pub) {
124            (Some(pub_key), Some(sig_pub)) => {
125                let mut result = [0u8; 64];
126                result[..32].copy_from_slice(&pub_key.public_bytes());
127                result[32..].copy_from_slice(&sig_pub.public_bytes());
128                Some(result)
129            }
130            _ => None,
131        }
132    }
133
134    pub fn hash(&self) -> &[u8; 16] {
135        &self.hash
136    }
137
138    pub fn encrypt(&self, plaintext: &[u8], rng: &mut dyn Rng) -> Result<Vec<u8>, CryptoError> {
139        let pub_key = self.pub_key.as_ref().ok_or(CryptoError::NoPublicKey)?;
140
141        let ephemeral = X25519PrivateKey::generate(rng);
142        let ephemeral_pub_bytes = ephemeral.public_key().public_bytes();
143        let shared_key = ephemeral.exchange(pub_key);
144
145        let derived_key = hkdf::hkdf(DERIVED_KEY_LENGTH, &shared_key, Some(&self.hash), None)
146            .map_err(CryptoError::HkdfError)?;
147
148        let token = Token::new(&derived_key).map_err(CryptoError::TokenError)?;
149        let ciphertext = token.encrypt(plaintext, rng);
150
151        let mut result = Vec::with_capacity(32 + ciphertext.len());
152        result.extend_from_slice(&ephemeral_pub_bytes);
153        result.extend_from_slice(&ciphertext);
154        Ok(result)
155    }
156
157    /// Encrypt with a specific ephemeral key and IV for deterministic testing
158    pub fn encrypt_deterministic(
159        &self,
160        plaintext: &[u8],
161        ephemeral_prv: &[u8; 32],
162        iv: &[u8; 16],
163    ) -> Result<Vec<u8>, CryptoError> {
164        let pub_key = self.pub_key.as_ref().ok_or(CryptoError::NoPublicKey)?;
165
166        let ephemeral = X25519PrivateKey::from_bytes(ephemeral_prv);
167        let ephemeral_pub_bytes = ephemeral.public_key().public_bytes();
168        let shared_key = ephemeral.exchange(pub_key);
169
170        let derived_key = hkdf::hkdf(DERIVED_KEY_LENGTH, &shared_key, Some(&self.hash), None)
171            .map_err(CryptoError::HkdfError)?;
172
173        let token = Token::new(&derived_key).map_err(CryptoError::TokenError)?;
174        let ciphertext = token.encrypt_with_iv(plaintext, iv);
175
176        let mut result = Vec::with_capacity(32 + ciphertext.len());
177        result.extend_from_slice(&ephemeral_pub_bytes);
178        result.extend_from_slice(&ciphertext);
179        Ok(result)
180    }
181
182    pub fn decrypt(&self, ciphertext_token: &[u8]) -> Result<Vec<u8>, CryptoError> {
183        let prv = self.prv.as_ref().ok_or(CryptoError::NoPrivateKey)?;
184
185        if ciphertext_token.len() <= KEYSIZE / 8 / 2 {
186            return Err(CryptoError::InvalidCiphertext);
187        }
188
189        let peer_pub_bytes: [u8; 32] = ciphertext_token[..32].try_into().unwrap();
190        let peer_pub = X25519PublicKey::from_bytes(&peer_pub_bytes);
191        let ciphertext = &ciphertext_token[32..];
192
193        let shared_key = prv.exchange(&peer_pub);
194
195        let derived_key = hkdf::hkdf(DERIVED_KEY_LENGTH, &shared_key, Some(&self.hash), None)
196            .map_err(CryptoError::HkdfError)?;
197
198        let token = Token::new(&derived_key).map_err(CryptoError::TokenError)?;
199        token.decrypt(ciphertext).map_err(CryptoError::TokenError)
200    }
201
202    pub fn sign(&self, message: &[u8]) -> Result<[u8; 64], CryptoError> {
203        let sig_prv = self.sig_prv.as_ref().ok_or(CryptoError::NoPrivateKey)?;
204        Ok(sig_prv.sign(message))
205    }
206
207    pub fn verify(&self, signature: &[u8; 64], message: &[u8]) -> bool {
208        match &self.sig_pub {
209            Some(sig_pub) => sig_pub.verify(signature, message),
210            None => false,
211        }
212    }
213}
214
215fn truncated_hash(data: &[u8]) -> [u8; 16] {
216    let full = sha256::sha256(data);
217    let mut result = [0u8; 16];
218    result.copy_from_slice(&full[..16]);
219    result
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use crate::FixedRng;
226
227    #[test]
228    fn test_identity_key_roundtrip() {
229        let mut rng = FixedRng::new(&(0..64).collect::<Vec<u8>>());
230        let id = Identity::new(&mut rng);
231        let prv_bytes = id.get_private_key().unwrap();
232        let id2 = Identity::from_private_key(&prv_bytes);
233        assert_eq!(id.get_public_key().unwrap(), id2.get_public_key().unwrap());
234    }
235
236    #[test]
237    fn test_identity_hash() {
238        let mut rng = FixedRng::new(&(0..64).collect::<Vec<u8>>());
239        let id = Identity::new(&mut rng);
240        let pub_key = id.get_public_key().unwrap();
241        let expected_hash = truncated_hash(&pub_key);
242        assert_eq!(*id.hash(), expected_hash);
243    }
244
245    #[test]
246    fn test_identity_encrypt_decrypt_roundtrip() {
247        let mut rng = FixedRng::new(&(0..128).collect::<Vec<u8>>());
248        let id = Identity::new(&mut rng);
249        let plaintext = b"Hello, Reticulum! This is a test of the encrypt/decrypt pipeline.";
250        let mut rng2 = FixedRng::new(&(128..255).collect::<Vec<u8>>());
251        let ciphertext = id.encrypt(plaintext, &mut rng2).unwrap();
252        let decrypted = id.decrypt(&ciphertext).unwrap();
253        assert_eq!(decrypted, plaintext);
254    }
255
256    #[test]
257    fn test_identity_sign_verify() {
258        let mut rng = FixedRng::new(&(0..64).collect::<Vec<u8>>());
259        let id = Identity::new(&mut rng);
260        let msg = b"Sign this message";
261        let sig = id.sign(msg).unwrap();
262        assert!(id.verify(&sig, msg));
263        assert!(!id.verify(&sig, b"Wrong message"));
264    }
265}