1use 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}