sentinel_tonic/
lib.rs

1//! This crate provides the [sentinel](https://docs.rs/sentinel-core) middleware for [tonic](https://docs.rs/tonic).
2//! The are two kinds of middlewares in tonic.
3//! - [`SentinelInterceptor`] based on [`tonic::service::interceptor::Interceptor`](https://docs.rs/tonic/latest/tonic/service/interceptor/trait.Interceptor.html)
4//! - [`SentinelService`] and [`SentinelLayer`] based on [`tower::Service`](https://docs.rs/tower/latest/tower/trait.Service.html)
5
6use sentinel_core::EntryBuilder;
7pub use sentinel_tower::{SentinelLayer, SentinelService, ServiceRole};
8use tonic::service::interceptor::Interceptor;
9use tonic::{Request, Status};
10
11type Extractor = fn(&Request<()>) -> String;
12type Fallback = fn(&Request<()>, &sentinel_core::Error) -> Result<Request<()>, Status>;
13
14#[derive(Clone)]
15pub struct SentinelInterceptor {
16    traffic_type: sentinel_core::base::TrafficType,
17    extractor: Option<Extractor>,
18    fallback: Option<Fallback>,
19}
20
21impl SentinelInterceptor {
22    pub fn new(role: ServiceRole) -> Self {
23        Self {
24            extractor: None,
25            fallback: None,
26            traffic_type: {
27                match role {
28                    ServiceRole::Server => sentinel_core::base::TrafficType::Inbound,
29                    ServiceRole::Client => sentinel_core::base::TrafficType::Outbound,
30                }
31            },
32        }
33    }
34
35    pub fn with_extractor(mut self, extractor: Extractor) -> Self {
36        self.extractor = Some(extractor);
37        self
38    }
39
40    pub fn with_fallback(mut self, fallback: Fallback) -> Self {
41        self.fallback = Some(fallback);
42        self
43    }
44}
45
46impl Interceptor for SentinelInterceptor {
47    fn call(&mut self, request: Request<()>) -> Result<Request<()>, Status> {
48        let resource = match self.extractor {
49            Some(extractor) => extractor(&request),
50            None => format!("{:?}", request.metadata()),
51        };
52        let entry_builder = EntryBuilder::new(resource).with_traffic_type(self.traffic_type);
53        match entry_builder.build() {
54            Ok(entry) => {
55                entry.exit();
56                Ok(request)
57            }
58            Err(err) => match self.fallback {
59                Some(fallback) => fallback(&request, &err),
60                None => Err(Status::resource_exhausted(format!(
61                    "Blocked by Sentinel: {:?}",
62                    err
63                ))),
64            },
65        }
66    }
67}