Repository navigation
Expand file tree
/
Copy pathstate.rs
More file actions
167 lines (138 loc) · 4.29 KB
/
Copy pathstate.rs
File metadata and controls
167 lines (138 loc) · 4.29 KB
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
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
//! Network proxy state and metrics.
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
/// Metrics for network proxy requests.
#[derive(Debug, Default)]
pub struct RequestMetrics {
/// Total number of requests.
pub total_requests: AtomicU64,
/// Number of allowed requests.
pub allowed_requests: AtomicU64,
/// Number of blocked requests.
pub blocked_requests: AtomicU64,
/// Total bytes transferred (approximate).
pub bytes_transferred: AtomicU64,
}
impl RequestMetrics {
/// Create new metrics.
pub fn new() -> Self {
Self::default()
}
/// Record an allowed request.
pub fn record_allowed(&self, bytes: u64) {
self.total_requests.fetch_add(1, Ordering::Relaxed);
self.allowed_requests.fetch_add(1, Ordering::Relaxed);
self.bytes_transferred.fetch_add(bytes, Ordering::Relaxed);
}
/// Record a blocked request.
pub fn record_blocked(&self) {
self.total_requests.fetch_add(1, Ordering::Relaxed);
self.blocked_requests.fetch_add(1, Ordering::Relaxed);
}
/// Get total requests.
pub fn total(&self) -> u64 {
self.total_requests.load(Ordering::Relaxed)
}
/// Get allowed requests count.
pub fn allowed(&self) -> u64 {
self.allowed_requests.load(Ordering::Relaxed)
}
/// Get blocked requests count.
pub fn blocked(&self) -> u64 {
self.blocked_requests.load(Ordering::Relaxed)
}
/// Get bytes transferred.
pub fn bytes(&self) -> u64 {
self.bytes_transferred.load(Ordering::Relaxed)
}
/// Get a snapshot of the metrics.
pub fn snapshot(&self) -> MetricsSnapshot {
MetricsSnapshot {
total_requests: self.total(),
allowed_requests: self.allowed(),
blocked_requests: self.blocked(),
bytes_transferred: self.bytes(),
}
}
}
/// Snapshot of metrics at a point in time.
#[derive(Debug, Clone)]
pub struct MetricsSnapshot {
pub total_requests: u64,
pub allowed_requests: u64,
pub blocked_requests: u64,
pub bytes_transferred: u64,
}
/// State for the network proxy.
pub struct NetworkProxyState {
/// Policy engine.
policy: super::PolicyEngine,
/// Request metrics.
metrics: Arc<RequestMetrics>,
/// Whether the proxy is active.
active: std::sync::atomic::AtomicBool,
}
impl NetworkProxyState {
/// Create new proxy state.
pub fn new(policy: super::PolicyEngine) -> Self {
Self {
policy,
metrics: Arc::new(RequestMetrics::new()),
active: std::sync::atomic::AtomicBool::new(true),
}
}
/// Get the policy engine.
pub fn policy(&self) -> &super::PolicyEngine {
&self.policy
}
/// Get the metrics.
pub fn metrics(&self) -> &Arc<RequestMetrics> {
&self.metrics
}
/// Check if the proxy is active.
pub fn is_active(&self) -> bool {
self.active.load(Ordering::Acquire)
}
/// Set the proxy active state.
pub fn set_active(&self, active: bool) {
self.active.store(active, Ordering::Release);
}
/// Check if a host is blocked and record metrics.
pub async fn check_and_record(&self, method: &str, host: &str, port: u16) -> super::Result<()> {
if !self.is_active() {
self.metrics.record_blocked();
return Err(super::NetworkProxyError::Internal(
"Proxy is not active".to_string(),
));
}
match self.policy.validate_request(method, host, port).await {
Ok(()) => {
// Will record bytes after transfer
Ok(())
}
Err(e) => {
self.metrics.record_blocked();
Err(e)
}
}
}
/// Record successful transfer.
pub fn record_success(&self, bytes: u64) {
self.metrics.record_allowed(bytes);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_metrics() {
let metrics = RequestMetrics::new();
metrics.record_allowed(1000);
metrics.record_allowed(500);
metrics.record_blocked();
assert_eq!(metrics.total(), 3);
assert_eq!(metrics.allowed(), 2);
assert_eq!(metrics.blocked(), 1);
assert_eq!(metrics.bytes(), 1500);
}
}