1use blake3::Hasher;
48use chacha20poly1305::{
49 ChaCha20Poly1305, Nonce,
50 aead::{Aead, KeyInit},
51};
52use curve25519_dalek::{
53 constants::RISTRETTO_BASEPOINT_TABLE,
54 ristretto::{CompressedRistretto, RistrettoPoint},
55 scalar::Scalar,
56};
57use rand::Rng;
58use serde::{Deserialize, Serialize};
59
60#[derive(Debug, Clone, PartialEq, Eq)]
62pub enum ProxyReError {
63 InvalidCiphertext,
65 DecryptionFailed,
67 EncryptionFailed,
69 InvalidPublicKey,
71 InvalidReKey,
73 SerializationError,
75}
76
77impl std::fmt::Display for ProxyReError {
78 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
79 match self {
80 ProxyReError::InvalidCiphertext => write!(f, "Invalid ciphertext format"),
81 ProxyReError::DecryptionFailed => write!(f, "Decryption failed"),
82 ProxyReError::EncryptionFailed => write!(f, "Encryption failed"),
83 ProxyReError::InvalidPublicKey => write!(f, "Invalid public key"),
84 ProxyReError::InvalidReKey => write!(f, "Invalid re-encryption key"),
85 ProxyReError::SerializationError => write!(f, "Serialization/deserialization error"),
86 }
87 }
88}
89
90impl std::error::Error for ProxyReError {}
91
92pub type ProxyReResult<T> = Result<T, ProxyReError>;
94
95#[derive(Clone, Serialize, Deserialize)]
97pub struct ProxyReSecretKey(Scalar);
98
99impl ProxyReSecretKey {
100 pub fn generate() -> Self {
102 let mut rng = rand::thread_rng();
103 let mut bytes = [0u8; 32];
104 rng.fill(&mut bytes);
105 Self(Scalar::from_bytes_mod_order(bytes))
106 }
107
108 pub fn to_bytes(&self) -> [u8; 32] {
110 self.0.to_bytes()
111 }
112
113 pub fn from_bytes(bytes: &[u8; 32]) -> Self {
115 Self(Scalar::from_bytes_mod_order(*bytes))
116 }
117}
118
119#[derive(Clone, Copy, Debug, Serialize, Deserialize, PartialEq, Eq)]
121pub struct ProxyRePublicKey(RistrettoPoint);
122
123impl ProxyRePublicKey {
124 pub fn from_secret(secret: &ProxyReSecretKey) -> Self {
126 Self(&secret.0 * RISTRETTO_BASEPOINT_TABLE)
127 }
128
129 pub fn to_bytes(&self) -> [u8; 32] {
131 self.0.compress().to_bytes()
132 }
133
134 pub fn from_bytes(bytes: &[u8; 32]) -> ProxyReResult<Self> {
136 CompressedRistretto(*bytes)
137 .decompress()
138 .map(Self)
139 .ok_or(ProxyReError::InvalidPublicKey)
140 }
141}
142
143#[derive(Clone)]
145pub struct ProxyReKeypair {
146 secret: ProxyReSecretKey,
147 public: ProxyRePublicKey,
148}
149
150impl ProxyReKeypair {
151 pub fn generate() -> Self {
153 let secret = ProxyReSecretKey::generate();
154 let public = ProxyRePublicKey::from_secret(&secret);
155 Self { secret, public }
156 }
157
158 pub fn public_key(&self) -> ProxyRePublicKey {
160 self.public
161 }
162
163 pub fn secret_key(&self) -> &ProxyReSecretKey {
165 &self.secret
166 }
167
168 pub fn encrypt(&self, plaintext: &[u8]) -> ProxyReResult<ProxyReCiphertext> {
170 encrypt(&self.public, plaintext)
171 }
172
173 pub fn decrypt(&self, ciphertext: &ProxyReCiphertext) -> ProxyReResult<Vec<u8>> {
175 decrypt(&self.secret, ciphertext)
176 }
177
178 pub fn generate_re_key(&self, target_pk: &ProxyRePublicKey) -> ProxyReReKey {
180 generate_re_key(&self.secret, target_pk)
181 }
182}
183
184#[derive(Clone, Serialize, Deserialize)]
186pub struct ProxyReReKey {
187 re_key: Scalar,
189 target_pk: ProxyRePublicKey,
191}
192
193#[derive(Clone, Serialize, Deserialize)]
195pub struct ProxyReCiphertext {
196 ephemeral_pk: RistrettoPoint,
198 encrypted_key: Vec<u8>,
200 ciphertext: Vec<u8>,
202 nonce: [u8; 12],
204}
205
206pub fn encrypt(pk: &ProxyRePublicKey, plaintext: &[u8]) -> ProxyReResult<ProxyReCiphertext> {
208 let mut rng = rand::thread_rng();
209
210 let ephemeral_sk = ProxyReSecretKey::generate();
212 let ephemeral_pk = ProxyRePublicKey::from_secret(&ephemeral_sk);
213
214 let shared_point = pk.0 * ephemeral_sk.0;
216
217 let sym_key = derive_symmetric_key(&shared_point);
219
220 let mut nonce_bytes = [0u8; 12];
222 rng.fill(&mut nonce_bytes);
223 let nonce = Nonce::from_slice(&nonce_bytes);
224
225 let cipher = ChaCha20Poly1305::new(&sym_key.into());
227 let ciphertext = cipher
228 .encrypt(nonce, plaintext)
229 .map_err(|_| ProxyReError::EncryptionFailed)?;
230
231 let encrypted_key = vec![0u8; 32];
235
236 Ok(ProxyReCiphertext {
237 ephemeral_pk: ephemeral_pk.0,
238 encrypted_key,
239 ciphertext,
240 nonce: nonce_bytes,
241 })
242}
243
244pub fn decrypt(sk: &ProxyReSecretKey, ciphertext: &ProxyReCiphertext) -> ProxyReResult<Vec<u8>> {
246 let shared_point = ciphertext.ephemeral_pk * sk.0;
248
249 let sym_key = derive_symmetric_key(&shared_point);
251
252 let cipher = ChaCha20Poly1305::new(&sym_key.into());
254 let nonce = Nonce::from_slice(&ciphertext.nonce);
255
256 cipher
257 .decrypt(nonce, ciphertext.ciphertext.as_ref())
258 .map_err(|_| ProxyReError::DecryptionFailed)
259}
260
261pub fn generate_re_key(
263 delegator_sk: &ProxyReSecretKey,
264 delegatee_pk: &ProxyRePublicKey,
265) -> ProxyReReKey {
266 let re_key = delegator_sk.0.invert();
269
270 ProxyReReKey {
271 re_key,
272 target_pk: *delegatee_pk,
273 }
274}
275
276pub fn re_encrypt(
278 ciphertext: &ProxyReCiphertext,
279 re_key: &ProxyReReKey,
280) -> ProxyReResult<ProxyReCiphertext> {
281 let mut rng = rand::thread_rng();
299
300 let re_ephemeral_sk = ProxyReSecretKey::generate();
302 let re_ephemeral_pk = ProxyRePublicKey::from_secret(&re_ephemeral_sk);
303
304 let new_shared_point = re_key.target_pk.0 * re_ephemeral_sk.0;
306 let new_sym_key = derive_symmetric_key(&new_shared_point);
307
308 let mut nonce_bytes = [0u8; 12];
310 rng.fill(&mut nonce_bytes);
311 let nonce = Nonce::from_slice(&nonce_bytes);
312
313 let cipher = ChaCha20Poly1305::new(&new_sym_key.into());
314
315 let original_serialized =
317 crate::codec::encode(ciphertext).map_err(|_| ProxyReError::SerializationError)?;
318
319 let new_ciphertext = cipher
320 .encrypt(nonce, original_serialized.as_ref())
321 .map_err(|_| ProxyReError::EncryptionFailed)?;
322
323 Ok(ProxyReCiphertext {
324 ephemeral_pk: re_ephemeral_pk.0,
325 encrypted_key: vec![1u8; 32], ciphertext: new_ciphertext,
327 nonce: nonce_bytes,
328 })
329}
330
331fn derive_symmetric_key(point: &RistrettoPoint) -> [u8; 32] {
333 let mut hasher = Hasher::new();
334 hasher.update(b"chie-proxy-re-v1");
335 hasher.update(&point.compress().to_bytes());
336 let hash = hasher.finalize();
337 *hash.as_bytes()
338}
339
340#[cfg(test)]
341mod tests {
342 use super::*;
343
344 #[test]
345 fn test_keypair_generation() {
346 let keypair = ProxyReKeypair::generate();
347 let pk_derived = ProxyRePublicKey::from_secret(keypair.secret_key());
348 assert_eq!(pk_derived, keypair.public_key());
349 }
350
351 #[test]
352 fn test_basic_encryption_decryption() {
353 let keypair = ProxyReKeypair::generate();
354 let plaintext = b"Hello, proxy re-encryption!";
355
356 let ciphertext = keypair.encrypt(plaintext).unwrap();
357 let decrypted = keypair.decrypt(&ciphertext).unwrap();
358
359 assert_eq!(decrypted, plaintext);
360 }
361
362 #[test]
363 fn test_encryption_produces_different_ciphertexts() {
364 let keypair = ProxyReKeypair::generate();
365 let plaintext = b"Test message";
366
367 let ct1 = keypair.encrypt(plaintext).unwrap();
368 let ct2 = keypair.encrypt(plaintext).unwrap();
369
370 assert_ne!(ct1.ephemeral_pk.compress(), ct2.ephemeral_pk.compress());
372 assert_ne!(ct1.ciphertext, ct2.ciphertext);
373 }
374
375 #[test]
376 fn test_wrong_key_decryption_fails() {
377 let alice = ProxyReKeypair::generate();
378 let bob = ProxyReKeypair::generate();
379
380 let plaintext = b"Secret message";
381 let ciphertext = alice.encrypt(plaintext).unwrap();
382
383 assert!(bob.decrypt(&ciphertext).is_err());
385 }
386
387 #[test]
388 fn test_proxy_re_encryption() {
389 let alice = ProxyReKeypair::generate();
390 let bob = ProxyReKeypair::generate();
391
392 let plaintext = b"Delegated content";
393
394 let ciphertext = alice.encrypt(plaintext).unwrap();
396
397 let alice_decrypted = alice.decrypt(&ciphertext).unwrap();
399 assert_eq!(alice_decrypted, plaintext);
400
401 let re_key = alice.generate_re_key(&bob.public_key());
403
404 let re_encrypted = re_encrypt(&ciphertext, &re_key).unwrap();
406
407 let outer_decrypted = bob.decrypt(&re_encrypted).unwrap();
410
411 let inner_ciphertext: ProxyReCiphertext = crate::codec::decode(&outer_decrypted).unwrap();
413
414 let final_plaintext = alice.decrypt(&inner_ciphertext).unwrap();
416 assert_eq!(final_plaintext, plaintext);
417 }
418
419 #[test]
420 fn test_public_key_serialization() {
421 let keypair = ProxyReKeypair::generate();
422 let pk = keypair.public_key();
423
424 let bytes = pk.to_bytes();
425 let pk_restored = ProxyRePublicKey::from_bytes(&bytes).unwrap();
426
427 assert_eq!(pk, pk_restored);
428 }
429
430 #[test]
431 fn test_secret_key_serialization() {
432 let keypair = ProxyReKeypair::generate();
433 let sk = keypair.secret_key();
434
435 let bytes = sk.to_bytes();
436 let sk_restored = ProxyReSecretKey::from_bytes(&bytes);
437
438 let pk1 = ProxyRePublicKey::from_secret(sk);
440 let pk2 = ProxyRePublicKey::from_secret(&sk_restored);
441 assert_eq!(pk1, pk2);
442 }
443
444 #[test]
445 fn test_invalid_public_key() {
446 let invalid_bytes = [255u8; 32];
447 assert!(ProxyRePublicKey::from_bytes(&invalid_bytes).is_err());
448 }
449
450 #[test]
451 fn test_ciphertext_serialization() {
452 let keypair = ProxyReKeypair::generate();
453 let plaintext = b"Serialize this";
454
455 let ciphertext = keypair.encrypt(plaintext).unwrap();
456 let serialized = crate::codec::encode(&ciphertext).unwrap();
457 let deserialized: ProxyReCiphertext = crate::codec::decode(&serialized).unwrap();
458
459 let decrypted = keypair.decrypt(&deserialized).unwrap();
460 assert_eq!(decrypted, plaintext);
461 }
462
463 #[test]
464 fn test_empty_plaintext() {
465 let keypair = ProxyReKeypair::generate();
466 let plaintext = b"";
467
468 let ciphertext = keypair.encrypt(plaintext).unwrap();
469 let decrypted = keypair.decrypt(&ciphertext).unwrap();
470
471 assert_eq!(decrypted, plaintext);
472 }
473
474 #[test]
475 fn test_large_plaintext() {
476 let keypair = ProxyReKeypair::generate();
477 let plaintext = vec![42u8; 10_000];
478
479 let ciphertext = keypair.encrypt(&plaintext).unwrap();
480 let decrypted = keypair.decrypt(&ciphertext).unwrap();
481
482 assert_eq!(decrypted, plaintext);
483 }
484
485 #[test]
486 fn test_multiple_delegations() {
487 let alice = ProxyReKeypair::generate();
488 let bob = ProxyReKeypair::generate();
489 let carol = ProxyReKeypair::generate();
490
491 let plaintext = b"Multi-hop delegation";
492
493 let ct_alice = alice.encrypt(plaintext).unwrap();
495
496 let re_key_alice_to_bob = alice.generate_re_key(&bob.public_key());
498 let ct_bob = re_encrypt(&ct_alice, &re_key_alice_to_bob).unwrap();
499
500 let re_key_alice_to_carol = alice.generate_re_key(&carol.public_key());
502 let ct_carol = re_encrypt(&ct_alice, &re_key_alice_to_carol).unwrap();
503
504 let bob_outer = bob.decrypt(&ct_bob).unwrap();
506 let carol_outer = carol.decrypt(&ct_carol).unwrap();
507
508 assert!(crate::codec::decode::<ProxyReCiphertext>(&bob_outer).is_ok());
509 assert!(crate::codec::decode::<ProxyReCiphertext>(&carol_outer).is_ok());
510 }
511
512 #[test]
513 fn test_re_key_serialization() {
514 let alice = ProxyReKeypair::generate();
515 let bob = ProxyReKeypair::generate();
516
517 let re_key = alice.generate_re_key(&bob.public_key());
518 let serialized = crate::codec::encode(&re_key).unwrap();
519 let deserialized: ProxyReReKey = crate::codec::decode(&serialized).unwrap();
520
521 assert_eq!(re_key.target_pk, deserialized.target_pk);
522 }
523}