use actix_web::{ dev::{Service, ServiceRequest, Transform}, http::StatusCode, HttpResponse, ResponseError, }; use futures::future::{ok, LocalBoxFuture, Ready}; use std::task::{Context, Poll}; use tracing_futures::{Instrument, Instrumented}; use uuid::Uuid; pub(crate) struct Tracing; pub(crate) struct TracingMiddleware { inner: S, } pub(crate) struct Internal(pub(crate) Option); pub(crate) struct InternalMiddleware(Option, S); #[derive(Clone, Debug, thiserror::Error)] #[error("Invalid API Key")] struct ApiError; impl ResponseError for ApiError { fn status_code(&self) -> StatusCode { StatusCode::UNAUTHORIZED } fn error_response(&self) -> HttpResponse { HttpResponse::build(self.status_code()).json(serde_json::json!({ "msg": self.to_string() })) } } impl Transform for Tracing where S: Service, S::Future: 'static, { type Request = S::Request; type Response = S::Response; type Error = S::Error; type InitError = (); type Transform = TracingMiddleware; type Future = Ready>; fn new_transform(&self, service: S) -> Self::Future { ok(TracingMiddleware { inner: service }) } } impl Service for TracingMiddleware where S: Service, S::Future: 'static, { type Request = S::Request; type Response = S::Response; type Error = S::Error; type Future = Instrumented; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(cx) } fn call(&mut self, req: S::Request) -> Self::Future { let uuid = Uuid::new_v4(); self.inner .call(req) .instrument(tracing::info_span!("request", ?uuid)) } } impl Transform for Internal where S: Service, S::Future: 'static, { type Request = S::Request; type Response = S::Response; type Error = S::Error; type InitError = (); type Transform = InternalMiddleware; type Future = Ready>; fn new_transform(&self, service: S) -> Self::Future { ok(InternalMiddleware(self.0.clone(), service)) } } impl Service for InternalMiddleware where S: Service, S::Future: 'static, { type Request = S::Request; type Response = S::Response; type Error = S::Error; type Future = LocalBoxFuture<'static, Result>; fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll> { self.1.poll_ready(cx) } fn call(&mut self, req: S::Request) -> Self::Future { if let Some(value) = req.headers().get("x-api-token") { if value.to_str().is_ok() && value.to_str().ok() == self.0.as_ref().map(|s| s.as_str()) { let fut = self.1.call(req); return Box::pin(async move { fut.await }); } } Box::pin(async move { Err(ApiError.into()) }) } }