zon_router/

method_router.rs

use bytes::{Bytes, BytesMut};
use http::{header, HeaderValue, Method, StatusCode};
use zon_core::{Body, HttpMiddleware, HttpService, IntoHttpService, Response};

use crate::service::{BoxedHttpService, Fallback};

#[inline]
pub fn get<B, X>(
    svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
) -> MethodRouter<B>
where
    B: Send + 'static,
{
    MethodRouter::new().get(svc)
}

#[inline]
pub fn head<B, X>(
    svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
) -> MethodRouter<B>
where
    B: Send + 'static,
{
    MethodRouter::new().head(svc)
}

#[inline]
pub fn post<B, X>(
    svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
) -> MethodRouter<B>
where
    B: Send + 'static,
{
    MethodRouter::new().post(svc)
}

#[inline]
pub fn put<B, X>(
    svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
) -> MethodRouter<B>
where
    B: Send + 'static,
{
    MethodRouter::new().put(svc)
}

#[inline]
pub fn patch<B, X>(
    svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
) -> MethodRouter<B>
where
    B: Send + 'static,
{
    MethodRouter::new().patch(svc)
}

#[inline]
pub fn delete<B, X>(
    svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
) -> MethodRouter<B>
where
    B: Send + 'static,
{
    MethodRouter::new().delete(svc)
}

#[inline]
pub fn options<B, X>(
    svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
) -> MethodRouter<B>
where
    B: Send + 'static,
{
    MethodRouter::new().options(svc)
}

#[inline]
pub fn trace<B, X>(
    svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
) -> MethodRouter<B>
where
    B: Send + 'static,
{
    MethodRouter::new().trace(svc)
}

pub struct MethodRouter<B> {
    // Services for specific HTTP methods
    get: Option<BoxedHttpService<B>>, // also handles head if no head service set
    head: Option<BoxedHttpService<B>>,
    post: Option<BoxedHttpService<B>>,
    put: Option<BoxedHttpService<B>>,
    patch: Option<BoxedHttpService<B>>,
    delete: Option<BoxedHttpService<B>>,
    options: Option<BoxedHttpService<B>>,
    trace: Option<BoxedHttpService<B>>,

    /// The service to call when no service for a given request's method is set
    fallback: Fallback<B>,

    /// The allow header's value to add to the fallback service's response.
    allow_header: Option<Bytes>,
}

impl<B> MethodRouter<B>
where
    B: Send + 'static,
{
    #[inline]
    pub fn new() -> Self {
        Self {
            get: None,
            head: None,
            post: None,
            put: None,
            patch: None,
            delete: None,
            options: None,
            trace: None,
            fallback: Fallback::Default(BoxedHttpService::new(|| async {
                StatusCode::METHOD_NOT_ALLOWED
            })),
            allow_header: None,
        }
    }

    /// Register a service that's called for `GET` requests.
    ///
    /// If you don't specify a separate service for `HEAD` requests, the same
    /// service will be called for them, but any body set by it will be
    /// stripped.
    #[track_caller]
    pub fn get<X>(
        mut self,
        svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
    ) -> Self {
        register_service(
            &mut self.get,
            BoxedHttpService::new(svc),
            "GET",
            &mut self.allow_header,
            "GET,HEAD",
        );
        self
    }

    /// Register a service that's called for `HEAD` requests.
    #[track_caller]
    pub fn head<X>(
        mut self,
        svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
    ) -> Self {
        register_service(
            &mut self.head,
            BoxedHttpService::new(svc),
            "HEAD",
            &mut self.allow_header,
            "HEAD",
        );
        self
    }

    /// Register a service that's called for `POST` requests.
    #[track_caller]
    pub fn post<X>(
        mut self,
        svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
    ) -> Self {
        register_service(
            &mut self.post,
            BoxedHttpService::new(svc),
            "POST",
            &mut self.allow_header,
            "POST",
        );
        self
    }

    /// Register a service that's called for `PUT` requests.
    #[track_caller]
    pub fn put<X>(
        mut self,
        svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
    ) -> Self {
        register_service(
            &mut self.put,
            BoxedHttpService::new(svc),
            "PUT",
            &mut self.allow_header,
            "PUT",
        );
        self
    }

    /// Register a service that's called for `PATCH` requests.
    #[track_caller]
    pub fn patch<X>(
        mut self,
        svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
    ) -> Self {
        register_service(
            &mut self.patch,
            BoxedHttpService::new(svc),
            "PATCH",
            &mut self.allow_header,
            "PATCH",
        );
        self
    }

    /// Register a service that's called for `DELETE` requests.
    #[track_caller]
    pub fn delete<X>(
        mut self,
        svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
    ) -> Self {
        register_service(
            &mut self.delete,
            BoxedHttpService::new(svc),
            "DELETE",
            &mut self.allow_header,
            "DELETE",
        );
        self
    }

    /// Register a service that's called for `OPTIONS` requests.
    #[track_caller]
    pub fn options<X>(
        mut self,
        svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
    ) -> Self {
        register_service(
            &mut self.options,
            BoxedHttpService::new(svc),
            "OPTIONS",
            &mut self.allow_header,
            "OPTIONS",
        );
        self
    }

    /// Register a service that's called for `TRACE` requests.
    #[track_caller]
    pub fn trace<X>(
        mut self,
        svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
    ) -> Self {
        register_service(
            &mut self.trace,
            BoxedHttpService::new(svc),
            "TRACE",
            &mut self.allow_header,
            "TRACE",
        );
        self
    }

    /// Register a service that's called for requests that don't match any of
    /// the other services registered in this `MethodRouter`.
    ///
    /// This method is useful if you want to add a custom body to your
    /// method-not-allowed responses. Note that an `Accept` header will still be
    /// added on the response based on the other services registered, unless the
    /// fallback service already set this header.
    #[track_caller]
    pub fn fallback<X>(
        mut self,
        svc: impl IntoHttpService<B, X, ResponseBody = Body, Service: Send + 'static>,
    ) -> Self {
        if self.fallback.is_custom() {
            panic!("A fallback service was already set.");
        }
        self.fallback = Fallback::Custom(BoxedHttpService::new(svc));
        self
    }

    pub fn middleware(
        self,
        mid: impl HttpMiddleware<BoxedHttpService<B>, Service: HttpService<B, ResponseBody = Body> + Send>
            + Clone
            + 'static,
    ) -> Self {
        let apply_middleware = |svc| BoxedHttpService::new(mid.clone().apply(svc));
        Self {
            get: self.get.map(apply_middleware),
            head: self.head.map(apply_middleware),
            post: self.post.map(apply_middleware),
            put: self.put.map(apply_middleware),
            patch: self.patch.map(apply_middleware),
            delete: self.delete.map(apply_middleware),
            options: self.options.map(apply_middleware),
            trace: self.trace.map(apply_middleware),
            fallback: self.fallback.middleware(mid),
            allow_header: self.allow_header,
        }
    }
}

