p2p_chat/crypto/
hpke.rs

1//! This module implements a simplified version of Hybrid Public Key Encryption (HPKE).
2//!
3//! It uses X25519 for key exchange and ChaCha20Poly1305 for symmetric encryption.
4//! The key derivation is done using SHA-256.
5use 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
14/// A context for performing E2E encryption using X25519 and ChaCha20Poly1305.
15pub struct HpkeContext {
16    private_key: StaticSecret,
17}
18
19impl HpkeContext {
20    /// Generates a new keypair.
21    ///
22    /// This will create a new `HpkeContext` with a randomly generated X25519 keypair.
23    ///
24    /// # Errors
25    ///
26    /// This function will return an error if the key generation fails.
27    pub fn new() -> Result<Self> {
28        let private_key = StaticSecret::random_from_rng(OsRng);
29        Ok(Self { private_key })
30    }
31
32    /// Loads a keypair from existing private key bytes.
33    ///
34    /// # Arguments
35    ///
36    /// * `private_key_bytes` - A 32-byte slice representing the private key.
37    ///
38    /// # Errors
39    ///
40    /// This function will return an error if the private key is not 32 bytes long.
41    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    /// Returns the public key bytes corresponding to our private key.
51    pub fn public_key_bytes(&self) -> Vec<u8> {
52        PublicKey::from(&self.private_key).as_bytes().to_vec()
53    }
54
55    /// Returns the private key bytes.
56    pub fn private_key_bytes(&self) -> Vec<u8> {
57        self.private_key.to_bytes().to_vec()
58    }
59
60    /// Derives a shared secret and uses it to encrypt a message for a recipient.
61    ///
62    /// # Arguments
63    ///
64    /// * `recipient_pub` - The public key of the recipient.
65    /// * `plaintext` - The data to encrypt.
66    ///
67    /// # Returns
68    ///
69    /// A vector containing the nonce and the ciphertext.
70    ///
71    /// # Errors
72    ///
73    /// This function will return an error if encryption fails.
74    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    /// Derives a shared secret and uses it to decrypt a message from a sender.
94    ///
95    /// # Arguments
96    ///
97    /// * `sender_pub` - The public key of the sender.
98    /// * `ciphertext` - The data to decrypt (nonce prepended).
99    ///
100    /// # Returns
101    ///
102    /// A vector containing the plaintext.
103    ///
104    /// # Errors
105    ///
106    /// This function will return an error if decryption fails.
107    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    /// Parses a public key from a byte slice.
130    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    /// Derives a symmetric key from a shared secret.
139    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}