sentinel_actix/
lib.rs

1#![doc(html_logo_url = "https://avatars.githubusercontent.com/u/43955412")]
2//! This crate provides the [sentinel](https://docs.rs/sentinel-core) middleware for [actix-web](https://docs.rs/actix-web).
3//! See [examples](https://github.com/sentinel-group/sentinel-rust/tree/main/middleware) for help.
4
5use actix_utils::future::{ok, Ready};
6use actix_web::{
7    body::EitherBody,
8    dev::{forward_ready, Service, ServiceRequest, ServiceResponse, Transform},
9    http::StatusCode,
10    Error, HttpResponse,
11};
12use sentinel_core::EntryBuilder;
13use std::{future::Future, pin::Pin, rc::Rc};
14
15/// It is used to extractor a resource name from requests for Sentinel. If you the service request is [`http::Request`](https://docs.rs/http/latest/http/request/struct.Request.html),
16/// and you are using nightly toolchain, you don't need to provide a sentinel resource name extractor. The middleware will automatically extract the request uri.
17pub type Extractor = fn(&ServiceRequest) -> String;
18
19/// The fallback function when service is rejected by sentinel.
20pub type Fallback<B> =
21    fn(ServiceRequest, &sentinel_core::Error) -> Result<ServiceResponse<EitherBody<B>>, Error>;
22
23/// Sentinel wrapper
24pub struct Sentinel<B> {
25    extractor: Option<Extractor>,
26    fallback: Option<Fallback<B>>,
27}
28
29// rustc cannot derive `Default` trait for function pointers correctly,
30// implement it by hands
31impl<B> Default for Sentinel<B> {
32    fn default() -> Self {
33        Self {
34            extractor: None,
35            fallback: None,
36        }
37    }
38}
39
40impl<B> Sentinel<B> {
41    pub fn with_extractor(mut self, extractor: Extractor) -> Self {
42        self.extractor = Some(extractor);
43        self
44    }
45
46    pub fn with_fallback(mut self, fallback: Fallback<B>) -> Self {
47        self.fallback = Some(fallback);
48        self
49    }
50}
51
52impl<S, B> Transform<S, ServiceRequest> for Sentinel<B>
53where
54    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
55    S::Future: 'static,
56    B: 'static,
57{
58    type Response = ServiceResponse<EitherBody<B>>;
59    type Error = Error;
60    type Transform = SentinelMiddleware<S, B>;
61    type InitError = ();
62    type Future = Ready<Result<Self::Transform, Self::InitError>>;
63
64    fn new_transform(&self, service: S) -> Self::Future {
65        ok(SentinelMiddleware {
66            service: Rc::new(service),
67            extractor: self.extractor.clone(),
68            fallback: self.fallback.clone(),
69        })
70    }
71}
72
73/// Sentinel middleware
74pub struct SentinelMiddleware<S, B> {
75    service: Rc<S>,
76    extractor: Option<Extractor>,
77    fallback: Option<Fallback<B>>,
78}
79
80impl<S, B> Service<ServiceRequest> for SentinelMiddleware<S, B>
81where
82    S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
83    S::Future: 'static,
84    B: 'static,
85{
86    type Response = ServiceResponse<EitherBody<B>>;
87    type Error = Error;
88    type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>>>>;
89
90    forward_ready!(service);
91
92    fn call(&self, req: ServiceRequest) -> Self::Future {
93        let resource = match self.extractor {
94            Some(extractor) => extractor(&req),
95            None => req.uri().to_string(),
96        };
97        let service = Rc::clone(&self.service);
98        let entry_builder = EntryBuilder::new(resource)
99            .with_traffic_type(sentinel_core::base::TrafficType::Inbound);
100
101        match entry_builder.build() {
102            Ok(entry) => Box::pin(async move {
103                let response = service
104                    .call(req)
105                    .await
106                    .map(ServiceResponse::map_into_left_body);
107                entry.exit();
108                response
109            }),
110            Err(err) => match self.fallback {
111                Some(fallback) => Box::pin(async move { fallback(req, &err) }),
112                None => Box::pin(async move {
113                    Ok(req.into_response(
114                        HttpResponse::new(StatusCode::TOO_MANY_REQUESTS).map_into_right_body(),
115                    ))
116                }),
117            },
118        }
119    }
120}