Add support for encrypting/decrypting with the PublicKey/SecretKey interface

Cargo.toml

@@ -32,6 +32,8 @@ fjall = "2.9.0"
 futures = "0.3.31"
 futures-util = "0.3.31"
 hex = "0.4.3"
+chacha20poly1305 = "0.10.1"
+hkdf = "0.12.4"
 jiff = { version = "0.2.3", features = ["serde"] }
 k256 = "0.13.4"
 kademlia-discovery = { path = "packages/examples/kademlia-discovery" }

packages/shared/Cargo.toml

@@ -6,19 +6,30 @@ edition = "2021"
 license = "MIT"
 
 [dependencies]
+chacha20poly1305 = { workspace = true, optional = true }
 borsh = { workspace = true, features = ["derive"], optional = true }
 cosmwasm-std = { workspace = true, optional = true }
 cw-storage-plus = { workspace = true, optional = true }
+hkdf = { workspace = true, optional = true }
 hex = { workspace = true, optional = true }
-k256 = { workspace = true, features = ["ecdsa", "serde", "sha256"], optional = true }
+k256 = { workspace = true, features = ["ecdsa", "ecdh", "serde", "sha256"], optional = true }
 rand = { workspace = true, optional = true }
 serde = { workspace = true }
+sha2 = { workspace = true, optional = true }
 thiserror = { workspace = true }
 
 [features]
 chaincryptography = []
 cosmwasm = ["dep:cosmwasm-std", "dep:cw-storage-plus"]
-realcryptography = ["borsh?/std", "dep:hex", "dep:k256", "dep:rand"]
+realcryptography = [
+  "borsh?/std",
+  "dep:chacha20poly1305",
+  "dep:hex",
+  "dep:hkdf",
+  "dep:k256",
+  "dep:rand",
+  "dep:sha2",
+]
 solana = ["dep:borsh", "dep:hex"]
 
 [lints]

packages/shared/src/cryptography/real.rs

@@ -1,7 +1,16 @@
 pub use rand::rngs::ThreadRng;
 
+use chacha20poly1305::{
+    aead::{AeadInPlace, KeyInit},
+    ChaCha20Poly1305, Key, Nonce, Tag,
+};
+use k256::ecdh::{diffie_hellman, EphemeralSecret, SharedSecret};
+use rand::{CryptoRng, RngCore};
+use sha2::Sha256;
 use std::{fmt::Display, str::FromStr};
 
+const KEY_DERIVATION_INFO: &[u8] = b"kolme/chacha20poly1305-v1";
+
 /// Newtype wrapper around [k256::PublicKey] to provide consistent serialization.
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 pub struct PublicKey(k256::PublicKey);
@@ -55,6 +64,41 @@ impl PublicKey {
                 bytes: bytes.to_owned(),
             })
     }
+
+    /// Encrypt data so it can be decrypted by the corresponding [SecretKey].
+    pub fn encrypt(
+        &self,
+        plaintext: impl AsRef<[u8]>,
+    ) -> Result<EncryptedMessage, EncryptionError> {
+        self.encrypt_with(&mut rand::thread_rng(), plaintext)
+    }
+
+    pub fn encrypt_with(
+        &self,
+        rng: &mut (impl CryptoRng + RngCore),
+        plaintext: impl AsRef<[u8]>,
+    ) -> Result<EncryptedMessage, EncryptionError> {
+        let ephemeral_secret = EphemeralSecret::random(rng);
+        let ephemeral_public_key: PublicKey = k256::PublicKey::from(&ephemeral_secret).into();
+        let shared_secret = ephemeral_secret.diffie_hellman(&self.0);
+        let key = derive_encryption_key(&shared_secret)?;
+
+        let mut nonce_bytes = [0u8; 12];
+        rng.fill_bytes(&mut nonce_bytes);
+
+        let cipher = ChaCha20Poly1305::new(&key);
+        let mut buffer = plaintext.as_ref().to_vec();
+        let tag = cipher
+            .encrypt_in_place_detached(Nonce::from_slice(&nonce_bytes), &[], &mut buffer)
+            .map_err(|_| EncryptionError::EncryptionFailed)?;
+
+        Ok(EncryptedMessage {
+            ephemeral_public_key,
+            nonce: nonce_bytes,
+            ciphertext: buffer,
+            tag: tag.into(),
+        })
+    }
 }
 
 impl From<k256::PublicKey> for PublicKey {
@@ -157,6 +201,24 @@ pub enum SecretKeyError {
     InvalidBytes { source: k256::elliptic_curve::Error },
 }
 
+#[derive(thiserror::Error, Debug)]
+pub enum EncryptionError {
+    #[error("Key derivation failed")]
+    KeyDerivationFailed,
+    #[error("Encryption failed")]
+    EncryptionFailed,
+    #[error("Decryption failed")]
+    DecryptionFailed,
+}
+
+#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+pub struct EncryptedMessage {
+    pub ephemeral_public_key: PublicKey,
+    pub nonce: [u8; 12],
+    pub ciphertext: Vec<u8>,
+    pub tag: [u8; 16],
+}
+
 #[derive(Clone, PartialEq, Eq)]
 pub struct SecretKey(k256::SecretKey);
 
