malice

use std::{

    collections::HashMap,

    path::{Path, PathBuf},

    sync::{Mutex, MutexGuard, PoisonError},

    time::Duration,

};

use axum::{

    Router,

    body::Bytes,

    extract::{DefaultBodyLimit, State},

    http::{HeaderMap, StatusCode},

    response::{IntoResponse, Response},

    routing::{get, post},

};

use macltok::{HttpPusher, PushFormat, PushNotificationStatus, PusherKey};

use macltok_push_gateway::{GatewayDevice, GatewayNotification, notify_response, parse_notify};

use malice_core::push::{notice, unread_total};

use malice_push::{Message, Outcome, Sender, Subscription, Urgency};

use malice_static::push::{NOTIFY_PATH, json_response};

use rand_core::{OsRng, RngCore};

use serde::{Deserialize, Serialize};

use tokio::{sync::Semaphore, time::timeout};

use crate::{

    error::{ApiError, json_ok, no_content},

    routes::auth,

    state::{AppState, RestoreError},

};

pub const HOSTED_APP_ID: &str = "dev.malice.hosted";

const CONCURRENT_NOTIFIES: usize = 8;

const RESOLVE_TIMEOUT: Duration = Duration::from_secs(20);

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

struct Entry {

    sid: String,

    subscription: Subscription,

}

/// Browser subscriptions by the random pushkey each one's pusher carries,

/// kept in `push.json` beside the sessions.

pub struct HostedPush {

    sender: Sender,

    public_url: Option<String>,

    path: PathBuf,

    entries: Mutex<HashMap<String, Entry>>,

    permits: Semaphore,

}

impl HostedPush {

    pub fn new(data_dir: &Path, public_url: Option<String>, sender: Sender) -> Self {

        let path = data_dir.join("push.json");

        let entries = match std::fs::read(&path) {

            Ok(bytes) => serde_json::from_slice(&bytes).unwrap_or_else(|error| {

                tracing::warn!(%error, "push subscriptions are unreadable, starting empty");

                HashMap::new()

            }),

            Err(_) => HashMap::new(),

        };

        Self {

            sender,

            public_url,

            path,

            entries: Mutex::new(entries),

            permits: Semaphore::new(CONCURRENT_NOTIFIES),

        }

    }

    /// The VAPID key lives in `data_dir`. Subscribing needs `MALICE_PUBLIC_URL`,

    /// since the homeserver must reach the notify endpoint.

    pub fn from_env(data_dir: &Path) -> Result<Self, String> {

        let public_url = malice_push::env::public_url();

        if public_url.is_none() {

            tracing::info!("web push is off until MALICE_PUBLIC_URL is set");

        }

        let sender = malice_push::env::sender(data_dir, public_url.as_deref())?;

        Ok(Self::new(data_dir, public_url, sender))

    }

    fn entries(&self) -> MutexGuard<'_, HashMap<String, Entry>> {

        self.entries.lock().unwrap_or_else(PoisonError::into_inner)

    }

    fn save(&self, entries: &HashMap<String, Entry>) {

        let result = serde_json::to_vec(entries)

            .map_err(std::io::Error::other)

            .and_then(|bytes| {

                let temp = self.path.with_extension("json.tmp");

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

                options.write(true).create(true).truncate(true);

                #[cfg(unix)]

                {

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

                    options.mode(0o600);

                }

                let mut file = options.open(&temp)?;

                std::io::Write::write_all(&mut file, &bytes)?;

                file.sync_all()?;

                std::fs::rename(&temp, &self.path)

            });

        if let Err(error) = result {

            tracing::warn!(%error, "saving push subscriptions failed");

        }

    }

    /// Files `subscription` under `sid`, keeping the pushkey if this session

    /// already had it. Returns the pushkey, and the session and pushkey that

    /// held the endpoint before when another session did.

    fn claim(&self, sid: &str, subscription: Subscription) -> (String, Option<(String, String)>) {

        let mut entries = self.entries();

        let held = entries

            .iter()

            .find(|(_, entry)| entry.subscription.endpoint == subscription.endpoint)

            .map(|(pushkey, entry)| (pushkey.clone(), entry.sid.clone()));

        let (pushkey, displaced) = match held {

            Some((pushkey, owner)) if owner == sid => (pushkey, None),

            Some((pushkey, owner)) => {

                entries.remove(&pushkey);

                (new_pushkey(), Some((owner, pushkey)))

            }

            None => (new_pushkey(), None),

        };

        entries.insert(

            pushkey.clone(),

            Entry {

                sid: sid.to_owned(),

                subscription,

            },

        );

        self.save(&entries);

        (pushkey, displaced)

    }

    fn get(&self, pushkey: &str) -> Option<Entry> {

        self.entries().get(pushkey).cloned()

    }

    fn find(&self, sid: &str, endpoint: &str) -> Option<String> {

        self.entries()

            .iter()

            .find(|(_, entry)| entry.sid == sid && entry.subscription.endpoint == endpoint)

            .map(|(pushkey, _)| pushkey.clone())

    }

    fn forget(&self, pushkey: &str) {

        let mut entries = self.entries();

        if entries.remove(pushkey).is_some() {

            self.save(&entries);

        }

    }

    /// The pushers a logout of `sid` should take down with it.

    pub fn pusher_keys(&self, sid: &str) -> Vec<PusherKey> {

        self.entries()

            .iter()

            .filter(|(_, entry)| entry.sid == sid)

            .map(|(pushkey, _)| PusherKey {

                app_id: HOSTED_APP_ID.to_owned(),

                pushkey: pushkey.clone(),

            })

            .collect()

    }

    pub fn forget_session(&self, sid: &str) {

        let mut entries = self.entries();

        let before = entries.len();

        entries.retain(|_, entry| entry.sid != sid);

        if entries.len() != before {

            self.save(&entries);

        }

    }

}

