malice

use std::{

    io::{ErrorKind, Write},

    path::Path,

    time::{SystemTime, UNIX_EPOCH},

};

use aws_lc_rs::{

    rand::SystemRandom,

    signature::{ECDSA_P256_SHA256_FIXED_SIGNING, EcdsaKeyPair, KeyPair},

};

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

const TOKEN_LIFETIME_SECS: u64 = 12 * 60 * 60;

/// The server's VAPID identity (RFC 8292), an ES256 key kept on disk.

pub struct Vapid {

    key: EcdsaKeyPair,

    public: String,

}

impl Vapid {

    /// Reads the PKCS#8 key at `path`, or creates it owner-only.

    pub fn load_or_create(path: &Path) -> Result<Self, String> {

        match std::fs::read(path) {

            Ok(bytes) => return Self::from_pkcs8(&bytes),

            Err(error) if error.kind() != ErrorKind::NotFound => {

                return Err(format!("reading the VAPID key: {error}"));

            }

            Err(_) => {}

        }

        let document =

            EcdsaKeyPair::generate_pkcs8(&ECDSA_P256_SHA256_FIXED_SIGNING, &SystemRandom::new())

                .map_err(|_| "generating the VAPID key failed".to_owned())?;

        let mut options = std::fs::OpenOptions::new();

        options.write(true).create_new(true);

        #[cfg(unix)]

        {

            use std::os::unix::fs::OpenOptionsExt;

            options.mode(0o600);

        }

        match options.open(path) {

            Ok(mut file) => {

                if let Err(error) = file

                    .write_all(document.as_ref())

                    .and_then(|_| file.sync_all())

                {

                    let _ = std::fs::remove_file(path);

                    return Err(format!("writing the VAPID key: {error}"));

                }

                Self::from_pkcs8(document.as_ref())

            }

            Err(error) if error.kind() == ErrorKind::AlreadyExists => {

                let bytes = std::fs::read(path)

                    .map_err(|error| format!("reading the VAPID key: {error}"))?;

                Self::from_pkcs8(&bytes)

            }

            Err(error) => Err(format!("creating the VAPID key: {error}")),

        }

    }

    fn from_pkcs8(bytes: &[u8]) -> Result<Self, String> {

        let key = EcdsaKeyPair::from_pkcs8(&ECDSA_P256_SHA256_FIXED_SIGNING, bytes)

            .map_err(|_| "the VAPID key file is not a P-256 PKCS#8 key".to_owned())?;

        let public = URL_SAFE_NO_PAD.encode(key.public_key().as_ref());

        Ok(Self { key, public })

    }

    /// The uncompressed public point, base64url, as `applicationServerKey` takes it.

    pub fn public_key(&self) -> &str {

        &self.public

    }

    /// The `Authorization` header for a push service at `audience` (its origin).

    pub fn authorization(&self, audience: &str, contact: &str) -> Result<String, String> {

        let now = SystemTime::now()

            .duration_since(UNIX_EPOCH)

            .map_err(|_| "the clock is before 1970".to_owned())?

            .as_secs();

        let header = URL_SAFE_NO_PAD.encode(br#"{"typ":"JWT","alg":"ES256"}"#);

        let claims = serde_json::json!({

            "aud": audience,

            "exp": now + TOKEN_LIFETIME_SECS,

            "sub": contact,

        });

        let claims = URL_SAFE_NO_PAD.encode(claims.to_string());

        let signing_input = format!("{header}.{claims}");

        let signature = self

            .key

            .sign(&SystemRandom::new(), signing_input.as_bytes())

            .map_err(|_| "signing the VAPID token failed".to_owned())?;

        Ok(format!(

            "vapid t={signing_input}.{}, k={}",

            URL_SAFE_NO_PAD.encode(signature.as_ref()),

            self.public

        ))

    }

}

#[cfg(test)]

mod tests {

    use aws_lc_rs::signature::{ECDSA_P256_SHA256_FIXED, UnparsedPublicKey};

    use super::*;

    #[test]

    fn key_persists_owner_only_and_signs_verifiable_tokens() {

        let dir = tempfile::tempdir().unwrap();

        let path = dir.path().join("vapid.pk8");

        let first = Vapid::load_or_create(&path).unwrap();

        let again = Vapid::load_or_create(&path).unwrap();

        assert_eq!(first.public_key(), again.public_key());

        let public = URL_SAFE_NO_PAD.decode(first.public_key()).unwrap();

        assert_eq!(public.len(), 65);

        assert_eq!(public[0], 4);

        #[cfg(unix)]

        {

            use std::os::unix::fs::PermissionsExt;

            let mode = std::fs::metadata(&path).unwrap().permissions().mode();

            assert_eq!(mode & 0o777, 0o600);

        }

        let header = first

            .authorization("https://fcm.googleapis.com", "mailto:a@example.org")

            .unwrap();

        let (token, key) = header

            .strip_prefix("vapid t=")

            .unwrap()

            .split_once(", k=")

            .unwrap();

        assert_eq!(key, first.public_key());

        let (input, signature) = token.rsplit_once('.').unwrap();

        UnparsedPublicKey::new(&ECDSA_P256_SHA256_FIXED, &public)

            .verify(

                input.as_bytes(),

                &URL_SAFE_NO_PAD.decode(signature).unwrap(),

            )

            .unwrap();

        let claims: serde_json::Value = serde_json::from_slice(

            &URL_SAFE_NO_PAD

                .decode(input.split_once('.').unwrap().1)

                .unwrap(),

        )

        .unwrap();

        assert_eq!(claims["aud"], "https://fcm.googleapis.com");

        assert_eq!(claims["sub"], "mailto:a@example.org");

    }

}