ts_webapi/middleware/

api_key.rs

//! API key middleware, validates an API key is present and in the allow list.

use alloc::sync::Arc;
use core::{
    pin::Pin,
    task::{Context, Poll, ready},
};

use http::{HeaderName, HeaderValue, Request, Response};
use http_body::Body;
use pin_project_lite::pin_project;
use tower_layer::Layer;
use tower_service::Service;

use crate::{ErrorResponse, middleware::response_body::ResponseBody};

/// API key validation layer.
#[derive(Debug, Clone)]
pub struct ApiKeyAuth {
    /// The header to find the API key in.
    pub header: Arc<HeaderName>,
    /// The list of allowed keys.
    pub allowed_keys: Arc<Vec<String>>,
}
impl ApiKeyAuth {
    /// Create new API key auth.
    pub fn new(header: HeaderName, allowed_keys: Vec<String>) -> Self {
        Self {
            allowed_keys: Arc::new(allowed_keys),
            header: Arc::new(header),
        }
    }

    /// Authenticate a request against this auth
    pub fn authenticate<T>(&self, request: &Request<T>) -> bool {
        let Some(Ok(api_key)) = request
            .headers()
            .get(self.header.as_ref())
            .map(HeaderValue::to_str)
        else {
            return false;
        };

        self.allowed_keys
            .iter()
            .map(String::as_str)
            .any(|key| key.eq(api_key))
    }
}

impl<S> Layer<S> for ApiKeyAuth {
    type Service = ApiKeyAuthService<S>;

    fn layer(&self, inner: S) -> Self::Service {
        ApiKeyAuthService {
            inner,
            auth: self.clone(),
        }
    }
}

/// Tower service behind the API key auth middleware.
#[derive(Debug, Clone)]
pub struct ApiKeyAuthService<S> {
    /// Inner service.
    inner: S,
    /// The logic layer.
    auth: ApiKeyAuth,
}

impl<S> ApiKeyAuthService<S> {
    /// Create a new service.
    pub fn new(inner: S, auth: ApiKeyAuth) -> Self {
        Self { inner, auth }
    }
}

impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for ApiKeyAuthService<S>
where
    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
    ResBody: Body,
{
    type Response = Response<ResponseBody<ResBody>>;
    type Error = S::Error;
    type Future = ResponseFuture<S::Future>;

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

    fn call(&mut self, request: Request<ReqBody>) -> Self::Future {
        if !self.auth.authenticate(&request) {
            ResponseFuture {
                inner: Kind::Unauthorized,
            }
        } else {
            ResponseFuture {
                inner: Kind::Ok {
                    future: self.inner.call(request),
                },
            }
        }
    }
}

pin_project! {
    /// Response future for [`ApiKeyAuthService`].
    pub struct ResponseFuture<F> {
        #[pin]
        inner: Kind<F>,
    }
}

pin_project! {
    #[project = KindProj]
    enum Kind<F> {
        Ok {
            #[pin]
            future: F,
        },
        Unauthorized,
    }
}

impl<ResBody, F, E> Future for ResponseFuture<F>
where
    ResBody: Body,
    F: Future<Output = Result<Response<ResBody>, E>>,
{
    type Output = Result<Response<ResponseBody<ResBody>>, E>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let response = match self.project().inner.project() {
            KindProj::Ok { future } => ready!(future.poll(cx))?.map(ResponseBody::new),
            KindProj::Unauthorized => ErrorResponse::unauthenticated().as_response(),
        };
        Poll::Ready(Ok(response))
    }
}

#[cfg(test)]
mod test {
    use axum::{Router, routing::get};
    use bytes::Bytes;
    use http::{HeaderName, HeaderValue, Request, Response, StatusCode};
    use http_body_util::Full;
    use tower::{ServiceBuilder, ServiceExt};
    use tower_http::BoxError;

    use crate::middleware::api_key::ApiKeyAuth;

    async fn echo(req: Request<Full<Bytes>>) -> Result<Response<Full<Bytes>>, BoxError> {
        Ok(Response::new(req.into_body()))
    }

    #[tokio::test]
    async fn axum_compat() {
        let api_key_auth = ApiKeyAuth::new(
            HeaderName::from_static("x-api-key"),
            vec!["api-key-1".to_string()],
        );

        let router = Router::new()
            .route("/", get(|| async { StatusCode::OK }))
            .layer(api_key_auth);

        let response = router
            .oneshot(
                Request::builder()
                    .uri("/")
                    .body(axum::body::Body::empty())
                    .unwrap(),
            )
            .await
            .unwrap();
        assert_eq!(StatusCode::UNAUTHORIZED, response.status());
    }

    #[tokio::test]
    async fn blocks_no_header() {
        let api_key_auth = ApiKeyAuth::new(
            HeaderName::from_static("x-api-key"),
            vec!["api-key-1".to_string()],
        );

        let mut service = ServiceBuilder::new().layer(api_key_auth).service_fn(echo);

        let body: Bytes = Bytes::new();
        let request = Request::new(Full::new(body));

        let response = service
            .ready()
            .await
            .unwrap()
            .oneshot(request)
            .await
            .unwrap();
        assert_eq!(StatusCode::UNAUTHORIZED, response.status());
    }

    #[tokio::test]
    async fn blocks_not_allowed() {
        let api_key_auth = ApiKeyAuth::new(
            HeaderName::from_static("x-api-key"),
            vec!["api-key-1".to_string()],
        );

        let mut service = ServiceBuilder::new().layer(api_key_auth).service_fn(echo);

        let body: Bytes = Bytes::new();
        let mut request = Request::new(Full::new(body));
        request
            .headers_mut()
            .insert("x-api-key", HeaderValue::from_static("not-allowed-key"));

        let response = service
            .ready()
            .await
            .unwrap()
            .oneshot(request)
            .await
            .unwrap();
        assert_eq!(StatusCode::UNAUTHORIZED, response.status());
    }

    #[tokio::test]
    async fn allows_allowed() {
        let api_key_auth = ApiKeyAuth::new(
            HeaderName::from_static("x-api-key"),
            vec!["api-key-1".to_string()],
        );

        let mut service = ServiceBuilder::new().layer(api_key_auth).service_fn(echo);

        let body: Bytes = Bytes::new();
        let mut request = Request::new(Full::new(body));
        request
            .headers_mut()
            .insert("x-api-key", HeaderValue::from_static("api-key-1"));

        let response = service
            .ready()
            .await
            .unwrap()
            .oneshot(request)
            .await
            .unwrap();
        assert_eq!(StatusCode::OK, response.status());
    }
}