jelq

use std::{collections::HashMap, time::Instant};

use axum::{

    body::Bytes,

    extract::{FromRequest, FromRequestParts, Query, RawPathParams, Request},

    http::{HeaderMap, HeaderName, Method, StatusCode, Uri, Version, header},

    response::{IntoResponse, Response},

};

use hyper::upgrade::OnUpgrade;

use crate::problem::Problem;

const BODY_LIMIT: usize = 64 * 1024;

/// Everything a handler may look at, gathered once per request.

pub struct Ctx {

    pub method: Method,

    pub uri: Uri,

    pub version: Version,

    pub headers: HeaderMap,

    pub query: HashMap<String, String>,

    pub params: HashMap<String, String>,

    pub body: Bytes,

    pub now: Instant,

    pub upgrade: Option<OnUpgrade>,

}

impl Ctx {

    pub fn header(&self, name: HeaderName) -> Option<&str> {

        self.headers.get(name).and_then(|v| v.to_str().ok())

    }

    pub fn cookie(&self, name: &str) -> Option<&str> {

        self.headers

            .get_all(header::COOKIE)

            .iter()

            .filter_map(|v| v.to_str().ok())

            .flat_map(|v| v.split(';'))

            .filter_map(|pair| pair.trim().split_once('='))

            .find(|(key, _)| *key == name)

            .map(|(_, value)| value)

    }

    pub fn param(&self, name: &str) -> &str {

        self.params.get(name).map_or("", String::as_str)

    }

    pub fn query(&self, name: &str) -> Option<&str> {

        self.query.get(name).map(String::as_str)

    }

}

impl<S: Send + Sync> FromRequest<S> for Ctx {

    type Rejection = Response;

    async fn from_request(mut request: Request, state: &S) -> Result<Self, Self::Rejection> {

        let upgrade = request.extensions_mut().remove::<OnUpgrade>();

        let (mut parts, body) = request.into_parts();

        let params = RawPathParams::from_request_parts(&mut parts, state)

            .await

            .map(|raw| {

                raw.iter()

                    .map(|(k, v)| (k.to_owned(), v.to_owned()))

                    .collect()

            })

            .unwrap_or_default();

        let query = Query::<HashMap<String, String>>::try_from_uri(&parts.uri)

            .map(|q| q.0)

            .unwrap_or_default();

        let body = axum::body::to_bytes(body, BODY_LIMIT).await.map_err(|_| {

            Problem::new(

                StatusCode::PAYLOAD_TOO_LARGE,

                "body-too-large",

                "The request body is larger than anything the server is prepared to receive.",

            )

            .into_response()

        })?;

        Ok(Self {

            method: parts.method,

            uri: parts.uri,

            version: parts.version,

            headers: parts.headers,

            query,

            params,

            body,

            now: Instant::now(),

            upgrade,

        })

    }

}