malice

use std::{sync::Arc, time::Duration};

use aws_lc_rs::digest::{SHA256, digest};

use base64::{

    Engine,

    engine::general_purpose::{STANDARD_NO_PAD, URL_SAFE_NO_PAD},

};

use reqwest::{StatusCode, header};

use serde::{Deserialize, Serialize};

use url::Url;

use crate::{

    ece::{self, EceError},

    vapid::Vapid,

};

const PUSH_HOSTS: &[&str] = &[

    "*.push.apple.com",

    "fcm.googleapis.com",

    "updates.push.services.mozilla.com",

    "*.notify.windows.com",

];

const ALLOW_HOSTS_ENV: &str = "MALICE_PUSH_ALLOW_HOSTS";

const MAX_ENDPOINT_LEN: usize = 2048;

const TTL_SECS: u32 = 12 * 60 * 60;

/// A browser's `PushSubscription`, keys as the browser's `toJSON` gives them.

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]

pub struct Subscription {

    pub endpoint: String,

    pub p256dh: String,

    pub auth: String,

}

impl Subscription {

    /// Whether the keys decode to a P-256 point and a 16-byte secret.

    pub fn keys_valid(&self) -> bool {

        matches!(

            (decode_key(&self.p256dh), decode_key(&self.auth)),

            (Some(p256dh), Some(auth)) if p256dh.len() == 65 && p256dh[0] == 4 && auth.len() == 16

        )

    }

}

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

    let value = value.trim().trim_end_matches('=');

    URL_SAFE_NO_PAD

        .decode(value)

        .or_else(|_| STANDARD_NO_PAD.decode(value))

        .ok()

}

/// Push service hosts this server will POST to, so a pusher can't aim it at

/// an arbitrary URL.

#[derive(Debug, Clone, Default)]

pub struct Allowlist {

    /// `host:port` authorities also allowed over plain http, for local testing

    extra: Vec<String>,

}

impl Allowlist {

    /// The built-in push services, plus any `host:port` entries in

    /// `MALICE_PUSH_ALLOW_HOSTS`.

    pub fn from_env() -> Self {

        let extra = std::env::var(ALLOW_HOSTS_ENV)

            .unwrap_or_default()

            .split(',')

            .map(|host| host.trim().to_ascii_lowercase())

            .filter(|host| !host.is_empty())

            .collect::<Vec<_>>();

        if !extra.is_empty() {

            tracing::warn!(hosts = ?extra, "{ALLOW_HOSTS_ENV} allows extra push hosts, for testing only");

        }

        Self { extra }

    }

    pub fn with_extra(extra: impl IntoIterator<Item = String>) -> Self {

        Self {

            extra: extra.into_iter().collect(),

        }

    }

    pub fn check(&self, endpoint: &str) -> Option<Url> {

        if endpoint.len() > MAX_ENDPOINT_LEN {

            return None;

        }

        let url = Url::parse(endpoint).ok()?;

        if !url.username().is_empty() || url.password().is_some() {

            return None;

        }

        let port = url.port_or_known_default()?;

        let builtin = url.scheme() == "https"

            && port == 443

            && url.domain().is_some_and(|host| {

                PUSH_HOSTS

                    .iter()

                    .any(|pattern| match pattern.strip_prefix('*') {

                        Some(suffix) => host.len() > suffix.len() && host.ends_with(suffix),

                        None => host == *pattern,

                    })

            });

        let extra = matches!(url.scheme(), "http" | "https")

            && self.extra.contains(&format!("{}:{port}", url.host_str()?));

        (builtin || extra).then_some(url)

    }

}

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

pub enum Urgency {

    Normal,

    High,

}

pub struct Message<'a> {

    pub payload: &'a [u8],

    pub urgency: Urgency,

    /// pushes sharing a topic replace each other while still queued

    pub topic: &'a str,

}

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

pub enum Outcome {

    Delivered,

    /// the push service no longer knows the subscription (404 or 410)

    Gone,

    /// the endpoint is not an allowed push service, or the keys are unusable

    Refused,

    Retryable,

    Failed(String),

}

impl Outcome {

    pub fn drops_subscription(&self) -> bool {

        matches!(self, Self::Gone | Self::Refused)

    }

}

/// Encrypts and delivers pushes, signed with this server's VAPID key.

#[derive(Clone)]

pub struct Sender {

    http: reqwest::Client,

    vapid: Arc<Vapid>,

    contact: String,

    allow: Allowlist,

}

impl Sender {

    /// `contact` is the VAPID `sub`, a `mailto:` or `https:` URL.

    pub fn new(vapid: Vapid, contact: String, allow: Allowlist) -> Result<Self, String> {

        let http = reqwest::Client::builder()

            .redirect(reqwest::redirect::Policy::none())

            .connect_timeout(Duration::from_secs(5))

            .timeout(Duration::from_secs(10))

            .build()

            .map_err(|error| format!("building the push client: {error}"))?;

        Ok(Self {

            http,

            vapid: Arc::new(vapid),

            contact,

            allow,

        })

    }

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

