1#![doc(html_logo_url = "https://avatars.githubusercontent.com/u/43955412")]
2use 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
15pub type Extractor = fn(&ServiceRequest) -> String;
18
19pub type Fallback<B> =
21 fn(ServiceRequest, &sentinel_core::Error) -> Result<ServiceResponse<EitherBody<B>>, Error>;
22
23pub struct Sentinel<B> {
25 extractor: Option<Extractor>,
26 fallback: Option<Fallback<B>>,
27}
28
29impl<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
73pub 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}