fn new_pushkey() -> String {

    let mut bytes = [0u8; 16];

    OsRng.fill_bytes(&mut bytes);

    hex::encode(bytes)

}

#[derive(Debug, Deserialize)]

struct SubscribeKeys {

    p256dh: String,

    auth: String,

}

#[derive(Debug, Deserialize)]

struct SubscribeBody {

    endpoint: String,

    keys: SubscribeKeys,

}

#[derive(Debug, Deserialize)]

struct UnsubscribeBody {

    endpoint: String,

}

pub fn api_routes() -> Router<AppState> {

    Router::new()

        .route("/push/vapid", get(vapid))

        .route("/push/subscribe", post(subscribe))

        .route("/push/unsubscribe", post(unsubscribe))

}

pub fn notify_route() -> Router<AppState> {

    Router::new()

        .route(NOTIFY_PATH, post(notify))

        .layer(DefaultBodyLimit::max(64 * 1024))

}

fn push_off() -> ApiError {

    ApiError::new(StatusCode::SERVICE_UNAVAILABLE, "web push is not set up")

}

/// The push service, the public URL to point pushers at, or why push is off.

fn configured(state: &AppState) -> Result<(&HostedPush, &str), ApiError> {

    let push = state.push().ok_or_else(push_off)?;

    let url = push.public_url.as_deref().ok_or_else(push_off)?;

    Ok((push, url))

}

async fn vapid(State(state): State<AppState>) -> Response {

    match configured(&state) {

        Ok((push, _)) => json_ok(

            &serde_json::json!({ "public_key": push.sender.vapid_public_key() }),

            None,

        ),

        Err(error) => error.into_response(),

    }

}

async fn subscribe(

    State(state): State<AppState>,

    headers: HeaderMap,

    body: axum::Json<SubscribeBody>,

) -> Response {

    let (push, public_url) = match configured(&state) {

        Ok(configured) => configured,

        Err(error) => return error.into_response(),

    };

    let (user, sid) = match auth(&state, &headers).await {

        Ok(pair) => pair,

        Err(error) => return error.into_response(),

    };

    let subscription = Subscription {

        endpoint: body.0.endpoint,

        p256dh: body.0.keys.p256dh,

        auth: body.0.keys.auth,

    };

    if !push.sender.allows(&subscription.endpoint) {

        return ApiError::bad_request("endpoint is not a known push service").into_response();

    }

    if !subscription.keys_valid() {

        return ApiError::bad_request("subscription keys are not valid").into_response();

    }

    let (pushkey, displaced) = push.claim(&sid, subscription);

    if let Some((owner, pushkey)) = displaced {

        let state = state.clone();

        tokio::spawn(async move {

            if let Ok(owner) = state.get_or_restore(&owner).await {

                let _ = owner

                    .handle

                    .remove_pusher(PusherKey {

                        app_id: HOSTED_APP_ID.to_owned(),

                        pushkey,

                    })

                    .await;

            }

        });

    }

    let pusher = HttpPusher {

        app_id: HOSTED_APP_ID.to_owned(),

        pushkey: pushkey.clone(),

        app_display_name: "malice".to_owned(),

        device_display_name: "Browser".to_owned(),

        lang: "en".to_owned(),

        url: format!("{public_url}{NOTIFY_PATH}"),

        format: PushFormat::EventIdOnly,

        data: serde_json::Map::new(),

    };

    match user.handle.set_pusher(pusher).await {

        Ok(()) => no_content(None),

        Err(error) => {

            push.forget(&pushkey);

            ApiError::from(malice_core::error::CoreError::from(error)).into_response()

        }

    }

}

async fn unsubscribe(

    State(state): State<AppState>,

    headers: HeaderMap,

    body: axum::Json<UnsubscribeBody>,

) -> Response {

    let Some(push) = state.push() else {

        return push_off().into_response();

    };

    let (user, sid) = match auth(&state, &headers).await {

        Ok(pair) => pair,

        Err(error) => return error.into_response(),

    };

    let Some(pushkey) = push.find(&sid, &body.0.endpoint) else {

        return no_content(None);

    };

    let removed = user

        .handle

        .remove_pusher(PusherKey {

            app_id: HOSTED_APP_ID.to_owned(),

            pushkey: pushkey.clone(),

        })

        .await;

    match removed {

        Ok(()) => {

            push.forget(&pushkey);

            no_content(None)

        }

        Err(error) => ApiError::from(malice_core::error::CoreError::from(error)).into_response(),

    }

}

