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