reqkey 0.1.0

//! Warp 0.4 composable filter integration.
//!
//! Warp filters do not have an after-response middleware hook. The filter
//! returns [`WarpRequest`]; handlers call [`WarpRequest::record`] around their
//! reply to add decision headers and complete analytics.

use std::{convert::Infallible, future::Ready, net::SocketAddr};

use ::warp::{
    filters::path::FullPath,
    http::{HeaderMap, StatusCode},
    reject::Reject,
    Filter, Rejection, Reply,
};

use crate::{
    middleware::{
        AuthorizationOutcome, AuthorizedRequest, DeniedResponse, Middleware, ReqKeyFailure,
        RequestContext, ResponseContext,
    },
    VerificationResult,
};

/// Cloneable ReqKey filter factory.
#[derive(Clone)]
pub struct ReqKeyFilter {
    middleware: Middleware,
}

impl ReqKeyFilter {
    /// Create a Warp filter factory from the shared middleware engine.
    pub const fn new(middleware: Middleware) -> Self {
        Self { middleware }
    }

    /// Build a filter that authorizes and extracts [`WarpRequest`].
    pub fn filter(
        &self,
    ) -> impl Filter<Extract = (WarpRequest,), Error = Rejection> + Clone + Send + Sync + 'static
    {
        let middleware = self.middleware.clone();
        let query = ::warp::query::raw()
            .or(::warp::any().map(String::new))
            .unify();
        ::warp::method()
            .and(::warp::path::full())
            .and(query)
            .and(::warp::header::headers_cloned())
            .and(::warp::addr::remote())
            .and_then(
                move |method: http::Method,
                      path: FullPath,
                      query: String,
                      headers: HeaderMap,
                      remote: Option<SocketAddr>| {
                    let middleware = middleware.clone();
                    async move {
                        let mut context =
                            RequestContext::new(method, path.as_str()).with_headers(headers);
                        if !query.is_empty() {
                            context = context.with_query(query);
                        }
                        if let Some(remote) = remote {
                            context = context.with_client_ip(remote.ip().to_string());
                        }
                        match middleware.authorize(context).await {
                            AuthorizationOutcome::Bypass => Ok(WarpRequest {
                                middleware,
                                authorized: None,
                            }),
                            AuthorizationOutcome::Authorized(authorized) => Ok(WarpRequest {
                                middleware,
                                authorized: Some(authorized),
                            }),
                            AuthorizationOutcome::Denied(denial) => {
                                Err(::warp::reject::custom(ReqKeyRejection(denial)))
                            }
                        }
                    }
                },
            )
    }
}

/// Authorized request value extracted by [`ReqKeyFilter`].
pub struct WarpRequest {
    middleware: Middleware,
    authorized: Option<AuthorizedRequest>,
}

impl WarpRequest {
    /// Return the validation decision in validation/both mode.
    pub fn decision(&self) -> Option<&VerificationResult> {
        self.authorized
            .as_ref()
            .and_then(|authorized| authorized.decision.as_ref())
    }

    /// Return the fail-open service error, when present.
    pub fn failure(&self) -> Option<&ReqKeyFailure> {
        self.authorized
            .as_ref()
            .and_then(|authorized| authorized.failure.as_ref())
    }

    /// Wrap a handler reply, add decision headers, and await analytics.
    pub async fn record(self, reply: impl Reply) -> ::warp::reply::Response {
        let mut response = reply.into_response();
        let Some(authorized) = self.authorized else {
            return response;
        };
        let original_headers = response.headers().clone();
        response.headers_mut().extend(authorized.response_headers());
        self.middleware
            .record(
                authorized.clone(),
                ResponseContext::new(response.status().as_u16())
                    .with_headers(original_headers)
                    .with_latency_ms(authorized.elapsed_ms()),
            )
            .await;
        response
    }
}

/// Custom rejection generated by ReqKey policy.
#[derive(Debug)]
pub struct ReqKeyRejection(pub DeniedResponse);

impl Reject for ReqKeyRejection {}

/// Convert a ReqKey rejection into its stable JSON response.
///
/// Chain this with `Filter::recover`. Unknown rejections are returned so
/// another recovery layer can handle them.
///
/// # Errors
///
/// Returns the original rejection when it was not generated by ReqKey.
pub fn recover(rejection: Rejection) -> Ready<Result<::warp::reply::Response, Rejection>> {
    std::future::ready(recover_inner(rejection))
}

fn recover_inner(rejection: Rejection) -> Result<::warp::reply::Response, Rejection> {
    let Some(rejection) = rejection.find::<ReqKeyRejection>() else {
        return Err(rejection);
    };
    let denial = &rejection.0;
    let status =
        StatusCode::from_u16(denial.status_code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
    let mut response = ::warp::reply::with_status(
        ::warp::reply::json(&serde_json::json!({
            "error": denial.error.as_str(),
            "message": denial.message,
        })),
        status,
    )
    .into_response();
    response.headers_mut().extend(denial.headers());
    Ok(response)
}

/// Helper return type for infallible Warp handlers.
pub type InfallibleReply<T> = Result<T, Infallible>;