tide-tracing/src/lib.rs

#![feature(const_trait_impl)]
#![feature(const_fn_trait_bound)]

#[macro_use]
extern crate tracing;

use nrid::Nrid;
use std::time::Instant;
use tide::Middleware;
use tide::Next;
use tide::Request;
use tracing::Instrument;

#[derive(Clone, Copy, Debug, Eq)]
#[repr(transparent)]
pub struct ReqComps(u8);

#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, num_enum::FromPrimitive)]
#[repr(u8)]
pub enum ReqComp {
    #[num_enum(default)]
    __Default = 0b00000000,
    Method = 0b00000001,
    Path = 0b00000010,
    IP = 0b00000100,
    Query = 0b00001000,
}

impl ReqComps {
    pub const fn of(comp: ReqComp) -> Self {
        Self(comp as u8)
    }

    pub const fn contains(self, mask: ReqComp) -> bool {
        self & mask == mask
    }

    pub const fn contains_many(self, mask: ReqComps) -> bool {
        self & mask == mask
    }
}

impl const Into<ReqComps> for ReqComp {
    fn into(self) -> ReqComps {
        ReqComps::of(self)
    }
}

impl const core::cmp::PartialEq<ReqComp> for ReqComps {
    fn eq(&self, comp: &ReqComp) -> bool {
        self.0 & (*comp as u8) == *comp as u8
    }
}

impl const core::cmp::PartialEq for ReqComps {
    fn eq(&self, comp: &Self) -> bool {
        self.0 == comp.0
    }
}

impl const core::ops::BitOr for ReqComp {
    type Output = ReqComps;

    fn bitor(self, rhs: Self) -> Self::Output {
        ReqComps((self as u8) | (rhs as u8))
    }
}

impl const core::ops::BitOr for ReqComps {
    type Output = ReqComps;

    fn bitor(self, rhs: Self) -> Self::Output {
        ReqComps((self.0 as u8) | (rhs.0 as u8))
    }
}

impl const core::ops::BitOr<ReqComp> for ReqComps {
    type Output = ReqComps;

    fn bitor(self, rhs: ReqComp) -> Self::Output {
        ReqComps((self.0 as u8) | (rhs as u8))
    }
}

// TODO make `const` when is "stable" #57349 <https://github.com/rust-lang/rust/issues/57349>
impl core::ops::BitOrAssign for ReqComps {
    fn bitor_assign(&mut self, rhs: ReqComps) {
        self.0 = (self.0 as u8) | (rhs.0 as u8)
    }
}

// TODO make `const` when is "stable" #57349 <https://github.com/rust-lang/rust/issues/57349>
impl core::ops::BitOrAssign<ReqComp> for ReqComps {
    fn bitor_assign(&mut self, rhs: ReqComp) {
        self.0 = (self.0 as u8) | (rhs as u8)
    }
}

impl const core::ops::BitAnd for ReqComp {
    type Output = ReqComps;

    fn bitand(self, rhs: Self) -> Self::Output {
        ReqComps((self as u8) & (rhs as u8))
    }
}

impl const core::ops::BitAnd for ReqComps {
    type Output = ReqComps;

    fn bitand(self, rhs: Self) -> Self::Output {
        ReqComps((self.0 as u8) & (rhs.0 as u8))
    }
}

impl const core::ops::BitAnd<ReqComp> for ReqComps {
    type Output = ReqComps;

    fn bitand(self, rhs: ReqComp) -> Self::Output {
        ReqComps((self.0 as u8) & (rhs as u8))
    }
}

// TODO make `const` when is "stable" #57349 <https://github.com/rust-lang/rust/issues/57349>
impl core::ops::BitAndAssign for ReqComps {
    fn bitand_assign(&mut self, rhs: ReqComps) {
        self.0 = (self.0 as u8) & (rhs.0 as u8)
    }
}

// TODO make `const` when is "stable" #57349 <https://github.com/rust-lang/rust/issues/57349>
impl core::ops::BitAndAssign<ReqComp> for ReqComps {
    fn bitand_assign(&mut self, rhs: ReqComp) {
        self.0 = (self.0 as u8) & (rhs as u8)
    }
}

#[derive(Debug, Clone)]
pub struct LogMiddleware {
    pub components: ReqComps,
}

impl Default for LogMiddleware {
    fn default() -> Self {
        Self {
            components: ReqComp::Method | ReqComp::Path,
        }
    }
}

impl LogMiddleware {
    async fn log<'a, State: Clone + Send + Sync + 'static>(
        &'a self,
        req: Request<State>,
        next: Next<'a, State>,
    ) -> tide::Result {
        let id = Nrid::now();
        let span = debug_span!("Request", id = %id);
        let comp = self.components;
        let url = req.url();
        if comp.contains(ReqComp::Method) {
            span.record("http.method", &tracing::field::display(req.method()));
        }
        if comp.contains(ReqComp::Path) {
            span.record("http.path", &url.path());
        }
        if comp.contains(ReqComp::IP) {
            span.record("conn.ip", &tracing::field::display(req.remote().unwrap_or("unknown").to_string()));
        }
        if comp.contains(ReqComp::Query) {
            span.record("http.query", &url.query().unwrap_or(""));
        }
        Ok(async {
            let start = Instant::now();
            let res = next.run(req).await;
            let duration = start.elapsed();
            let status = res.status();
            debug_span!("Response", status = %status, duration = ?duration).in_scope(|| {});
            res
        }
        .instrument(span)
        .await)
    }
}

#[async_trait::async_trait]
impl<State: Clone + Send + Sync + 'static> Middleware<State> for LogMiddleware {
    async fn handle(&self, req: Request<State>, next: Next<'_, State>) -> tide::Result {
        self.log(req, next).await
    }
}