sentinel-tonic 0.1.0

Sentinel middleware for Tonic
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

//! This crate provides the [sentinel](https://docs.rs/sentinel-core) middleware for [tonic](https://docs.rs/tonic).
//! The are two kinds of middlewares in tonic.
//! - [`SentinelInterceptor`] based on [`tonic::service::interceptor::Interceptor`](https://docs.rs/tonic/latest/tonic/service/interceptor/trait.Interceptor.html)
//! - [`SentinelService`] and [`SentinelLayer`] based on [`tower::Service`](https://docs.rs/tower/latest/tower/trait.Service.html)

use sentinel_core::EntryBuilder;
pub use sentinel_tower::{SentinelLayer, SentinelService, ServiceRole};
use tonic::service::interceptor::Interceptor;
use tonic::{Request, Status};

type Extractor = fn(&Request<()>) -> String;
type Fallback = fn(&Request<()>, &sentinel_core::Error) -> Result<Request<()>, Status>;

#[derive(Clone)]
pub struct SentinelInterceptor {
    traffic_type: sentinel_core::base::TrafficType,
    extractor: Option<Extractor>,
    fallback: Option<Fallback>,
}

impl SentinelInterceptor {
    pub fn new(role: ServiceRole) -> Self {
        Self {
            extractor: None,
            fallback: None,
            traffic_type: {
                match role {
                    ServiceRole::Server => sentinel_core::base::TrafficType::Inbound,
                    ServiceRole::Client => sentinel_core::base::TrafficType::Outbound,
                }
            },
        }
    }

    pub fn with_extractor(mut self, extractor: Extractor) -> Self {
        self.extractor = Some(extractor);
        self
    }

    pub fn with_fallback(mut self, fallback: Fallback) -> Self {
        self.fallback = Some(fallback);
        self
    }
}

impl Interceptor for SentinelInterceptor {
    fn call(&mut self, request: Request<()>) -> Result<Request<()>, Status> {
        let resource = match self.extractor {
            Some(extractor) => extractor(&request),
            None => format!("{:?}", request.metadata()),
        };
        let entry_builder = EntryBuilder::new(resource).with_traffic_type(self.traffic_type);
        match entry_builder.build() {
            Ok(entry) => {
                entry.exit();
                Ok(request)
            }
            Err(err) => match self.fallback {
                Some(fallback) => fallback(&request, &err),
                None => Err(Status::resource_exhausted(format!(
                    "Blocked by Sentinel: {:?}",
                    err
                ))),
            },
        }
    }
}