#[track_caller]
fn register_service<B: Send + 'static>(
    slot: &mut Option<BoxedHttpService<B>>,
    svc: BoxedHttpService<B>,
    method_name: &str,
    allow_header: &mut Option<Bytes>,
    allow: &str,
) {
    if slot.is_some() {
        panic!("Overlapping service. Cannot add two services that both handle `{method_name}`");
    }
    *slot = Some(svc);

    *allow_header = {
        let capacity = allow_header.as_ref().map_or(0, |hdr| hdr.len() + 1) + allow.len();
        let mut new_allow_hdr = BytesMut::with_capacity(capacity);
        if let Some(old_allow_hdr) = allow_header {
            new_allow_hdr.extend_from_slice(old_allow_hdr);
            new_allow_hdr.extend_from_slice(b",");
        }
        new_allow_hdr.extend_from_slice(allow.as_bytes());
        Some(new_allow_hdr.freeze())
    };
}

impl<B> Default for MethodRouter<B>
where
    B: Send + 'static,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<B> HttpService<B> for MethodRouter<B>
where
    B: Send + 'static,
{
    type ResponseBody = Body;

    async fn call(&self, request: http::Request<B>) -> Response {
        // Make sure nothing is forgotten via exhaustive pattern match
        let Self { get, head, post, put, patch, delete, options, trace, fallback, allow_header } =
            self;

        let method = request.method();
        macro_rules! call {
            ($method:expr, $maybe_svc:expr) => {
                if method == $method {
                    if let Some(svc) = $maybe_svc {
                        return svc.call(request).await;
                    }
                }
            };
        }

        if method == Method::HEAD {
            if let Some(svc) = head {
                return svc.call(request).await;
            }
            // GET services implicitly also handle HEAD requests
            if let Some(svc) = get {
                // strip bodies of responses to HEAD requests if they're handled
                // by the GET service
                let mut res = svc.call(request).await;
                *res.body_mut() = Body::empty();
                return res;
            }
        }

        call!(Method::GET, get);
        call!(Method::POST, post);
        call!(Method::PUT, put);
        call!(Method::PATCH, patch);
        call!(Method::DELETE, delete);
        call!(Method::OPTIONS, options);
        call!(Method::TRACE, trace);

        let mut res = fallback.call(request).await;
        maybe_add_allow_header(&mut res, allow_header);
        res
    }
}

fn maybe_add_allow_header(res: &mut http::Response<Body>, allow_header: &Option<Bytes>) {
    if let Some(allow_header) = allow_header {
        if let header::Entry::Vacant(v) = res.headers_mut().entry(header::ACCEPT) {
            if let Ok(value) = HeaderValue::from_maybe_shared(allow_header.clone()) {
                v.insert(value);
            } else {
                #[cfg(debug_assertions)]
                panic!("Invalid allow header value. This should never happen.");
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use zon_core::Body;

    use super::MethodRouter;

    #[test]
    fn method_router_hyper_compatible() {
        fn assert_sync<S: Sync>() {}

        // Need to be able to clone zon services for hyper integration.
        // Not realistic for MethodRouter directly, usual workaround is wrapping
        // in an `Arc`. Needs MethodRouter to be `Send` in addition to `Sync`.
        // This was previously not the case, this test ensures we don't regress.
        assert_sync::<Arc<MethodRouter<Body>>>();
    }
}