synd-api 0.4.0

use std::{
    convert::Infallible,
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};

use axum::{
    extract::Request,
    http::StatusCode,
    response::{IntoResponse, Response},
};
use pin_project::pin_project;
use tower::{Layer, Service};

use crate::principal::Principal;

pub trait Authenticate {
    // how to implementor fill this associate type
    // need impl trait in associate type ?
    // https://github.com/rust-lang/rust/issues/63063
    type Output: Future<Output = Result<Principal, ()>>;

    fn authenticate(&self, token: Option<String>) -> Self::Output;
}

#[expect(clippy::large_enum_variant)]
#[pin_project(project = AuthFutureProj)]
pub enum AuthenticateFuture<AuthFut, S, F> {
    Authenticate {
        req: Option<Request>,
        #[pin]
        auth_fut: AuthFut,
        inner: S,
    },
    ServiceCall {
        #[pin]
        service_fut: F,
    },
}

impl<AuthFut, S, F> AuthenticateFuture<AuthFut, S, F> {
    fn new(req: Request, auth_fut: AuthFut, inner: S) -> Self {
        AuthenticateFuture::Authenticate {
            req: Some(req),
            auth_fut,
            inner,
        }
    }
}

impl<AuthFut, S> Future for AuthenticateFuture<AuthFut, S, S::Future>
where
    AuthFut: Future<Output = Result<Principal, ()>>,
    S: Service<Request, Response = Response, Error = Infallible>,
{
    type Output = Result<Response, Infallible>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        match self.as_mut().project() {
            AuthFutureProj::Authenticate {
                req,
                auth_fut,
                inner,
            } => match auth_fut.poll(cx) {
                Poll::Ready(Ok(principal)) => {
                    let mut req = req.take().unwrap();
                    req.extensions_mut().insert(principal);
                    let service_fut = inner.call(req);

                    self.set(AuthenticateFuture::ServiceCall { service_fut });
                    self.poll(cx)
                }
                Poll::Ready(Err(())) => Poll::Ready(Ok(StatusCode::UNAUTHORIZED.into_response())),
                Poll::Pending => Poll::Pending,
            },
            AuthFutureProj::ServiceCall { service_fut } => service_fut.poll(cx),
        }
    }
}

#[derive(Clone)]
pub struct AuthenticateLayer<A> {
    authenticator: A,
}

impl<A> AuthenticateLayer<A> {
    pub fn new(authenticator: A) -> Self {
        Self { authenticator }
    }
}

impl<S, A> Layer<S> for AuthenticateLayer<A>
where
    A: Authenticate + Clone,
{
    type Service = AuthenticateService<S, A>;

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

#[derive(Clone)]
pub struct AuthenticateService<S, A> {
    inner: S,
    authenticator: A,
}

impl<S, A> Service<Request> for AuthenticateService<S, A>
where
    S: Service<Request, Response = Response, Error = Infallible> + Clone,
    A: Authenticate,
{
    type Response = Response;
    type Error = Infallible;
    type Future = AuthenticateFuture<A::Output, S, S::Future>;

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

    fn call(&mut self, req: Request) -> Self::Future {
        let token = req
            .headers()
            .get(axum::http::header::AUTHORIZATION)
            .and_then(|header| header.to_str().ok())
            .map(ToOwned::to_owned);

        let auth_fut = self.authenticator.authenticate(token);
        let inner = self.inner.clone();

        AuthenticateFuture::new(req, auth_fut, inner)
    }
}

#[cfg(test)]
mod tests {
    use std::{
        future::{Ready, ready},
        task::{Context, Poll},
    };

    use axum::{
        body::Body,
        http::{StatusCode, header::AUTHORIZATION},
    };

    use crate::principal::{Principal, User};

    use super::*;

    #[derive(Clone)]
    struct TestAuthenticator {
        token: &'static str,
    }

    impl Authenticate for TestAuthenticator {
        type Output = Ready<Result<Principal, ()>>;

        fn authenticate(&self, token: Option<String>) -> Self::Output {
            if token.as_deref() == Some(self.token) {
                ready(Ok(Principal::User(User::local())))
            } else {
                ready(Err(()))
            }
        }
    }

    #[derive(Clone)]
    struct OkService;

    impl Service<Request> for OkService {
        type Response = Response;
        type Error = Infallible;
        type Future = Ready<Result<Response, Infallible>>;

        fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
            Poll::Ready(Ok(()))
        }

        fn call(&mut self, _req: Request) -> Self::Future {
            ready(Ok(StatusCode::OK.into_response()))
        }
    }

    fn request(token: Option<&str>) -> Request {
        let mut builder = Request::builder();
        if let Some(token) = token {
            builder = builder.header(AUTHORIZATION, token);
        }

        builder.body(Body::empty()).unwrap()
    }

    #[tokio::test]
    async fn valid_token_passes_request() {
        let layer = AuthenticateLayer::new(TestAuthenticator {
            token: "Bearer secret",
        });
        let mut service = layer.layer(OkService);

        let response = service.call(request(Some("Bearer secret"))).await.unwrap();

        assert_eq!(response.status(), StatusCode::OK);
    }

    #[tokio::test]
    async fn missing_or_wrong_token_returns_unauthorized() {
        let layer = AuthenticateLayer::new(TestAuthenticator {
            token: "Bearer secret",
        });

        let mut missing_token_service = layer.layer(OkService);
        let response = missing_token_service.call(request(None)).await.unwrap();
        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);

        let mut wrong_token_service = layer.layer(OkService);
        let response = wrong_token_service
            .call(request(Some("Bearer wrong")))
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
    }
}