@@ -215,6 +277,23 @@ impl SecretKey {
     pub fn reveal_as_hex(&self) -> String {
         hex::encode(self.0.to_bytes())
     }
+
+    pub fn decrypt(&self, message: &EncryptedMessage) -> Result<Vec<u8>, EncryptionError> {
+        let shared_secret = diffie_hellman(
+            self.0.to_nonzero_scalar(),
+            message.ephemeral_public_key.0.as_affine(),
+        );
+        let key = derive_encryption_key(&shared_secret)?;
+
+        let cipher = ChaCha20Poly1305::new(&key);
+        let mut buffer = message.ciphertext.clone();
+        let tag = Tag::from_slice(&message.tag);
+        cipher
+            .decrypt_in_place_detached(Nonce::from_slice(&message.nonce), &[], &mut buffer, tag)
+            .map_err(|_| EncryptionError::DecryptionFailed)?;
+
+        Ok(buffer)
+    }
 }
 
 impl FromStr for SecretKey {
@@ -231,6 +310,15 @@ impl std::fmt::Debug for SecretKey {
     }
 }
 
+fn derive_encryption_key(shared_secret: &SharedSecret) -> Result<Key, EncryptionError> {
+    let hkdf = shared_secret.extract::<Sha256>(None);
+    let mut okm = [0u8; 32];
+    hkdf.expand(KEY_DERIVATION_INFO, &mut okm)
+        .map_err(|_| EncryptionError::KeyDerivationFailed)?;
+
+    Ok(Key::clone_from_slice(&okm))
+}
+
 mod sigerr {
     #[derive(thiserror::Error, Debug)]
     pub enum SignatureError {
@@ -416,3 +504,42 @@ impl SignatureWithRecovery {
         PublicKey::recover_from_msg(msg, &self.sig, self.recid)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn encrypt_decrypt_roundtrip() {
+        let receiver = SecretKey::random();
+        let plaintext = b"hello world";
+
+        let encrypted = receiver.public_key().encrypt(plaintext).unwrap();
+        let decrypted = receiver.decrypt(&encrypted).unwrap();
+
+        assert_eq!(decrypted, plaintext);
+    }
+
+    #[test]
+    fn encrypt_is_randomized() {
+        let receiver = SecretKey::random();
+        let plaintext = b"stable text";
+
+        let first = receiver.public_key().encrypt(plaintext).unwrap();
+        let second = receiver.public_key().encrypt(plaintext).unwrap();
+
+        assert_ne!(first.nonce, second.nonce);
+        assert_ne!(first.ciphertext, second.ciphertext);
+        assert_ne!(first.ephemeral_public_key, second.ephemeral_public_key);
+    }
+
+    #[test]
+    fn tampering_fails() {
+        let receiver = SecretKey::random();
+        let mut encrypted = receiver.public_key().encrypt(b"super secret").unwrap();
+
+        encrypted.ciphertext[0] ^= 0b0000_0001;
+
+        assert!(receiver.decrypt(&encrypted).is_err());
+    }
+}