zon_middleware 0.0.6

use std::future::Future;

use http::{Request, Response};
use tokio::task::LocalKey;
use zon_core::{HttpMiddleware, HttpService};

/// Middleware to set tokio task-local data.
#[derive(Debug, Clone, Copy)]
pub struct SetTaskLocal<T: 'static> {
    key: &'static LocalKey<T>,
    value: T,
}

impl<T> SetTaskLocal<T>
where
    T: Clone + Send + Sync + 'static,
{
    pub fn new(key: &'static LocalKey<T>, value: T) -> Self {
        SetTaskLocal { key, value }
    }
}

impl<S, T> HttpMiddleware<S> for SetTaskLocal<T>
where
    T: Clone + Send + Sync + 'static,
{
    type Service = SetTaskLocalService<S, T>;

    fn apply(self, inner: S) -> Self::Service {
        SetTaskLocalService::new(inner, self.key, self.value)
    }
}

/// Service for [`SetTaskLocal`] middleware.
#[derive(Debug, Clone, Copy)]
pub struct SetTaskLocalService<S, T: 'static> {
    inner: S,
    key: &'static LocalKey<T>,
    value: T,
}

impl<S, T> SetTaskLocalService<S, T>
where
    T: Clone + Send + Sync + 'static,
{
    pub fn new(inner: S, key: &'static LocalKey<T>, value: T) -> Self {
        Self { inner, key, value }
    }
}

impl<S, T, B> HttpService<B> for SetTaskLocalService<S, T>
where
    S: HttpService<B>,
    T: Clone + Send + Sync + 'static,
    B: Send,
{
    type ResponseBody = S::ResponseBody;

    fn call(
        &self,
        request: Request<B>,
    ) -> impl Future<Output = Response<Self::ResponseBody>> + Send {
        self.key.scope(self.value.clone(), self.inner.call(request))
    }
}