malice

use std::sync::Arc;

use axum::{

    Router,

    extract::{DefaultBodyLimit, Multipart, Query, State},

    http::{HeaderMap, StatusCode},

    middleware,

    response::{IntoResponse, Response, Sse, sse::Event as SseEvent},

    routing::{get, post},

};

use futures_util::StreamExt;

use macltok::Retain;

use malice_core::{

    feed::subscribe,

    login::LoginInput,

    request::{Reply, Request, SendKind, dispatch},

};

use serde::{Deserialize, Serialize};

use tower_http::{

    compression::{

        CompressionLayer,

        predicate::{DefaultPredicate, NotForContentType, Predicate},

    },

    services::{ServeDir, ServeFile},

    trace::TraceLayer,

};

use crate::{

    csrf::{csrf_guard, session_sid, set_cookie_value},

    error::{ApiError, json_ok, no_content},

    state::{AppState, RestoreError, UserSession},

};

#[derive(Debug, Deserialize)]

struct RoomIdBody {

    #[serde(default)]

    room_id: String,

}

#[derive(Debug, Deserialize)]

struct SendBody {

    #[serde(default)]

    room_id: String,

    #[serde(default)]

    body: String,

    #[serde(default)]

    formatted_body: Option<String>,

    #[serde(default)]

    reply_to: Option<String>,

    #[serde(default)]

    mentions: Vec<String>,

    #[serde(default)]

    kind: SendKind,

    #[serde(default)]

    thread_root: Option<String>,

}

#[derive(Debug, Deserialize)]

struct PaginateBody {

    #[serde(default)]

    room_id: String,

    #[serde(default)]

    thread_root: Option<String>,

}

#[derive(Debug, Deserialize)]

struct JoinBody {

    #[serde(default)]

    room: String,

}

#[derive(Debug, Deserialize)]

struct CreateBody {

    #[serde(default)]

    name: String,

    #[serde(default)]

    topic: Option<String>,

}

#[derive(Debug, Deserialize)]

struct ReactBody {

    #[serde(default)]

    room_id: String,

    #[serde(default)]

    event_id: String,

    #[serde(default)]

    key: String,

    #[serde(default)]

    shortcode: Option<String>,

    #[serde(default)]

    thread_root: Option<String>,

}

#[derive(Debug, Deserialize)]

struct EditBody {

    #[serde(default)]

    room_id: String,

    #[serde(default)]

    event_id: String,

    #[serde(default)]

    body: String,

    #[serde(default)]

    formatted_body: Option<String>,

    #[serde(default)]

    mentions: Vec<String>,

    #[serde(default)]

    thread_root: Option<String>,

}

#[derive(Debug, Deserialize)]

struct ReadBody {

    #[serde(default)]

    room_id: String,

    #[serde(default)]

    event_id: String,

    #[serde(default)]

    thread_root: Option<String>,

}

#[derive(Debug, Deserialize)]

struct MembersQuery {

    #[serde(default)]

    room_id: String,

}

#[derive(Debug, Deserialize)]

struct RecoverBody {

    recovery_key: String,

}

#[derive(Debug, Deserialize)]

struct VerifyBody {

    action: String,

}

#[derive(Debug, Deserialize)]

struct MediaQuery {

    mxc: String,

    size: Option<u32>,

}

#[derive(Debug, Deserialize)]

struct MessageMediaQuery {

    media_id: String,

    size: String,

}

#[derive(Debug, Deserialize)]

struct TypingBody {

    #[serde(default)]

    room_id: String,

    #[serde(default)]

    typing: bool,

}

#[derive(Debug, Deserialize)]

struct AccountDataQuery {

    event_type: String,

}

#[derive(Debug, Deserialize)]

struct UrlPreviewQuery {

    #[serde(default)]

    room_id: String,

    #[serde(default)]

    url: String,

}

#[derive(Debug, Deserialize)]

struct AccountDataBody {

    event_type: String,

    content: serde_json::Value,

}

#[derive(Debug, Serialize)]

struct SoftLogoutInfo {

    user_id: String,

    homeserver: String,

}

#[derive(Debug, Serialize)]

struct SessionResponse {

