malice

//! RFC 8291 message encryption with the aes128gcm content coding of RFC 8188,

//! as a single record.

use aws_lc_rs::{

    aead::{AES_128_GCM, Aad, LessSafeKey, Nonce, UnboundKey},

    agreement::{self, ECDH_P256, PrivateKey, UnparsedPublicKey},

    hkdf::{HKDF_SHA256, KeyType, Salt},

};

const RECORD_SIZE: u32 = 4096;

const TAG_LEN: usize = 16;

const HEADER_LEN: usize = 16 + 4 + 1 + 65;

/// Push services cap the encrypted body at 4096 bytes.

pub const MAX_PLAINTEXT: usize = 4096 - HEADER_LEN - TAG_LEN - 1;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]

pub enum EceError {

    BadKey,

    TooLong,

    Crypto,

}

impl std::fmt::Display for EceError {

    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

        f.write_str(match self {

            Self::BadKey => "the subscription's keys are not valid",

            Self::TooLong => "the push payload is too long",

            Self::Crypto => "push encryption failed",

        })

    }

}

impl std::error::Error for EceError {}

struct Len(usize);

impl KeyType for Len {

    fn len(&self) -> usize {

        self.0

    }

}

fn expand(salt: &[u8], secret: &[u8], info: &[&[u8]], out: &mut [u8]) -> Result<(), EceError> {

    Salt::new(HKDF_SHA256, salt)

        .extract(secret)

        .expand(info, Len(out.len()))

        .and_then(|okm| okm.fill(out))

        .map_err(|_| EceError::Crypto)

}

struct Keys {

    cek: [u8; 16],

    nonce: [u8; 12],

}

fn derive(

    ecdh_secret: &[u8],

    auth: &[u8],

    ua_public: &[u8],

    as_public: &[u8],

    salt: &[u8],

) -> Result<Keys, EceError> {

    let mut ikm = [0u8; 32];

    expand(

        auth,

        ecdh_secret,

        &[b"WebPush: info\0", ua_public, as_public],

        &mut ikm,

    )?;

    let mut keys = Keys {

        cek: [0; 16],

        nonce: [0; 12],

    };

    expand(

        salt,

        &ikm,

        &[b"Content-Encoding: aes128gcm\0"],

        &mut keys.cek,

    )?;

    expand(salt, &ikm, &[b"Content-Encoding: nonce\0"], &mut keys.nonce)?;

    Ok(keys)

}

fn ecdh(private: &PrivateKey, public: &[u8]) -> Result<Vec<u8>, EceError> {

    agreement::agree(

        private,

        UnparsedPublicKey::new(&ECDH_P256, public),

        EceError::BadKey,

        |secret| Ok(secret.to_vec()),

    )

}

/// Encrypts `plaintext` for the subscription whose `p256dh` key is

/// `ua_public` (uncompressed, 65 bytes) and whose auth secret is `auth`.

pub fn encrypt(ua_public: &[u8], auth: &[u8], plaintext: &[u8]) -> Result<Vec<u8>, EceError> {

    let as_private = PrivateKey::generate(&ECDH_P256).map_err(|_| EceError::Crypto)?;

    let mut salt = [0u8; 16];

    aws_lc_rs::rand::fill(&mut salt).map_err(|_| EceError::Crypto)?;

    encrypt_with(&as_private, salt, ua_public, auth, plaintext)

}

fn encrypt_with(

    as_private: &PrivateKey,

    salt: [u8; 16],

    ua_public: &[u8],

    auth: &[u8],

    plaintext: &[u8],

) -> Result<Vec<u8>, EceError> {

    if ua_public.len() != 65 || auth.len() != 16 {

        return Err(EceError::BadKey);

    }

    if plaintext.len() > MAX_PLAINTEXT {

        return Err(EceError::TooLong);

    }

    let as_public = as_private

        .compute_public_key()

        .map_err(|_| EceError::Crypto)?;

    let as_public = as_public.as_ref();

    let secret = ecdh(as_private, ua_public)?;

    let keys = derive(&secret, auth, ua_public, as_public, &salt)?;

    let mut out = Vec::with_capacity(HEADER_LEN + plaintext.len() + 1 + TAG_LEN);

    out.extend_from_slice(&salt);

    out.extend_from_slice(&RECORD_SIZE.to_be_bytes());

    out.push(as_public.len() as u8);

    out.extend_from_slice(as_public);

    let mut record = Vec::with_capacity(plaintext.len() + 1 + TAG_LEN);

    record.extend_from_slice(plaintext);

    record.push(2);

    LessSafeKey::new(UnboundKey::new(&AES_128_GCM, &keys.cek).map_err(|_| EceError::Crypto)?)

        .seal_in_place_append_tag(

            Nonce::assume_unique_for_key(keys.nonce),

            Aad::empty(),

            &mut record,

        )

        .map_err(|_| EceError::Crypto)?;

    out.extend_from_slice(&record);

    Ok(out)

}

