use core::fmt;
use bech32::primitives::decode::CheckedHrpstring;
use chacha20poly1305::{
aead::{self, generic_array::typenum::Unsigned, Aead, AeadCore, KeyInit},
ChaCha20Poly1305,
};
use hkdf::Hkdf;
use sha2::Sha256;
pub fn aead_encrypt(key: &[u8; 32], plaintext: &[u8]) -> Vec<u8> {
let c = ChaCha20Poly1305::new(key.into());
c.encrypt(&[0; 12].into(), plaintext)
.expect("we won't overflow the ChaCha20 block counter")
}
pub fn aead_decrypt(
key: &[u8; 32],
size: usize,
ciphertext: &[u8],
) -> Result<Vec<u8>, aead::Error> {
if ciphertext.len() != size + <ChaCha20Poly1305 as AeadCore>::TagSize::to_usize() {
return Err(aead::Error);
}
let c = ChaCha20Poly1305::new(key.into());
c.decrypt(&[0; 12].into(), ciphertext)
}
pub fn hkdf(salt: &[u8], label: &[u8], ikm: &[u8]) -> [u8; 32] {
let mut okm = [0; 32];
Hkdf::<Sha256>::new(Some(salt), ikm)
.expand(label, &mut okm)
.expect("okm is the correct length");
okm
}
#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
enum Bech32Long {}
impl bech32::Checksum for Bech32Long {
type MidstateRepr = u32;
const CODE_LENGTH: usize = usize::MAX;
const CHECKSUM_LENGTH: usize = bech32::Bech32::CHECKSUM_LENGTH;
const GENERATOR_SH: [u32; 5] = bech32::Bech32::GENERATOR_SH;
const TARGET_RESIDUE: u32 = bech32::Bech32::TARGET_RESIDUE;
}
pub fn bech32_encode(hrp: bech32::Hrp, data: &[u8]) -> String {
bech32::encode_lower::<Bech32Long>(hrp, data).expect("we don't enforce the Bech32 length limit")
}
pub fn bech32_encode_to_fmt(f: &mut impl fmt::Write, hrp: bech32::Hrp, data: &[u8]) -> fmt::Result {
bech32::encode_lower_to_fmt::<Bech32Long, _>(f, hrp, data).map_err(|e| match e {
bech32::EncodeError::Fmt(error) => error,
bech32::EncodeError::TooLong(_) => unreachable!("we don't enforce the Bech32 length limit"),
_ => panic!("Unexpected error: {e}"),
})
}
pub fn bech32_decode<E, F, G, H, T>(
s: &str,
parse_err: F,
hrp_filter: G,
data_parse: H,
) -> Result<T, E>
where
F: FnOnce(bech32::primitives::decode::CheckedHrpstringError) -> E,
G: FnOnce(bech32::Hrp) -> Result<(), E>,
H: FnOnce(bech32::Hrp, bech32::primitives::decode::ByteIter) -> Result<T, E>,
{
CheckedHrpstring::new::<Bech32Long>(s)
.map_err(parse_err)
.and_then(|parsed| {
hrp_filter(parsed.hrp()).and_then(|()| data_parse(parsed.hrp(), parsed.byte_iter()))
})
}
pub fn hpke_seal<Kem: hpke::Kem, R: rand::RngCore + rand::CryptoRng>(
pk_recip: &Kem::PublicKey,
info: &[u8],
plaintext: &[u8],
rng: &mut R,
) -> (Kem::EncappedKey, Vec<u8>) {
hpke::single_shot_seal::<hpke::aead::ChaCha20Poly1305, hpke::kdf::HkdfSha256, Kem, R>(
&hpke::OpModeS::Base,
pk_recip,
info,
plaintext,
&[],
rng,
)
.expect("no errors should occur with these HPKE parameters")
}
pub fn hpke_open<Kem: hpke::Kem>(
encapped_key: &Kem::EncappedKey,
sk_recip: &Kem::PrivateKey,
info: &[u8],
ciphertext: &[u8],
) -> Result<Vec<u8>, hpke::HpkeError> {
hpke::single_shot_open::<hpke::aead::ChaCha20Poly1305, hpke::kdf::HkdfSha256, Kem>(
&hpke::OpModeR::Base,
sk_recip,
encapped_key,
info,
ciphertext,
&[],
)
}
#[cfg(test)]
mod tests {
use hpke::Kem;
use rand::rngs::OsRng;
use super::{aead_decrypt, aead_encrypt, bech32_decode, bech32_encode, hpke_open, hpke_seal};
#[test]
fn aead_round_trip() {
let key = [14; 32];
let plaintext = b"12345678";
let encrypted = aead_encrypt(&key, plaintext);
let decrypted = aead_decrypt(&key, plaintext.len(), &encrypted).unwrap();
assert_eq!(decrypted, plaintext);
}
#[test]
fn bech32_round_trip() {
let hrp = bech32::Hrp::parse_unchecked("12345678");
let data = [14; 32];
let encoded = bech32_encode(hrp, &data);
let decoded = bech32_decode(
&encoded,
|_| (),
|parsed_hrp| (parsed_hrp == hrp).then_some(()).ok_or(()),
|_, bytes| Ok(bytes.collect::<Vec<_>>()),
)
.unwrap();
assert_eq!(decoded, data);
}
#[test]
fn hpke_round_trip() {
type Kem = hpke::kem::DhP256HkdfSha256;
let mut rng = OsRng;
let (sk_recip, pk_recip) = Kem::gen_keypair(&mut rng);
let info = b"foobar";
let plaintext = b"12345678";
let (encapped_key, ciphertext) = hpke_seal::<Kem, _>(&pk_recip, info, plaintext, &mut rng);
let decrypted = hpke_open::<Kem>(&encapped_key, &sk_recip, info, &ciphertext).unwrap();
assert_eq!(decrypted, plaintext);
}
}