        self.vapid.public_key()

    }

    pub fn allows(&self, endpoint: &str) -> bool {

        self.allow.check(endpoint).is_some()

    }

    /// Sends once, and once more after a second if the push service asks to retry.

    pub async fn send(&self, subscription: &Subscription, message: &Message<'_>) -> Outcome {

        let Some(url) = self.allow.check(&subscription.endpoint) else {

            return Outcome::Refused;

        };

        let (Some(p256dh), Some(auth)) = (

            decode_key(&subscription.p256dh),

            decode_key(&subscription.auth),

        ) else {

            return Outcome::Refused;

        };

        let body = match ece::encrypt(&p256dh, &auth, message.payload) {

            Ok(body) => body,

            Err(EceError::BadKey) => return Outcome::Refused,

            Err(error) => return Outcome::Failed(error.to_string()),

        };

        let authorization = match self

            .vapid

            .authorization(&url.origin().ascii_serialization(), &self.contact)

        {

            Ok(authorization) => authorization,

            Err(error) => return Outcome::Failed(error),

        };

        let topic = topic(message.topic);

        let urgency = match message.urgency {

            Urgency::Normal => "normal",

            Urgency::High => "high",

        };

        for attempt in 0..2 {

            if attempt > 0 {

                tokio::time::sleep(Duration::from_secs(1)).await;

            }

            let response = self

                .http

                .post(url.clone())

                .header(header::AUTHORIZATION, &authorization)

                .header(header::CONTENT_ENCODING, "aes128gcm")

                .header(header::CONTENT_TYPE, "application/octet-stream")

                .header("TTL", TTL_SECS)

                .header("Urgency", urgency)

                .header("Topic", &topic)

                .body(body.clone())

                .send()

                .await;

            let outcome = match response {

                Ok(response) => classify(response.status()),

                Err(error) => {

                    tracing::debug!(%error, "push service unreachable");

                    Outcome::Retryable

                }

            };

            if outcome != Outcome::Retryable {

                return outcome;

            }

        }

        Outcome::Retryable

    }

}

fn classify(status: StatusCode) -> Outcome {

    match status {

        status if status.is_success() => Outcome::Delivered,

        StatusCode::NOT_FOUND | StatusCode::GONE => Outcome::Gone,

        StatusCode::TOO_MANY_REQUESTS => Outcome::Retryable,

        status if status.is_server_error() => Outcome::Retryable,

        status => Outcome::Failed(format!("push service answered {status}")),

    }

}

/// RFC 8030 limits a topic to 32 characters of the base64url alphabet.

fn topic(raw: &str) -> String {

    let mut topic = URL_SAFE_NO_PAD.encode(digest(&SHA256, raw.as_bytes()));

    topic.truncate(32);

    topic

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn allowlist_admits_push_services_only() {

        let allow = Allowlist::default();

        assert!(

            allow

                .check("https://fcm.googleapis.com/fcm/send/abc")

                .is_some()

        );

        assert!(allow.check("https://web.push.apple.com/QGx").is_some());

        assert!(

            allow

                .check("https://wns2-par02p.notify.windows.com/w/?token=x")

                .is_some()

        );

        assert!(

            allow

                .check("https://updates.push.services.mozilla.com/wpush/v2/x")

                .is_some()

        );

        assert!(allow.check("http://fcm.googleapis.com/x").is_none());

        assert!(allow.check("https://fcm.googleapis.com:8443/x").is_none());

        assert!(allow.check("https://push.apple.com/x").is_none());

        assert!(allow.check("https://evilpush.apple.com/x").is_none());

        assert!(

            allow

                .check("https://fcm.googleapis.com.evil.org/x")

                .is_none()

        );

        assert!(allow.check("https://user@fcm.googleapis.com/x").is_none());

        assert!(allow.check("https://127.0.0.1/x").is_none());

        assert!(allow.check("http://127.0.0.1:9000/x").is_none());

        let allow = Allowlist::with_extra(["localhost:9000".to_owned()]);

        assert!(allow.check("http://localhost:9000/x").is_some());

        assert!(allow.check("http://localhost:9001/x").is_none());

        let allow = Allowlist::with_extra(["127.0.0.1:9000".to_owned()]);

        assert!(allow.check("http://127.0.0.1:9000/x").is_some());

    }

    #[test]

    fn topics_fit_rfc8030() {

        let topic = topic("!room:example.org");

        assert_eq!(topic.len(), 32);

        assert!(

            topic

                .chars()

                .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')

        );

    }

    #[test]

    fn keys_validate() {

        let good = Subscription {

            endpoint: String::new(),

            p256dh: "BCVxsr7N_eNgVRqvHtD0zTZsEc6-VV-JvLexhqUzORcxaOzi6-AYWXvTBHm4bjyPjs7Vd8pZGH6SRpkNtoIAiw4".to_owned(),

            auth: "BTBZMqHH6r4Tts7J_aSIgg==".to_owned(),

        };

        assert!(good.keys_valid());

        assert!(

            !Subscription {

                auth: "short".to_owned(),

                ..good

            }

            .keys_valid()

        );

    }

}