    user_id: Option<String>,

    #[serde(skip_serializing_if = "Option::is_none")]

    soft_logout: Option<SoftLogoutInfo>,

}

pub(crate) async fn auth(

    state: &AppState,

    headers: &HeaderMap,

) -> Result<(Arc<UserSession>, String), ApiError> {

    let Some(sid) = session_sid(headers) else {

        return Err(ApiError::unauthorized("not signed in"));

    };

    match state.get_or_restore(&sid).await {

        Ok(user) => Ok((user, sid)),

        Err(RestoreError::NotFound) => Err(ApiError::unauthorized_stale(

            "session expired; sign in again",

            state.cookie_secure(),

        )),

        Err(RestoreError::SoftLoggedOut { .. }) => Err(ApiError::unauthorized(

            "signed out by the server, sign in again",

        )),

        Err(RestoreError::Stale(message)) => {

            state.remove_session(&sid, Retain::Nothing).await;

            Err(ApiError::unauthorized_stale(message, state.cookie_secure()))

        }

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

            Err(ApiError::new(StatusCode::BAD_GATEWAY, message))

        }

    }

}

fn avatar_response(bytes: Vec<u8>, content_type: &'static str) -> Response {

    (

        [

            (axum::http::header::CONTENT_TYPE, content_type),

            // mxc uris name immutable content, a changed avatar gets a new uri

            (

                axum::http::header::CACHE_CONTROL,

                "private, max-age=31536000, immutable",

            ),

            (axum::http::header::X_CONTENT_TYPE_OPTIONS, "nosniff"),

        ],

        bytes,

    )

        .into_response()

}

/// commands answer when they're done, so the browser gets the outcome in the response

fn respond(result: Result<Reply, malice_core::error::CoreError>) -> Response {

    match result {

        Ok(Reply::Done) => no_content(None),

        Ok(Reply::Json(value)) => json_ok(&value, None),

        Ok(Reply::Avatar(source)) => {

            let (bytes, content_type) = source.render();

            avatar_response(bytes, content_type)

        }

        Ok(Reply::Attachment {

            bytes,

            content_type,

        }) => (

            [

                (

                    axum::http::header::CONTENT_TYPE,

                    content_type.unwrap_or("application/octet-stream"),

                ),

                (

                    axum::http::header::CONTENT_DISPOSITION,

                    if content_type.is_some() {

                        "inline"

                    } else {

                        "attachment"

                    },

                ),

                (axum::http::header::CACHE_CONTROL, "private, max-age=3600"),

                (axum::http::header::X_CONTENT_TYPE_OPTIONS, "nosniff"),

                (axum::http::header::CONTENT_SECURITY_POLICY, "sandbox"),

            ],

            bytes,

        )

            .into_response(),

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

    }

}

async fn get_session(State(state): State<AppState>, headers: HeaderMap) -> Response {

    let Some(sid) = session_sid(&headers) else {

        return json_ok(

            &SessionResponse {

                user_id: None,

                soft_logout: None,

            },

            None,

        );

    };

    match state.get_or_restore(&sid).await {

        Ok(user) => json_ok(

            &SessionResponse {

                user_id: Some(user.user_id.clone()),

                soft_logout: None,

            },

            None,

        ),

        Err(RestoreError::NotFound) => json_ok(

            &SessionResponse {

                user_id: None,

                soft_logout: None,

            },

            None,

        ),

        Err(RestoreError::SoftLoggedOut {

            user_id,

            homeserver,

        }) => json_ok(

            &SessionResponse {

                user_id: None,

                soft_logout: Some(SoftLogoutInfo {

                    user_id,

                    homeserver,

                }),

            },

            None,

        ),

        Err(RestoreError::Stale(message)) => {

            state.remove_session(&sid, Retain::Nothing).await;

            ApiError::unauthorized_stale(message, state.cookie_secure()).into_response()

        }

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

            ApiError::new(StatusCode::BAD_GATEWAY, message).into_response()

        }

    }

}