/// The user agent's side, for tests and the mock push service.

pub fn decrypt(ua_private: &PrivateKey, auth: &[u8], body: &[u8]) -> Result<Vec<u8>, EceError> {

    if body.len() < 21 {

        return Err(EceError::Crypto);

    }

    let salt = &body[..16];

    let id_len = body[20] as usize;

    let as_public = body.get(21..21 + id_len).ok_or(EceError::Crypto)?;

    let mut record = body[21 + id_len..].to_vec();

    let ua_public = ua_private

        .compute_public_key()

        .map_err(|_| EceError::Crypto)?;

    let secret = ecdh(ua_private, as_public)?;

    let keys = derive(&secret, auth, ua_public.as_ref(), as_public, salt)?;

    let plain =

        LessSafeKey::new(UnboundKey::new(&AES_128_GCM, &keys.cek).map_err(|_| EceError::Crypto)?)

            .open_in_place(

                Nonce::assume_unique_for_key(keys.nonce),

                Aad::empty(),

                &mut record,

            )

            .map_err(|_| EceError::Crypto)?;

    match plain.iter().rposition(|&byte| byte != 0) {

        Some(end) if plain[end] == 2 => Ok(plain[..end].to_vec()),

        _ => Err(EceError::Crypto),

    }

}

#[cfg(test)]

mod tests {

    use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};

    use super::*;

    fn b64(value: &str) -> Vec<u8> {

        URL_SAFE_NO_PAD.decode(value).unwrap()

    }

    // RFC 8291 section 5 and appendix A

    const PLAINTEXT: &str = "V2hlbiBJIGdyb3cgdXAsIEkgd2FudCB0byBiZSBhIHdhdGVybWVsb24";

    const AS_PRIVATE: &str = "yfWPiYE-n46HLnH0KqZOF1fJJU3MYrct3AELtAQ-oRw";

    const UA_PUBLIC: &str =

        "BCVxsr7N_eNgVRqvHtD0zTZsEc6-VV-JvLexhqUzORcxaOzi6-AYWXvTBHm4bjyPjs7Vd8pZGH6SRpkNtoIAiw4";

    const UA_PRIVATE: &str = "q1dXpw3UpT5VOmu_cf_v6ih07Aems3njxI-JWgLcM94";

    const SALT: &str = "DGv6ra1nlYgDCS1FRnbzlw";

    const AUTH: &str = "BTBZMqHH6r4Tts7J_aSIgg";

    const MESSAGE: &str = "DGv6ra1nlYgDCS1FRnbzlwAAEABBBP4z9KsN6nGRTbVYI_c7VJSPQTBtkgcy27mlmlMoZIIgDll6e3vCYLocInmYWAmS6TlzAC8wEqKK6PBru3jl7A_yl95bQpu6cVPTpK4Mqgkf1CXztLVBSt2Ks3oZwbuwXPXLWyouBWLVWGNWQexSgSxsj_Qulcy4a-fN";

    #[test]

    fn rfc8291_vector_encrypts_exactly() {

        let as_private = PrivateKey::from_private_key(&ECDH_P256, &b64(AS_PRIVATE)).unwrap();

        let salt: [u8; 16] = b64(SALT).try_into().unwrap();

        let body = encrypt_with(

            &as_private,

            salt,

            &b64(UA_PUBLIC),

            &b64(AUTH),

            &b64(PLAINTEXT),

        )

        .unwrap();

        assert_eq!(URL_SAFE_NO_PAD.encode(body), MESSAGE);

    }

    #[test]

    fn rfc8291_vector_decrypts() {

        let ua_private = PrivateKey::from_private_key(&ECDH_P256, &b64(UA_PRIVATE)).unwrap();

        let plain = decrypt(&ua_private, &b64(AUTH), &b64(MESSAGE)).unwrap();

        assert_eq!(plain, b64(PLAINTEXT));

    }

    #[test]

    fn random_keys_round_trip_and_bad_input_is_refused() {

        let ua_private = PrivateKey::generate(&ECDH_P256).unwrap();

        let ua_public = ua_private.compute_public_key().unwrap();

        let auth = [9u8; 16];

        let body = encrypt(ua_public.as_ref(), &auth, b"hello").unwrap();

        assert_eq!(decrypt(&ua_private, &auth, &body).unwrap(), b"hello");

        assert!(decrypt(&ua_private, &[1u8; 16], &body).is_err());

        assert_eq!(encrypt(&[4u8; 65], &auth, b"x"), Err(EceError::BadKey));

        assert_eq!(

            encrypt(ua_public.as_ref(), &auth, &vec![0; MAX_PLAINTEXT + 1]),

            Err(EceError::TooLong)

        );

    }

}