malice

use std::{path::PathBuf, sync::Arc};

use axum::{

    Router,

    body::Bytes,

    extract::{DefaultBodyLimit, State},

    http::{HeaderMap, HeaderValue, StatusCode, header},

    response::{IntoResponse, Response},

    routing::{get, post},

};

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

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

use serde_json::Value;

use tokio::sync::Semaphore;

pub const DEVICE_APP_ID: &str = "dev.malice.device";

pub const NOTIFY_PATH: &str = "/_matrix/push/v1/notify";

const NOTIFY_BODY_LIMIT: usize = 64 * 1024;

const CONCURRENT_NOTIFIES: usize = 16;

/// A push gateway that keeps nothing per user: each pusher's data carries its

/// browser subscription's endpoint and auth secret, and its pushkey is the

/// subscription's p256dh key, as Sygnal's webpush pushkin does it.

pub struct Gateway {

    sender: Sender,

    public_url: Option<String>,

    permits: Semaphore,

}

impl Gateway {

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

        Self {

            sender,

            public_url,

            permits: Semaphore::new(CONCURRENT_NOTIFIES),

        }

    }

    /// Keeps the VAPID key in `MALICE_DATA_DIR`, systemd's `STATE_DIRECTORY`

    /// or the XDG state dir. Push stays off when none of them is usable.

    pub fn from_env() -> Option<Self> {

        let Some(dir) = state_dir() else {

            tracing::warn!("no state directory for the VAPID key, push is off");

            return None;

        };

        if let Err(error) = create_private_dir(&dir) {

            tracing::warn!(%error, dir = %dir.display(), "creating the push state directory failed, push is off");

            return None;

        }

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

        match malice_push::env::sender(&dir, public_url.as_deref()) {

            Ok(sender) => Some(Self::new(sender, public_url)),

            Err(error) => {

                tracing::warn!(%error, "push is off");

                None

            }

        }

    }

    fn origin(&self, headers: &HeaderMap) -> String {

        if let Some(url) = &self.public_url {

            return url.clone();

        }

        headers

            .get(header::HOST)

            .and_then(|host| host.to_str().ok())

            .map(|host| format!("https://{host}"))

            .unwrap_or_default()

    }

}

fn state_dir() -> Option<PathBuf> {

    let var = |name| std::env::var_os(name).filter(|value| !value.is_empty());

    if let Some(dir) = var("MALICE_DATA_DIR") {

        return Some(dir.into());

    }

    if let Some(dirs) = var("STATE_DIRECTORY") {

        return std::env::split_paths(&dirs).next();

    }

    var("XDG_STATE_HOME")

        .map(PathBuf::from)

        .or_else(|| var("HOME").map(|home| PathBuf::from(home).join(".local/state")))

        .map(|dir| dir.join("malice"))

}

fn create_private_dir(path: &std::path::Path) -> std::io::Result<()> {

    let mut builder = std::fs::DirBuilder::new();

    builder.recursive(true);

    #[cfg(unix)]

    {

        use std::os::unix::fs::DirBuilderExt;

        builder.mode(0o700);

    }

    builder.create(path)

}

pub fn push_router(gateway: Arc<Gateway>) -> Router {

    Router::new()

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

        .route(NOTIFY_PATH, post(notify))

        .layer(DefaultBodyLimit::max(NOTIFY_BODY_LIMIT))

        .with_state(gateway)

}

/// Keeps the push paths from falling through to the app shell.

pub fn push_off_router() -> Router {

    let off = || async { StatusCode::NOT_FOUND };

    Router::new()

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

        .route(NOTIFY_PATH, post(off))

}

async fn vapid(State(gateway): State<Arc<Gateway>>) -> Response {

    json_response(

        StatusCode::OK,

        serde_json::to_vec(&serde_json::json!({ "public_key": gateway.sender.vapid_public_key() }))

            .unwrap_or_default(),

    )

}

pub fn json_response(status: StatusCode, body: Vec<u8>) -> Response {

    (

        status,

        [(

            header::CONTENT_TYPE,

            HeaderValue::from_static("application/json"),

        )],

        body,

    )

        .into_response()

}

async fn notify(State(gateway): State<Arc<Gateway>>, headers: HeaderMap, body: Bytes) -> Response {

    let notification = match parse_notify(&body) {

        Ok(notification) => notification,

        Err(error) => {

            return json_response(

                StatusCode::BAD_REQUEST,

                serde_json::to_vec(&serde_json::json!({ "error": error.to_string() }))

                    .unwrap_or_default(),

            );

        }

    };

    if !notification.is_event() {

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

    }

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

        return StatusCode::SERVICE_UNAVAILABLE.into_response();

    };

    let origin = gateway.origin(&headers);

    let mut rejected = Vec::new();

    for device in &notification.devices {

        if deliver(&gateway.sender, &notification, device, &origin)

            .await

            .drops_subscription()

        {

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

        }

    }

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

}

async fn deliver(

    sender: &Sender,

    notification: &GatewayNotification,

    device: &GatewayDevice,

    origin: &str,

) -> Outcome {

    let data = |key: &str| device.data.get(key).and_then(Value::as_str);

    let (DEVICE_APP_ID, Some(endpoint), Some(auth)) =

        (device.app_id.as_str(), data("endpoint"), data("auth"))

    else {

        return Outcome::Refused;

    };

    let subscription = Subscription {

        endpoint: endpoint.to_owned(),

        p256dh: device.pushkey.clone(),

        auth: auth.to_owned(),

    };

    let notice = device_notice(notification, device);

    let urgency = if notice.mention || notification.room_name.is_none() {

        Urgency::High

    } else {

        Urgency::Normal

    };

    let outcome = sender

        .send(

            &subscription,

            &Message {

                payload: &notice.to_bytes(origin),

                urgency,

                topic: &notice.room_id,

            },

        )

        .await;

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

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

    }

    outcome

}

/// Room and sender as the homeserver named them. A room without a name is

/// taken for a DM, so its title is the sender.

pub fn device_notice(notification: &GatewayNotification, device: &GatewayDevice) -> Notice {

    let sender = notification

        .sender_display_name

        .clone()

        .filter(|name| !name.trim().is_empty())

        .or_else(|| {

            notification

                .sender

                .as_ref()

                .map(|user| user.localpart().to_owned())

        })

        .unwrap_or_default();

    let content: Option<Value> = notification

        .content

        .as_ref()

        .and_then(|raw| serde_json::from_str(raw.get()).ok());

    let content = Content::of(

        notification.event_type.as_deref(),

        content.as_ref(),

        notification.user_is_target,

    );

    let direct = notification.room_name.is_none();

    Notice {

        title: notification

            .room_name

            .clone()

            .unwrap_or_else(|| sender.clone()),

        body: content.body(&sender, direct),

        room_id: notification

            .room_id

            .as_ref()

            .map(ToString::to_string)

            .unwrap_or_default(),

        event_id: notification

            .event_id

            .as_ref()

            .map(ToString::to_string)

            .unwrap_or_default(),

        encrypted: content == Content::Encrypted,

        mention: device.tweaks.get("highlight").and_then(Value::as_bool) == Some(true),

        silent: !device.tweaks.contains_key("sound"),

        unread: Some(notification.counts.unread),

    }

}