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(
146            DERIVED_KEY_LENGTH,
147            &shared_key,
148            Some(&self.hash),
149            None,
150        )
151        .map_err(CryptoError::HkdfError)?;
152
153        let token = Token::new(&derived_key).map_err(CryptoError::TokenError)?;
154        let ciphertext = token.encrypt(plaintext, rng);
155
156        let mut result = Vec::with_capacity(32 + ciphertext.len());
157        result.extend_from_slice(&ephemeral_pub_bytes);
158        result.extend_from_slice(&ciphertext);
159        Ok(result)
160    }
161
162    /// Encrypt with a specific ephemeral key and IV for deterministic testing
163    pub fn encrypt_deterministic(
164        &self,
165        plaintext: &[u8],
166        ephemeral_prv: &[u8; 32],
167        iv: &[u8; 16],
168    ) -> Result<Vec<u8>, CryptoError> {
169        let pub_key = self.pub_key.as_ref().ok_or(CryptoError::NoPublicKey)?;
170
171        let ephemeral = X25519PrivateKey::from_bytes(ephemeral_prv);
172        let ephemeral_pub_bytes = ephemeral.public_key().public_bytes();
173        let shared_key = ephemeral.exchange(pub_key);
174
175        let derived_key = hkdf::hkdf(
176            DERIVED_KEY_LENGTH,
177            &shared_key,
178            Some(&self.hash),
179            None,
180        )
181        .map_err(CryptoError::HkdfError)?;
182
183        let token = Token::new(&derived_key).map_err(CryptoError::TokenError)?;
184        let ciphertext = token.encrypt_with_iv(plaintext, iv);
185
186        let mut result = Vec::with_capacity(32 + ciphertext.len());
187        result.extend_from_slice(&ephemeral_pub_bytes);
188        result.extend_from_slice(&ciphertext);
189        Ok(result)
190    }
191
192    pub fn decrypt(&self, ciphertext_token: &[u8]) -> Result<Vec<u8>, CryptoError> {
193        let prv = self.prv.as_ref().ok_or(CryptoError::NoPrivateKey)?;
194
195        if ciphertext_token.len() <= KEYSIZE / 8 / 2 {
196            return Err(CryptoError::InvalidCiphertext);
197        }
198
199        let peer_pub_bytes: [u8; 32] = ciphertext_token[..32].try_into().unwrap();
200        let peer_pub = X25519PublicKey::from_bytes(&peer_pub_bytes);
201        let ciphertext = &ciphertext_token[32..];
202
203        let shared_key = prv.exchange(&peer_pub);
204
205        let derived_key = hkdf::hkdf(
206            DERIVED_KEY_LENGTH,
207            &shared_key,
208            Some(&self.hash),
209            None,
210        )
211        .map_err(CryptoError::HkdfError)?;
212
213        let token = Token::new(&derived_key).map_err(CryptoError::TokenError)?;
214        token.decrypt(ciphertext).map_err(CryptoError::TokenError)
215    }
216
217    pub fn sign(&self, message: &[u8]) -> Result<[u8; 64], CryptoError> {
218        let sig_prv = self.sig_prv.as_ref().ok_or(CryptoError::NoPrivateKey)?;
219        Ok(sig_prv.sign(message))
220    }
221
222    pub fn verify(&self, signature: &[u8; 64], message: &[u8]) -> bool {
223        match &self.sig_pub {
224            Some(sig_pub) => sig_pub.verify(signature, message),
225            None => false,
226        }
227    }
228}
229
230fn truncated_hash(data: &[u8]) -> [u8; 16] {
231    let full = sha256::sha256(data);
232    let mut result = [0u8; 16];
233    result.copy_from_slice(&full[..16]);
234    result
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240    use crate::FixedRng;
241
242    #[test]
243    fn test_identity_key_roundtrip() {
244        let mut rng = FixedRng::new(&(0..64).collect::<Vec<u8>>());
245        let id = Identity::new(&mut rng);
246        let prv_bytes = id.get_private_key().unwrap();
247        let id2 = Identity::from_private_key(&prv_bytes);
248        assert_eq!(
249            id.get_public_key().unwrap(),
250            id2.get_public_key().unwrap()
251        );
252    }
253
254    #[test]
255    fn test_identity_hash() {
256        let mut rng = FixedRng::new(&(0..64).collect::<Vec<u8>>());
257        let id = Identity::new(&mut rng);
258        let pub_key = id.get_public_key().unwrap();
259        let expected_hash = truncated_hash(&pub_key);
260        assert_eq!(*id.hash(), expected_hash);
261    }
262
263    #[test]
264    fn test_identity_encrypt_decrypt_roundtrip() {
265        let mut rng = FixedRng::new(&(0..128).collect::<Vec<u8>>());
266        let id = Identity::new(&mut rng);
267        let plaintext = b"Hello, Reticulum! This is a test of the encrypt/decrypt pipeline.";
268        let mut rng2 = FixedRng::new(&(128..255).collect::<Vec<u8>>());
269        let ciphertext = id.encrypt(plaintext, &mut rng2).unwrap();
270        let decrypted = id.decrypt(&ciphertext).unwrap();
271        assert_eq!(decrypted, plaintext);
272    }
273
274    #[test]
275    fn test_identity_sign_verify() {
276        let mut rng = FixedRng::new(&(0..64).collect::<Vec<u8>>());
277        let id = Identity::new(&mut rng);
278        let msg = b"Sign this message";
279        let sig = id.sign(msg).unwrap();
280        assert!(id.verify(&sig, msg));
281        assert!(!id.verify(&sig, b"Wrong message"));
282    }
283}