sentinel_core/core/base/
context.rs

1//! Context
2//!
3use super::{EntryWeakPtr, ResourceWrapper, StatNode, TokenResult};
4use crate::utils::time::curr_time_millis;
5use crate::Error;
6use std::collections::HashMap;
7use std::sync::Arc;
8
9cfg_async! {
10    use std::sync::RwLock;
11    pub type ContextPtr = Arc<RwLock<EntryContext>>;
12}
13
14cfg_not_async! {
15    use std::rc::Rc;
16    use std::cell::RefCell;
17    pub type ContextPtr = Rc<RefCell<EntryContext>>;
18}
19
20#[derive(Default)]
21pub struct EntryContext {
22    /// entry and context do not need to be `Send/Sync`
23    /// entry<->context, cycled reference, so need Weak
24    /// context should not change entry, so here we do not use RefCell
25    entry: Option<EntryWeakPtr>,
26    /// Use to calculate RT
27    start_time: u64,
28    /// The round trip time of this transaction
29    round_trip: u64,
30    resource: ResourceWrapper,
31    // todo: is it necessary to keep using trait object here?
32    // consider replacing by `crate::core::stat::ResourceNode`
33    stat_node: Option<Arc<dyn StatNode>>,
34    input: SentinelInput,
35    /// the result of rule slots check
36    rule_check_result: TokenResult,
37    err: Option<Error>,
38}
39
40impl EntryContext {
41    pub fn new() -> Self {
42        EntryContext {
43            start_time: curr_time_millis(),
44            ..Default::default()
45        }
46    }
47
48    pub fn set_entry(&mut self, entry: EntryWeakPtr) {
49        self.entry = Some(entry);
50    }
51
52    pub fn entry(&self) -> Option<&EntryWeakPtr> {
53        self.entry.as_ref()
54    }
55
56    pub fn start_time(&self) -> u64 {
57        self.start_time
58    }
59
60    pub fn is_blocked(&self) -> bool {
61        self.rule_check_result.is_blocked()
62    }
63
64    pub fn set_round_trip(&mut self, round_trip: u64) {
65        self.round_trip = round_trip
66    }
67
68    pub fn round_trip(&self) -> u64 {
69        self.round_trip
70    }
71
72    pub fn set_resource(&mut self, resource: ResourceWrapper) {
73        self.resource = resource;
74    }
75
76    pub fn resource(&self) -> &ResourceWrapper {
77        &self.resource
78    }
79
80    pub fn set_input(&mut self, input: SentinelInput) {
81        self.input = input;
82    }
83
84    pub fn input(&self) -> &SentinelInput {
85        &self.input
86    }
87
88    pub fn set_stat_node(&mut self, stat_node: Arc<dyn StatNode>) {
89        self.stat_node = Some(stat_node);
90    }
91
92    pub fn stat_node(&self) -> Option<Arc<dyn StatNode>> {
93        self.stat_node.clone()
94    }
95
96    pub fn set_result(&mut self, result: TokenResult) {
97        self.rule_check_result = result;
98    }
99
100    pub fn reset_result_to_pass(&mut self) {
101        self.rule_check_result.reset_to_pass();
102    }
103
104    pub fn result(&self) -> &TokenResult {
105        &self.rule_check_result
106    }
107
108    pub fn set_err(&mut self, err: Error) {
109        self.err = Some(err);
110    }
111
112    pub fn get_err(&self) -> &Option<Error> {
113        &self.err
114    }
115}
116
117pub type ParamKey = String;
118pub type ParamsList = Vec<ParamKey>;
119pub type ParamsMap = HashMap<String, ParamKey>;
120
121/// Input of policy algorithms
122#[derive(Debug)]
123pub struct SentinelInput {
124    batch_count: u32,
125    flag: i32,
126    /// following input items are used in hotspot module
127    args: Option<ParamsList>,
128    attachments: Option<ParamsMap>,
129}
130
131impl Default for SentinelInput {
132    fn default() -> Self {
133        SentinelInput {
134            batch_count: 1,
135            flag: 0,
136            args: None,
137            attachments: None,
138        }
139    }
140}
141
142impl SentinelInput {
143    pub fn new(batch_count: u32, flag: i32) -> Self {
144        SentinelInput {
145            batch_count,
146            flag,
147            ..Default::default()
148        }
149    }
150
151    pub fn set_batch_count(&mut self, batch_count: u32) {
152        self.batch_count = batch_count;
153    }
154
155    pub fn batch_count(&self) -> u32 {
156        self.batch_count
157    }
158
159    pub fn set_flag(&mut self, flag: i32) {
160        self.flag = flag;
161    }
162
163    pub fn flag(&self) -> i32 {
164        self.flag
165    }
166
167    pub fn set_args(&mut self, args: ParamsList) {
168        self.args = Some(args);
169    }
170
171    pub fn args(&self) -> Option<&ParamsList> {
172        self.args.as_ref()
173    }
174
175    pub fn set_attachments(&mut self, attachments: ParamsMap) {
176        self.attachments = Some(attachments);
177    }
178
179    pub fn attachments(&self) -> Option<&ParamsMap> {
180        self.attachments.as_ref()
181    }
182}
183
184#[cfg(test)]
185mod test {
186    use super::*;
187    use crate::base::result::BlockType;
188    #[test]
189    fn is_blocked() {
190        let mut ctx = EntryContext::new();
191        assert_eq!(ctx.is_blocked(), false);
192        ctx.set_result(TokenResult::new_blocked(BlockType::Other(1)));
193        assert_eq!(ctx.is_blocked(), true);
194    }
195}
sentinel_core/core/base/context.rs

sentinel_core/core/base/
context.rs
