1use anyhow::{anyhow, Result};
6use chacha20poly1305::{
7 aead::{Aead, AeadCore, KeyInit},
8 ChaCha20Poly1305, Key, Nonce,
9};
10use rand_core::OsRng;
11use sha2::{Digest, Sha256};
12use x25519_dalek::{PublicKey, SharedSecret, StaticSecret};
13
14pub struct HpkeContext {
16 private_key: StaticSecret,
17}
18
19impl HpkeContext {
20 pub fn new() -> Result<Self> {
28 let private_key = StaticSecret::random_from_rng(OsRng);
29 Ok(Self { private_key })
30 }
31
32 pub fn from_private_key(private_key_bytes: &[u8]) -> Result<Self> {
42 let key_array: [u8; 32] = private_key_bytes
43 .try_into()
44 .map_err(|_| anyhow!("Private key must be 32 bytes"))?;
45
46 let private_key = StaticSecret::from(key_array);
47 Ok(Self { private_key })
48 }
49
50 pub fn public_key_bytes(&self) -> Vec<u8> {
52 PublicKey::from(&self.private_key).as_bytes().to_vec()
53 }
54
55 pub fn private_key_bytes(&self) -> Vec<u8> {
57 self.private_key.to_bytes().to_vec()
58 }
59
60 pub fn seal(&self, recipient_pub: &[u8], plaintext: &[u8]) -> Result<Vec<u8>> {
75 let recipient_pk = self.parse_public_key(recipient_pub)?;
76 let shared_secret = self.private_key.diffie_hellman(&recipient_pk);
77
78 let key = self.derive_symmetric_key(&shared_secret);
79
80 let cipher = ChaCha20Poly1305::new(&key);
81 let nonce = ChaCha20Poly1305::generate_nonce(&mut OsRng);
82
83 let ciphertext = cipher
84 .encrypt(&nonce, plaintext)
85 .map_err(|e| anyhow!("Encryption failed: {}", e))?;
86
87 let mut result = nonce.to_vec();
88 result.extend_from_slice(&ciphertext);
89
90 Ok(result)
91 }
92
93 pub fn open(&self, sender_pub: &[u8], ciphertext: &[u8]) -> Result<Vec<u8>> {
108 if ciphertext.len() < 12 {
109 return Err(anyhow!("Ciphertext is too short"));
110 }
111
112 let (nonce_bytes, encrypted_data) = ciphertext.split_at(12);
113 let nonce = Nonce::from_slice(nonce_bytes);
114
115 let sender_pk = self.parse_public_key(sender_pub)?;
116 let shared_secret = self.private_key.diffie_hellman(&sender_pk);
117
118 let key = self.derive_symmetric_key(&shared_secret);
119
120 let cipher = ChaCha20Poly1305::new(&key);
121
122 let plaintext = cipher
123 .decrypt(nonce, encrypted_data)
124 .map_err(|e| anyhow!("Decryption failed: {}", e))?;
125
126 Ok(plaintext)
127 }
128
129 fn parse_public_key(&self, key_bytes: &[u8]) -> Result<PublicKey> {
131 let key_array: [u8; 32] = key_bytes
132 .try_into()
133 .map_err(|_| anyhow!("Public key must be 32 bytes"))?;
134
135 Ok(PublicKey::from(key_array))
136 }
137
138 fn derive_symmetric_key(&self, shared_secret: &SharedSecret) -> Key {
140 let mut hasher = Sha256::new();
141 hasher.update(shared_secret.as_bytes());
142 hasher.finalize()
143 }
144}