use alloc::vec::Vec;
use aes_gcm::Aes256Gcm;
use aes_gcm::aead::{Aead, KeyInit, Payload};
use chacha20poly1305::ChaCha20Poly1305 as ChaChaCipher;
use x25519_dalek::{PublicKey, StaticSecret};
use zeroize::Zeroizing;
use crate::alg::ContentAlgorithm;
use crate::claims::Claims;
use crate::codec::{self, ClaimsExpectation, NONCE_LEN};
use crate::error::{BuildError, DecodeError, OpenError};
use crate::kdf::{self, KdfParties};
use crate::keys::X25519RecipientPublic;
use crate::traits::{OpenRequest, Recipient};
use crate::types::{ContentType, CoseBytes, ExternalAad, KeyId};
#[derive(Debug, Clone)]
pub struct EncryptParams<'a> {
pub content_type: ContentType,
pub plaintext: &'a [u8],
pub recipient: X25519RecipientPublic,
pub content_algorithm: ContentAlgorithm,
pub external_aad: ExternalAad,
pub kdf_parties: KdfParties,
}
#[cfg(feature = "fixtures")]
#[derive(Debug)]
pub struct SealParts {
pub ephemeral_private: Zeroizing<[u8; 32]>,
pub nonce: [u8; NONCE_LEN],
}
pub fn random_array<const N: usize>() -> Result<Zeroizing<[u8; N]>, BuildError> {
let mut buf = Zeroizing::new([0u8; N]);
getrandom::fill(buf.as_mut_slice()).map_err(|_| BuildError::Rng)?;
Ok(buf)
}
pub fn aead_seal(
alg: ContentAlgorithm,
key: &Zeroizing<[u8; 32]>,
nonce: &[u8; NONCE_LEN],
plaintext: &[u8],
aad: &[u8],
) -> Result<Vec<u8>, BuildError> {
let payload = Payload {
msg: plaintext,
aad,
};
match alg {
ContentAlgorithm::A256Gcm => Aes256Gcm::new_from_slice(key.as_slice())
.map_err(|_| BuildError::SealFailed)?
.encrypt(aes_gcm::Nonce::from_slice(nonce), payload)
.map_err(|_| BuildError::SealFailed),
ContentAlgorithm::ChaCha20Poly1305 => ChaChaCipher::new_from_slice(key.as_slice())
.map_err(|_| BuildError::SealFailed)?
.encrypt(chacha20poly1305::Nonce::from_slice(nonce), payload)
.map_err(|_| BuildError::SealFailed),
}
}
pub fn aead_open(
alg: ContentAlgorithm,
key: &Zeroizing<[u8; 32]>,
nonce: &[u8; NONCE_LEN],
ciphertext: &[u8],
aad: &[u8],
) -> Result<Zeroizing<Vec<u8>>, OpenError> {
let payload = Payload {
msg: ciphertext,
aad,
};
let plaintext = match alg {
ContentAlgorithm::A256Gcm => Aes256Gcm::new_from_slice(key.as_slice())
.map_err(|_| OpenError::OpenFailed)?
.decrypt(aes_gcm::Nonce::from_slice(nonce), payload)
.map_err(|_| OpenError::OpenFailed)?,
ContentAlgorithm::ChaCha20Poly1305 => ChaChaCipher::new_from_slice(key.as_slice())
.map_err(|_| OpenError::OpenFailed)?
.decrypt(chacha20poly1305::Nonce::from_slice(nonce), payload)
.map_err(|_| OpenError::OpenFailed)?,
};
Ok(Zeroizing::new(plaintext))
}
pub struct EncryptCore<'a> {
pub content_algorithm: ContentAlgorithm,
pub content_type: &'a ContentType,
pub claims: Option<&'a Claims>,
pub plaintext: &'a [u8],
pub recipient: &'a X25519RecipientPublic,
pub external_aad: &'a ExternalAad,
pub kdf_parties: &'a KdfParties,
}
pub fn build_encrypt_core(
core: &EncryptCore<'_>,
ephemeral_private: &Zeroizing<[u8; 32]>,
nonce: [u8; NONCE_LEN],
) -> Result<Vec<u8>, BuildError> {
let ephemeral_secret = StaticSecret::from(**ephemeral_private);
let ephemeral_pub = PublicKey::from(&ephemeral_secret).to_bytes();
let recipient_pub = PublicKey::from(core.recipient.public);
let shared_secret = ephemeral_secret.diffie_hellman(&recipient_pub);
if !shared_secret.was_contributory() {
return Err(BuildError::SealFailed);
}
let shared = Zeroizing::new(shared_secret.to_bytes());
let recipient_protected = codec::encode_recipient_protected(core.kdf_parties)
.map_err(|codec::CodecError| BuildError::Codec)?;
let info = codec::kdf_context(
core.content_algorithm,
core.kdf_parties,
&recipient_protected,
)
.map_err(|codec::CodecError| BuildError::Codec)?;
let cek = kdf::derive_cek(&shared, &info).map_err(|kdf::KdfFailed| BuildError::SealFailed)?;
let protected =
codec::encode_encrypt_protected(core.content_algorithm, core.content_type, core.claims)
.map_err(|codec::CodecError| BuildError::Codec)?;
let aad = codec::enc_structure(&protected, core.external_aad.as_bytes())
.map_err(|codec::CodecError| BuildError::Codec)?;
let ciphertext = aead_seal(core.content_algorithm, &cek, &nonce, core.plaintext, &aad)?;
codec::assemble_encrypt(&codec::EncryptAssembly {
protected: &protected,
iv: &nonce,
ciphertext: &ciphertext,
recipient_protected: &recipient_protected,
recipient_kid: &core.recipient.key_id,
ephemeral_x: &ephemeral_pub,
})
.map_err(|codec::CodecError| BuildError::Codec)
}
impl EncryptParams<'_> {
const fn core(&self) -> EncryptCore<'_> {
EncryptCore {
content_algorithm: self.content_algorithm,
content_type: &self.content_type,
claims: None,
plaintext: self.plaintext,
recipient: &self.recipient,
external_aad: &self.external_aad,
kdf_parties: &self.kdf_parties,
}
}
}
pub fn build_encrypted(params: &EncryptParams<'_>) -> Result<CoseBytes, BuildError> {
let ephemeral = random_array::<32>()?;
let nonce = random_array::<NONCE_LEN>()?;
build_encrypt_core(¶ms.core(), &ephemeral, *nonce).map(CoseBytes::new)
}
#[cfg(feature = "fixtures")]
pub fn build_encrypted_with_parts(
params: &EncryptParams<'_>,
parts: &SealParts,
) -> Result<CoseBytes, BuildError> {
build_encrypt_core(¶ms.core(), &parts.ephemeral_private, parts.nonce).map(CoseBytes::new)
}
#[derive(Debug, Clone)]
pub struct EncryptedMessage {
pub content_type: ContentType,
pub recipient_key_id: KeyId,
pub content_algorithm: ContentAlgorithm,
pub parties: KdfParties,
bytes: Vec<u8>,
}
pub fn decode_encrypted(bytes: &[u8]) -> Result<EncryptedMessage, DecodeError> {
let decoded = codec::decode_encrypt_strict(bytes, ClaimsExpectation::Forbidden)?;
Ok(EncryptedMessage {
content_type: decoded.content_type,
recipient_key_id: decoded.recipient_kid,
content_algorithm: decoded.content_algorithm,
parties: decoded.parties,
bytes: bytes.to_vec(),
})
}
#[derive(Debug)]
pub struct Opened {
pub plaintext: Zeroizing<Vec<u8>>,
pub content_type: ContentType,
}
impl EncryptedMessage {
pub async fn open<R: Recipient>(
&self,
recipient: &R,
aad: &ExternalAad,
expected_parties: Option<&KdfParties>,
) -> Result<Opened, OpenError> {
if recipient.key_id() != &self.recipient_key_id {
return Err(OpenError::RecipientKeyMismatch);
}
if let Some(expected) = expected_parties
&& *expected != self.parties
{
return Err(OpenError::PartyMismatch);
}
let request = OpenRequest {
cose_encrypt: &self.bytes,
external_aad: aad,
expected_parties,
};
let plaintext = recipient.open(&request).await?;
Ok(Opened {
plaintext,
content_type: self.content_type.clone(),
})
}
}