reqkey 0.1.0

//! Actix Web 4 middleware integration.

use std::{rc::Rc, task::Context, task::Poll};

use ::actix_web::{
    body::{BoxBody, EitherBody, MessageBody},
    dev::{Service, ServiceRequest, ServiceResponse, Transform},
    http::{header as actix_header, StatusCode},
    Error as ActixError, HttpMessage, HttpResponse,
};
use futures_util::future::{ok, LocalBoxFuture, Ready};

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

/// Actix Web transform installed with `App::wrap`.
#[derive(Clone)]
pub struct ReqKeyMiddleware {
    middleware: Middleware,
}

impl ReqKeyMiddleware {
    /// Create Actix middleware from the shared engine.
    pub const fn new(middleware: Middleware) -> Self {
        Self { middleware }
    }
}

impl<S, B> Transform<S, ServiceRequest> for ReqKeyMiddleware
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = ActixError> + 'static,
    S::Future: 'static,
    B: MessageBody + 'static,
{
    type Response = ServiceResponse<EitherBody<B, BoxBody>>;
    type Error = ActixError;
    type InitError = ();
    type Transform = ReqKeyService<S>;
    type Future = Ready<Result<Self::Transform, Self::InitError>>;

    fn new_transform(&self, service: S) -> Self::Future {
        ok(ReqKeyService {
            service: Rc::new(service),
            middleware: self.middleware.clone(),
        })
    }
}

/// Per-worker Actix service produced by [`ReqKeyMiddleware`].
pub struct ReqKeyService<S> {
    service: Rc<S>,
    middleware: Middleware,
}

impl<S, B> Service<ServiceRequest> for ReqKeyService<S>
where
    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = ActixError> + 'static,
    S::Future: 'static,
    B: MessageBody + 'static,
{
    type Response = ServiceResponse<EitherBody<B, BoxBody>>;
    type Error = ActixError;
    type Future = LocalBoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&self, context: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.service.poll_ready(context)
    }

    fn call(&self, mut request: ServiceRequest) -> Self::Future {
        let service = Rc::clone(&self.service);
        let middleware = self.middleware.clone();
        let context = request_context(&request);

        Box::pin(async move {
            match middleware.authorize(context).await {
                AuthorizationOutcome::Bypass => service
                    .call(request)
                    .await
                    .map(ServiceResponse::map_into_left_body),
                AuthorizationOutcome::Denied(denial) => Ok(request
                    .into_response(denial_response(&denial))
                    .map_into_right_body()),
                AuthorizationOutcome::Authorized(authorized) => {
                    attach_state(&mut request, &authorized);
                    let result = service.call(request).await;
                    let mut response = match result {
                        Ok(response) => response,
                        Err(error) => {
                            middleware
                                .record(authorized, ResponseContext::new(500).with_latency_ms(0))
                                .await;
                            return Err(error);
                        }
                    };
                    let status = response.status().as_u16();
                    let original_headers = headers_to_http(response.headers());
                    append_headers(response.headers_mut(), &authorized.response_headers());
                    middleware
                        .record(
                            authorized.clone(),
                            ResponseContext::new(status)
                                .with_headers(original_headers)
                                .with_latency_ms(authorized.elapsed_ms()),
                        )
                        .await;
                    Ok(response.map_into_left_body())
                }
            }
        })
    }
}

fn request_context(request: &ServiceRequest) -> RequestContext {
    let method =
        http::Method::from_bytes(request.method().as_str().as_bytes()).unwrap_or(http::Method::GET);
    let mut context = RequestContext::new(method, request.path())
        .with_headers(headers_to_http(request.headers()));
    if !request.query_string().is_empty() {
        context = context.with_query(request.query_string());
    }
    if let Some(address) = request.peer_addr() {
        context = context.with_client_ip(address.ip().to_string());
    }
    context
}

fn attach_state(request: &mut ServiceRequest, authorized: &AuthorizedRequest) {
    if let Some(decision) = &authorized.decision {
        request.extensions_mut().insert(decision.clone());
    }
    if let Some(failure) = &authorized.failure {
        request.extensions_mut().insert(failure.clone());
    }
}

fn denial_response(denial: &DeniedResponse) -> HttpResponse<BoxBody> {
    let status =
        StatusCode::from_u16(denial.status_code).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR);
    let mut response = HttpResponse::build(status);
    response.content_type("application/json");
    if let Some(retry_after) = denial.retry_after {
        response.insert_header((actix_header::RETRY_AFTER, retry_after.to_string()));
    }
    response.body(denial.json_body())
}

fn headers_to_http(headers: &actix_header::HeaderMap) -> http::HeaderMap {
    let mut converted = http::HeaderMap::new();
    for (name, value) in headers {
        let Ok(name) = http::HeaderName::from_bytes(name.as_str().as_bytes()) else {
            continue;
        };
        let Ok(value) = http::HeaderValue::from_bytes(value.as_bytes()) else {
            continue;
        };
        converted.append(name, value);
    }
    converted
}

fn append_headers(target: &mut actix_header::HeaderMap, source: &http::HeaderMap) {
    for (name, value) in source {
        let Ok(name) = actix_header::HeaderName::try_from(name.as_str()) else {
            continue;
        };
        let Ok(value) = actix_header::HeaderValue::from_bytes(value.as_bytes()) else {
            continue;
        };
        target.insert(name, value);
    }
}