use actix_web::{ dev::{Service, ServiceRequest, Transform}, http::StatusCode, HttpResponse, ResponseError, }; use futures_util::future::LocalBoxFuture; use std::{ future::{ready, Ready}, 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()) .content_type("application/json") .body( serde_json::to_string(&serde_json::json!({ "msg": self.to_string() })) .unwrap_or(r#"{"msg":"unauthorized"}"#.to_string()), ) } } impl Transform for Tracing where S: Service, S::Future: 'static, { type Response = S::Response; type Error = S::Error; type InitError = (); type Transform = TracingMiddleware; type Future = Ready>; fn new_transform(&self, service: S) -> Self::Future { ready(Ok(TracingMiddleware { inner: service })) } } impl Service for TracingMiddleware where S: Service, S::Future: 'static, { type Response = S::Response; type Error = S::Error; type Future = Instrumented; fn poll_ready(&self, cx: &mut Context<'_>) -> Poll> { self.inner.poll_ready(cx) } fn call(&self, req: 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 Response = S::Response; type Error = S::Error; type InitError = (); type Transform = InternalMiddleware; type Future = Ready>; fn new_transform(&self, service: S) -> Self::Future { ready(Ok(InternalMiddleware(self.0.clone(), service))) } } impl Service for InternalMiddleware where S: Service, S::Future: 'static, { type Response = S::Response; type Error = S::Error; type Future = LocalBoxFuture<'static, Result>; fn poll_ready(&self, cx: &mut Context<'_>) -> Poll> { self.1.poll_ready(cx) } fn call(&self, req: ServiceRequest) -> 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_deref() { let fut = self.1.call(req); return Box::pin(async move { fut.await }); } } Box::pin(async move { Err(ApiError.into()) }) } }