use aes::cipher::{
BlockDecryptMut as _, BlockEncryptMut as _, KeyIvInit as _,
block_padding::{Pkcs7, UnpadError},
};
use hmac::{Mac as _, digest::MacError};
use matrix_pickle::{Decode, Encode};
use thiserror::Error;
use zeroize::{Zeroize, ZeroizeOnDrop};
use crate::{
Curve25519PublicKey, Curve25519SecretKey, KeyError, base64_decode,
cipher::{
Aes256CbcDec, Aes256CbcEnc, HmacSha256, Mac,
key::{CipherKeys, ExpandedKeys},
},
};
const PICKLE_VERSION: u32 = 1;
#[derive(Debug, Error)]
pub enum Error {
#[error(
"One or more keys lacked contributory behavior in the Diffie-Hellman operation, \
resulting in an insecure shared secret"
)]
NonContributoryKey,
#[error("failed to decrypt, invalid padding: {0}")]
InvalidPadding(#[from] UnpadError),
#[error("the MAC of the ciphertext didn't pass validation: {0}")]
Mac(#[from] MacError),
}
#[derive(Debug, Error)]
pub enum MessageDecodeError {
#[error(transparent)]
Base64(#[from] crate::Base64DecodeError),
#[error(transparent)]
Key(#[from] KeyError),
}
#[derive(Debug)]
pub struct Message {
pub ciphertext: Vec<u8>,
pub mac: Vec<u8>,
pub ephemeral_key: Curve25519PublicKey,
}
impl Message {
pub fn from_base64(
ciphertext: &str,
mac: &str,
ephemeral_key: &str,
) -> Result<Self, MessageDecodeError> {
Ok(Self {
ciphertext: base64_decode(ciphertext)?,
mac: base64_decode(mac)?,
ephemeral_key: Curve25519PublicKey::from_base64(ephemeral_key)?,
})
}
}
pub struct PkDecryption {
secret_key: Curve25519SecretKey,
public_key: Curve25519PublicKey,
}
impl PkDecryption {
pub fn new() -> Self {
let secret_key = Curve25519SecretKey::new();
let public_key = Curve25519PublicKey::from(&secret_key);
Self { secret_key, public_key }
}
pub fn from_key(secret_key: Curve25519SecretKey) -> Self {
let public_key = Curve25519PublicKey::from(&secret_key);
Self { secret_key, public_key }
}
pub const fn secret_key(&self) -> &Curve25519SecretKey {
&self.secret_key
}
pub const fn public_key(&self) -> Curve25519PublicKey {
self.public_key
}
pub fn from_libolm_pickle(
pickle: &str,
pickle_key: &[u8],
) -> Result<Self, crate::LibolmPickleError> {
use crate::utilities::unpickle_libolm;
unpickle_libolm::<PkDecryptionPickle, _>(pickle, pickle_key, PICKLE_VERSION)
}
pub fn to_libolm_pickle(&self, pickle_key: &[u8]) -> Result<String, crate::LibolmPickleError> {
use crate::utilities::pickle_libolm;
pickle_libolm::<PkDecryptionPickle>(self.into(), pickle_key)
}
pub fn decrypt(&self, message: &Message) -> Result<Vec<u8>, Error> {
let shared_secret = self
.secret_key
.diffie_hellman(&message.ephemeral_key)
.ok_or(Error::NonContributoryKey)?;
let expanded_keys = ExpandedKeys::new_helper(shared_secret.as_bytes(), b"");
let cipher_keys = CipherKeys::from_expanded_keys(expanded_keys);
#[allow(clippy::expect_used)]
let hmac = HmacSha256::new_from_slice(cipher_keys.mac_key())
.expect("We should be able to create a Hmac object from a 32 byte key");
hmac.verify_truncated_left(&message.mac)?;
let cipher = Aes256CbcDec::new(cipher_keys.aes_key(), cipher_keys.iv());
let decrypted = cipher.decrypt_padded_vec_mut::<Pkcs7>(&message.ciphertext)?;
Ok(decrypted)
}
}
impl Default for PkDecryption {
fn default() -> Self {
Self::new()
}
}
impl TryFrom<PkDecryptionPickle> for PkDecryption {
type Error = crate::LibolmPickleError;
fn try_from(pickle: PkDecryptionPickle) -> Result<Self, Self::Error> {
let secret_key = Curve25519SecretKey::from_slice(&pickle.private_curve25519_key);
let public_key = Curve25519PublicKey::from(&secret_key);
Ok(Self { secret_key, public_key })
}
}
#[derive(Encode, Decode, Zeroize, ZeroizeOnDrop)]
struct PkDecryptionPickle {
version: u32,
public_curve25519_key: [u8; 32],
#[secret]
private_curve25519_key: Box<[u8; 32]>,
}
impl From<&PkDecryption> for PkDecryptionPickle {
fn from(decrypt: &PkDecryption) -> Self {
Self {
version: PICKLE_VERSION,
public_curve25519_key: decrypt.public_key.to_bytes(),
private_curve25519_key: decrypt.secret_key.to_bytes(),
}
}
}
pub struct PkEncryption {
public_key: Curve25519PublicKey,
}
impl PkEncryption {
pub const fn from_key(public_key: Curve25519PublicKey) -> Self {
Self { public_key }
}
pub fn encrypt(&self, message: &[u8]) -> Result<Message, Error> {
let ephemeral_key = Curve25519SecretKey::new();
let shared_secret =
ephemeral_key.diffie_hellman(&self.public_key).ok_or(Error::NonContributoryKey)?;
let expanded_keys = ExpandedKeys::new_helper(shared_secret.as_bytes(), b"");
let cipher_keys = CipherKeys::from_expanded_keys(expanded_keys);
let cipher = Aes256CbcEnc::new(cipher_keys.aes_key(), cipher_keys.iv());
let ciphertext = cipher.encrypt_padded_vec_mut::<Pkcs7>(message);
#[allow(clippy::expect_used)]
let hmac = HmacSha256::new_from_slice(cipher_keys.mac_key())
.expect("We should be able to create a Hmac object from a 32 byte key");
let mut mac = hmac.finalize().into_bytes().to_vec();
mac.truncate(Mac::TRUNCATED_LEN);
Ok(Message { ciphertext, mac, ephemeral_key: Curve25519PublicKey::from(&ephemeral_key) })
}
}
impl From<&PkDecryption> for PkEncryption {
fn from(value: &PkDecryption) -> Self {
Self::from_key(value.public_key())
}
}
impl From<Curve25519PublicKey> for PkEncryption {
fn from(public_key: Curve25519PublicKey) -> Self {
Self { public_key }
}
}
#[cfg(test)]
mod tests {
use olm_rs::pk::{OlmPkDecryption, OlmPkEncryption, PkMessage};
use super::{Message, MessageDecodeError, PkDecryption, PkEncryption};
use crate::{Curve25519PublicKey, Curve25519SecretKey, base64_encode};
impl TryFrom<PkMessage> for Message {
type Error = MessageDecodeError;
fn try_from(value: PkMessage) -> Result<Self, Self::Error> {
Self::from_base64(&value.ciphertext, &value.mac, &value.ephemeral_key)
}
}
impl From<Message> for PkMessage {
fn from(val: Message) -> Self {
PkMessage {
ciphertext: base64_encode(val.ciphertext),
mac: base64_encode(val.mac),
ephemeral_key: val.ephemeral_key.to_base64(),
}
}
}
#[test]
fn decrypt_libolm_encrypted_message() {
let decryptor = PkDecryption::new();
let public_key = decryptor.public_key();
let encryptor = OlmPkEncryption::new(&public_key.to_base64());
let message = "It's a secret to everybody";
let encrypted = encryptor.encrypt(message);
let encrypted =
encrypted.try_into().expect("We should be able to decode a message libolm created");
let decrypted = decryptor
.decrypt(&encrypted)
.expect("We should be able to decrypt a message libolm encrypted");
assert_eq!(
message.as_bytes(),
decrypted,
"The plaintext should match the decrypted message"
);
}
#[test]
fn encrypt_for_libolm_pk_decryption() {
let decryptor = OlmPkDecryption::new();
let public_key = Curve25519PublicKey::from_base64(decryptor.public_key())
.expect("libolm should provide us with a valid Curve25519 public key");
let encryptor = PkEncryption::from_key(public_key);
let message = "It's a secret to everybody";
let encrypted = encryptor.encrypt(message.as_ref()).unwrap();
let encrypted = encrypted.into();
let decrypted = decryptor
.decrypt(encrypted)
.expect("We should be able to decrypt a message vodozemac encrypted using libolm");
assert_eq!(message, decrypted, "The plaintext should match the decrypted message");
}
#[test]
fn encryption_roundtrip() {
let decryptor = PkDecryption::new();
let public_key = decryptor.public_key();
let encryptor = PkEncryption::from_key(public_key);
let message = "It's a secret to everybody";
let encrypted = encryptor.encrypt(message.as_ref()).unwrap();
let decrypted = decryptor
.decrypt(&encrypted)
.expect("We should be able to decrypt a message we encrypted");
assert_eq!(message.as_ref(), decrypted, "The plaintext should match the decrypted message");
}
#[test]
fn from_bytes() {
let decryption = PkDecryption::default();
let bytes = decryption.secret_key().to_bytes();
let secret_key = Curve25519SecretKey::from_slice(&bytes);
let restored = PkDecryption::from_key(secret_key);
assert_eq!(
decryption.public_key(),
restored.public_key(),
"The public keys of the restored and original PK decryption should match"
);
}
#[test]
fn libolm_unpickling() {
let olm = OlmPkDecryption::new();
let key = b"DEFAULT_PICKLE_KEY";
let pickle = olm.pickle(olm_rs::PicklingMode::Encrypted { key: key.to_vec() });
let unpickled = PkDecryption::from_libolm_pickle(&pickle, key)
.expect("We should be able to unpickle a key pickled by libolm");
assert_eq!(
olm.public_key(),
unpickled.public_key().to_base64(),
"The public keys of libolm and vodozemac should match"
);
}
#[test]
fn libolm_pickle_cycle() {
let olm = OlmPkDecryption::new();
let key = b"DEFAULT_PICKLE_KEY";
let pickle = olm.pickle(olm_rs::PicklingMode::Encrypted { key: key.to_vec() });
let decrypt = PkDecryption::from_libolm_pickle(&pickle, key)
.expect("We should be able to unpickle a key pickled by libolm");
let vodozemac_pickle =
decrypt.to_libolm_pickle(key).expect("We should be able to pickle a key");
let _ = PkDecryption::from_libolm_pickle(&vodozemac_pickle, key)
.expect("We should be able to unpickle a key pickled by vodozemac");
let unpickled = OlmPkDecryption::unpickle(
vodozemac_pickle,
olm_rs::PicklingMode::Encrypted { key: key.to_vec() },
)
.expect("Libolm should be able to unpickle a key pickled by vodozemac");
assert_eq!(
olm.public_key(),
unpickled.public_key(),
"The public keys of the restored and original libolm PK decryption should match"
);
}
}