async fn post_login(

    State(state): State<AppState>,

    headers: HeaderMap,

    body: axum::Json<LoginInput>,

) -> Response {

    let username = body.0.username.trim().to_owned();

    let password = body.0.password.clone();

    let login = match body.0.into_login(state.allow_http_homeserver()) {

        Ok(login) => login,

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

    };

    let cookie_sid = session_sid(&headers);

    match state

        .login(cookie_sid.as_deref(), login, &username, password)

        .await

    {

        Ok(user) => json_ok(

            &serde_json::json!({ "user_id": user.user_id }),

            Some(set_cookie_value(&user.sid, state.cookie_secure())),

        ),

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

    }

}

async fn post_logout(State(state): State<AppState>, headers: HeaderMap) -> Response {

    if let Some(sid) = session_sid(&headers) {

        state.logout(&sid).await;

    }

    let mut response = no_content(Some(crate::csrf::clear_cookie_value(state.cookie_secure())));

    response.headers_mut().insert(

        axum::http::header::HeaderName::from_static("clear-site-data"),

        axum::http::HeaderValue::from_static("\"cache\""),

    );

    response

}

async fn get_rooms(State(state): State<AppState>, headers: HeaderMap) -> Response {

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

        Ok(pair) => pair,

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

    };

    respond(dispatch(&user.handle, &user.mirror, Request::Rooms).await)

}

async fn get_members(

    State(state): State<AppState>,

    headers: HeaderMap,

    Query(query): Query<MembersQuery>,

) -> Response {

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

        Ok(pair) => pair,

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

    };

    respond(

        dispatch(

            &user.handle,

            &user.mirror,

            Request::Members {

                room_id: query.room_id,

            },

        )

        .await,

    )

}

async fn rooms_open(

    State(state): State<AppState>,

    headers: HeaderMap,

    body: axum::Json<RoomIdBody>,

) -> Response {

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

        Ok(pair) => pair,

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

    };

    respond(

        dispatch(

            &user.handle,

            &user.mirror,

            Request::Open {

                room_id: body.0.room_id,

            },

        )

        .await,

    )

}

async fn rooms_close(

    State(state): State<AppState>,

    headers: HeaderMap,

    body: axum::Json<RoomIdBody>,

) -> Response {

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

        Ok(pair) => pair,

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

    };

    respond(

        dispatch(

            &user.handle,

            &user.mirror,

            Request::Close {

                room_id: body.0.room_id,

            },

        )

        .await,

    )

}

async fn rooms_send(

    State(state): State<AppState>,

    headers: HeaderMap,

    body: axum::Json<SendBody>,

) -> Response {

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

        Ok(pair) => pair,

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

    };

    respond(

        dispatch(

            &user.handle,

            &user.mirror,

            Request::Send {

                room_id: body.0.room_id,

                body: body.0.body,

                formatted_body: body.0.formatted_body,

                reply_to: body.0.reply_to,

                mentions: body.0.mentions,

                kind: body.0.kind,

                thread_root: body.0.thread_root,

            },

        )

        .await,

    )

}

async fn rooms_upload(

    State(state): State<AppState>,

    headers: HeaderMap,

    mut upload: Multipart,

) -> Response {

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

        Ok(pair) => pair,

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

    };

    let mut room_id = None;

    let mut caption = None;

    let mut formatted_caption = None;

    let mut reply_to = None;

    let mut thread_root = None;

    let mut mentions = Vec::new();

    let mut file = None;

    loop {

        let field = match upload.next_field().await {

            Ok(Some(field)) => field,

            Ok(None) => break,

            Err(_) => return ApiError::bad_request("invalid upload form").into_response(),

        };

        match field.name() {

            Some("room_id") if room_id.is_none() => {

                room_id = match field.text().await {

                    Ok(value) => Some(value),

                    Err(_) => return ApiError::bad_request("invalid room_id").into_response(),

                };

            }

            Some("caption") if caption.is_none() => {

                caption = match field.text().await {

                    Ok(value) => Some(value),

                    Err(_) => return ApiError::bad_request("invalid caption").into_response(),

                };

            }

            Some("formatted_caption") if formatted_caption.is_none() => {

                formatted_caption = match field.text().await {

                    Ok(value) => Some(value),

                    Err(_) => return ApiError::bad_request("invalid caption").into_response(),

                };

            }