async fn notify(State(state): State<AppState>, body: Bytes) -> Response {

    let Some(push) = state.push() else {

        return StatusCode::NOT_FOUND.into_response();

    };

    let notification = match parse_notify(&body) {

        Ok(notification) => notification,

        Err(error) => {

            return ApiError::bad_request(error.to_string()).into_response();

        }

    };

    if !notification.is_event() {

        return json_response(StatusCode::OK, notify_response(&[]));

    }

    let Ok(_permit) = push.permits.acquire().await else {

        return StatusCode::SERVICE_UNAVAILABLE.into_response();

    };

    let mut rejected = Vec::new();

    for device in &notification.devices {

        if deliver(&state, push, &notification, device).await {

            push.forget(&device.pushkey);

            rejected.push(device.pushkey.clone());

        }

    }

    json_response(StatusCode::OK, notify_response(&rejected))

}

/// Resolves and sends one device's push. True when its pushkey should be rejected.

async fn deliver(

    state: &AppState,

    push: &HostedPush,

    notification: &GatewayNotification,

    device: &GatewayDevice,

) -> bool {

    let (Some(room_id), Some(event_id)) = (&notification.room_id, &notification.event_id) else {

        return false;

    };

    if device.app_id != HOSTED_APP_ID {

        return true;

    }

    let Some(entry) = push.get(&device.pushkey) else {

        return true;

    };

    let user = match state.get_or_restore(&entry.sid).await {

        Ok(user) => user,

        Err(RestoreError::Unavailable(error)) => {

            tracing::warn!(%error, "push for a session that would not restore");

            return false;

        }

        Err(

            RestoreError::NotFound | RestoreError::Stale(_) | RestoreError::SoftLoggedOut { .. },

        ) => return true,

    };

    let status = match timeout(

        RESOLVE_TIMEOUT,

        user.handle.resolve_push_notification(room_id, event_id),

    )

    .await

    {

        Ok(Ok(status)) => status,

        Ok(Err(error)) => {

            tracing::warn!(%error, "resolving a push failed");

            return false;

        }

        Err(_) => {

            tracing::warn!("resolving a push timed out");

            return false;

        }

    };

    let PushNotificationStatus::Event(resolved) = status else {

        return false;

    };

    let unread = unread_total(&user.mirror.snapshot(), &resolved.room);

    let notice = notice(&resolved, unread);

    let urgency = if notice.mention || resolved.room.direct {

        Urgency::High

    } else {

        Urgency::Normal

    };

    let outcome = push

        .sender

        .send(

            &entry.subscription,

            &Message {

                payload: &notice.to_bytes(push.public_url.as_deref().unwrap_or_default()),

                urgency,

                topic: &notice.room_id,

            },

        )

        .await;

    if matches!(outcome, Outcome::Failed(_) | Outcome::Retryable) {

        tracing::warn!(?outcome, "web push delivery failed");

    }

    if !outcome.drops_subscription() {

        return false;

    }

    // continuwuity 26.9 ignores a notify response's rejected pushkeys

    let _ = user

        .handle

        .remove_pusher(PusherKey {

            app_id: HOSTED_APP_ID.to_owned(),

            pushkey: device.pushkey.clone(),

        })

        .await;

    true

}

#[cfg(test)]

mod tests {

    use axum::{body::Body, http::Request};

    use http_body_util::BodyExt;

    use malice_push::{Allowlist, Vapid};

    use serde_json::{Value, json};

    use tower::ServiceExt;

    use super::*;

    fn hosted(dir: &tempfile::TempDir) -> HostedPush {

        let sender = Sender::new(

            Vapid::load_or_create(&dir.path().join("vapid.pk8")).unwrap(),

            "mailto:test@example.org".to_owned(),

            Allowlist::default(),

        )

        .unwrap();

        HostedPush::new(

            dir.path(),

            Some("https://malice.example".to_owned()),

            sender,

        )

    }

    fn subscription(endpoint: &str) -> Subscription {

        Subscription {

            endpoint: endpoint.to_owned(),

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

            auth: "BTBZMqHH6r4Tts7J_aSIgg".to_owned(),

        }

    }

    #[test]

    fn registry_persists_and_moves_endpoints_between_sessions() {

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

        let push = hosted(&dir);

        let (first, displaced) = push.claim("alpha", subscription("https://fcm.googleapis.com/a"));

        assert!(displaced.is_none());

        let (again, _) = push.claim("alpha", subscription("https://fcm.googleapis.com/a"));

        assert_eq!(first, again);

        let (moved, displaced) = push.claim("beta", subscription("https://fcm.googleapis.com/a"));

        assert_ne!(moved, first);

        assert_eq!(displaced, Some(("alpha".to_owned(), first.clone())));

        push.claim("beta", subscription("https://fcm.googleapis.com/b"));

        let reloaded = hosted(&dir);

        assert_eq!(reloaded.pusher_keys("beta").len(), 2);

        assert!(reloaded.pusher_keys("alpha").is_empty());