mirror of https://github.com/ospab/ostp.git
591 lines
24 KiB
Rust
591 lines
24 KiB
Rust
use anyhow::Result;
|
||
use bytes::Bytes;
|
||
use ostp_core::{OstpEvent, ProtocolAction, ProtocolConfig, ProtocolMachine};
|
||
use std::collections::HashMap;
|
||
use std::net::SocketAddr;
|
||
use std::sync::{Arc, RwLock};
|
||
use std::sync::atomic::Ordering;
|
||
use portable_atomic::AtomicU64;
|
||
|
||
/// Maximum number of concurrent authenticated sessions.
|
||
/// Excess handshake attempts are silently dropped -- no response, no state allocated.
|
||
// const MAX_SESSIONS removed because dynamic limit is used
|
||
|
||
pub enum DispatchOutcome {
|
||
Accepted {
|
||
responses: Vec<Bytes>,
|
||
app_payloads: Vec<(u32, u16, Bytes)>, // session_id, stream_id, payload
|
||
peer_addr: SocketAddr,
|
||
},
|
||
Unauthorized(String),
|
||
Ignored,
|
||
}
|
||
|
||
/// Per-user traffic statistics.
|
||
pub struct UserStats {
|
||
pub bytes_up: AtomicU64,
|
||
pub bytes_down: AtomicU64,
|
||
pub connections: AtomicU64,
|
||
pub limit_bytes: Option<u64>,
|
||
pub created_at: std::time::SystemTime,
|
||
}
|
||
|
||
impl UserStats {
|
||
pub fn new(limit: Option<u64>) -> Self {
|
||
Self {
|
||
bytes_up: AtomicU64::new(0),
|
||
bytes_down: AtomicU64::new(0),
|
||
connections: AtomicU64::new(0),
|
||
limit_bytes: limit,
|
||
created_at: std::time::SystemTime::now(),
|
||
}
|
||
}
|
||
|
||
pub fn is_over_limit(&self) -> bool {
|
||
if let Some(limit) = self.limit_bytes {
|
||
let total = self.bytes_up.load(Ordering::Relaxed)
|
||
+ self.bytes_down.load(Ordering::Relaxed);
|
||
total >= limit
|
||
} else {
|
||
false
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Snapshot of user stats for API responses.
|
||
#[derive(Debug, Clone, serde::Serialize)]
|
||
pub struct UserStatsSnapshot {
|
||
pub access_key: String,
|
||
pub name: Option<String>,
|
||
pub bytes_up: u64,
|
||
pub bytes_down: u64,
|
||
pub connections: u64,
|
||
pub limit_bytes: Option<u64>,
|
||
pub online: bool,
|
||
pub last_seen: Option<u64>,
|
||
}
|
||
|
||
pub struct PeerState {
|
||
pub machine: ProtocolMachine,
|
||
pub last_addr: SocketAddr,
|
||
pub obfuscation_key: [u8; 8],
|
||
pub last_seen: std::time::Instant,
|
||
pub access_key: String,
|
||
}
|
||
|
||
pub struct Dispatcher {
|
||
peer_machines: HashMap<u32, PeerState>,
|
||
addr_to_session: HashMap<SocketAddr, u32>,
|
||
machine_config: ProtocolConfig,
|
||
access_keys: Arc<RwLock<HashMap<String, crate::api::UserMeta>>>,
|
||
user_stats: Arc<RwLock<HashMap<String, Arc<UserStats>>>>,
|
||
replay_cache: std::collections::HashMap<Vec<u8>, u64>,
|
||
roaming_tokens: f64,
|
||
last_token_regen: std::time::Instant,
|
||
}
|
||
|
||
impl Dispatcher {
|
||
pub fn new(machine_config: ProtocolConfig, access_keys: Arc<RwLock<HashMap<String, crate::api::UserMeta>>>) -> Self {
|
||
let mut initial_stats = HashMap::new();
|
||
for (key, meta) in access_keys.read().unwrap_or_else(|e| e.into_inner()).iter() {
|
||
initial_stats.insert(key.clone(), Arc::new(UserStats::new(meta.limit_bytes)));
|
||
}
|
||
Self {
|
||
peer_machines: HashMap::new(),
|
||
addr_to_session: HashMap::new(),
|
||
machine_config,
|
||
access_keys,
|
||
user_stats: Arc::new(RwLock::new(initial_stats)),
|
||
replay_cache: std::collections::HashMap::new(),
|
||
roaming_tokens: 50.0,
|
||
last_token_regen: std::time::Instant::now(),
|
||
}
|
||
}
|
||
|
||
/// Returns a shared reference to user stats for the Management API.
|
||
pub fn user_stats_ref(&self) -> Arc<RwLock<HashMap<String, Arc<UserStats>>>> {
|
||
self.user_stats.clone()
|
||
}
|
||
|
||
/// Snapshot all user stats for API responses.
|
||
#[allow(dead_code)]
|
||
pub fn snapshot_all_users(&self) -> Vec<UserStatsSnapshot> {
|
||
let stats = self.user_stats.read().unwrap_or_else(|e| e.into_inner());
|
||
let mut online_keys: HashMap<String, std::time::Instant> = HashMap::new();
|
||
for ps in self.peer_machines.values() {
|
||
let key = ps.access_key.clone();
|
||
if let Some(existing) = online_keys.get(&key) {
|
||
if ps.last_seen > *existing {
|
||
online_keys.insert(key, ps.last_seen);
|
||
}
|
||
} else {
|
||
online_keys.insert(key, ps.last_seen);
|
||
}
|
||
}
|
||
|
||
let now = std::time::Instant::now();
|
||
let current_sys_time = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_secs();
|
||
|
||
stats.iter().map(|(key, us)| {
|
||
let last_seen_unix = online_keys.get(key).map(|&instant| {
|
||
let diff = now.duration_since(instant).as_secs();
|
||
current_sys_time.saturating_sub(diff)
|
||
});
|
||
|
||
UserStatsSnapshot {
|
||
access_key: key.clone(),
|
||
name: None,
|
||
bytes_up: us.bytes_up.load(Ordering::Relaxed),
|
||
bytes_down: us.bytes_down.load(Ordering::Relaxed),
|
||
connections: us.connections.load(Ordering::Relaxed),
|
||
limit_bytes: us.limit_bytes,
|
||
online: online_keys.contains_key(key),
|
||
last_seen: last_seen_unix,
|
||
}
|
||
}).collect()
|
||
}
|
||
|
||
/// Get or create stats entry for a user key.
|
||
fn get_or_create_user_stats(&self, key: &str) -> Arc<UserStats> {
|
||
let stats = self.user_stats.read().unwrap_or_else(|e| e.into_inner());
|
||
if let Some(existing) = stats.get(key) {
|
||
return existing.clone();
|
||
}
|
||
drop(stats);
|
||
|
||
let limit_bytes = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).get(key).and_then(|m| m.limit_bytes);
|
||
|
||
let mut stats = self.user_stats.write().unwrap_or_else(|e| e.into_inner());
|
||
stats.entry(key.to_string())
|
||
.or_insert_with(|| Arc::new(UserStats::new(limit_bytes)))
|
||
.clone()
|
||
}
|
||
|
||
/// Set traffic limit for a user.
|
||
#[allow(dead_code)]
|
||
pub fn set_user_limit(&self, key: &str, limit: Option<u64>) {
|
||
let mut stats = self.user_stats.write().unwrap_or_else(|e| e.into_inner());
|
||
let entry = stats.entry(key.to_string())
|
||
.or_insert_with(|| Arc::new(UserStats::new(limit)));
|
||
// Replace the entry with new limit (stats reset)
|
||
*entry = Arc::new(UserStats {
|
||
bytes_up: AtomicU64::new(entry.bytes_up.load(Ordering::Relaxed)),
|
||
bytes_down: AtomicU64::new(entry.bytes_down.load(Ordering::Relaxed)),
|
||
connections: AtomicU64::new(entry.connections.load(Ordering::Relaxed)),
|
||
limit_bytes: limit,
|
||
created_at: entry.created_at,
|
||
});
|
||
}
|
||
|
||
/// Active session count.
|
||
#[allow(dead_code)]
|
||
pub fn active_sessions(&self) -> usize {
|
||
self.peer_machines.len()
|
||
}
|
||
|
||
pub fn on_datagram(&mut self, peer: SocketAddr, packet: Bytes) -> Result<DispatchOutcome> {
|
||
if packet.len() < 4 {
|
||
return Ok(DispatchOutcome::Unauthorized("packet too short".to_string()));
|
||
}
|
||
|
||
let mut session_id_opt = None;
|
||
|
||
if let Some(&sid) = self.addr_to_session.get(&peer) {
|
||
if let Some(peer_state) = self.peer_machines.get(&sid) {
|
||
let mut header = [0u8; 12];
|
||
if packet.len() >= 12 {
|
||
header.copy_from_slice(&packet[0..12]);
|
||
let ciphertext = &packet[12..];
|
||
ostp_core::crypto::deobfuscate_header_inplace(&mut header, ciphertext, &peer_state.obfuscation_key, false);
|
||
let candidate_sid = u32::from_be_bytes([header[0], header[1], header[2], header[3]]);
|
||
if candidate_sid == sid {
|
||
session_id_opt = Some(sid);
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
if session_id_opt.is_none() {
|
||
// Token Bucket rate limiter: mitigate seamless roaming CPU DoS vector
|
||
let now = std::time::Instant::now();
|
||
let elapsed = now.duration_since(self.last_token_regen).as_secs_f64();
|
||
self.last_token_regen = now;
|
||
self.roaming_tokens = (self.roaming_tokens + elapsed * 50.0).min(50.0);
|
||
|
||
if self.roaming_tokens >= 1.0 {
|
||
self.roaming_tokens -= 1.0;
|
||
|
||
// Try seamless roaming over all peers
|
||
for (&sid, peer_state) in &self.peer_machines {
|
||
if packet.len() >= 12 {
|
||
let mut header = [0u8; 12];
|
||
header.copy_from_slice(&packet[0..12]);
|
||
let ciphertext = &packet[12..];
|
||
ostp_core::crypto::deobfuscate_header_inplace(&mut header, ciphertext, &peer_state.obfuscation_key, false);
|
||
let candidate_sid = u32::from_be_bytes([header[0], header[1], header[2], header[3]]);
|
||
if candidate_sid == sid {
|
||
session_id_opt = Some(sid);
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
if let Some(session_id) = session_id_opt {
|
||
let key_opt = self.peer_machines.get(&session_id).map(|ps| ps.access_key.clone());
|
||
if let Some(access_key) = key_opt {
|
||
// Check if key is still valid and not over limit
|
||
let key_valid = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).contains_key(&access_key);
|
||
let user_stats = self.get_or_create_user_stats(&access_key);
|
||
if !key_valid || user_stats.is_over_limit() {
|
||
tracing::info!("Dropping session {} for key {} (valid={}, over_limit={})",
|
||
session_id, access_key, key_valid, user_stats.is_over_limit());
|
||
self.drop_session(session_id);
|
||
return Ok(DispatchOutcome::Unauthorized("key invalid or over limit".to_string()));
|
||
}
|
||
}
|
||
|
||
if let Some(peer_state) = self.peer_machines.get_mut(&session_id) {
|
||
|
||
if peer_state.last_addr != peer {
|
||
tracing::info!("Client roamed: session {} from {} to {}", session_id, peer_state.last_addr, peer);
|
||
self.addr_to_session.remove(&peer_state.last_addr);
|
||
}
|
||
peer_state.last_addr = peer;
|
||
peer_state.last_seen = std::time::Instant::now();
|
||
self.addr_to_session.insert(peer, session_id);
|
||
|
||
// Track inbound bytes per user
|
||
let key = peer_state.access_key.clone();
|
||
track_user_bytes_up(&self.user_stats, &self.access_keys, &key, packet.len() as u64);
|
||
|
||
let action = match peer_state.machine.on_event(OstpEvent::Inbound(packet)) {
|
||
Ok(a) => a,
|
||
Err(e) => {
|
||
return Ok(DispatchOutcome::Unauthorized(format!("protocol error: {}", e)));
|
||
}
|
||
};
|
||
|
||
let mut responses = Vec::new();
|
||
let mut app_payloads = Vec::new();
|
||
|
||
fn collect_action(
|
||
act: ProtocolAction,
|
||
sid: u32,
|
||
resps: &mut Vec<Bytes>,
|
||
loads: &mut Vec<(u32, u16, Bytes)>,
|
||
) {
|
||
match act {
|
||
ProtocolAction::SendDatagram(frame) => {
|
||
resps.push(frame);
|
||
}
|
||
ProtocolAction::DeliverApp(stream_id, data) => {
|
||
loads.push((sid, stream_id, data));
|
||
}
|
||
ProtocolAction::Multiple(list) => {
|
||
for item in list {
|
||
collect_action(item, sid, resps, loads);
|
||
}
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
collect_action(action, session_id, &mut responses, &mut app_payloads);
|
||
|
||
return Ok(DispatchOutcome::Accepted {
|
||
responses,
|
||
app_payloads,
|
||
peer_addr: peer,
|
||
});
|
||
}
|
||
}
|
||
|
||
// Not an existing session — try each registered access key's derived obfuscation key
|
||
let keys_snapshot: Vec<String> = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).keys().cloned().collect();
|
||
|
||
let mut failed_trials = Vec::new();
|
||
for candidate_key in keys_snapshot {
|
||
let secrets = ostp_core::crypto::derive_all_secrets(candidate_key.as_bytes());
|
||
|
||
// Decode the session_id using this key's obfuscation
|
||
// The handshake mask is derived from the Noise payload at bytes [6..],
|
||
// so we must deobfuscate the full packet, not just the header.
|
||
if packet.len() < 7 {
|
||
failed_trials.push(format!("key {}: packet too short", candidate_key));
|
||
continue;
|
||
}
|
||
let mut trial = packet.to_vec();
|
||
ostp_core::crypto::deobfuscate_packet_inplace(&mut trial, &secrets.obfuscation_key, true);
|
||
let candidate_session_id = u32::from_be_bytes([trial[0], trial[1], trial[2], trial[3]]);
|
||
|
||
let mut cfg = self.machine_config.clone();
|
||
cfg.session_id = candidate_session_id;
|
||
cfg.psk = secrets.psk;
|
||
cfg.handshake_payload = vec![];
|
||
cfg.obfuscation_key = secrets.obfuscation_key;
|
||
cfg.handshake_pad_min = secrets.handshake_pad_min;
|
||
cfg.handshake_pad_max = secrets.handshake_pad_max;
|
||
|
||
let mut machine = match ProtocolMachine::new(cfg) {
|
||
Ok(m) => m,
|
||
Err(e) => {
|
||
tracing::warn!("Failed to create protocol machine for key trial: {}", e);
|
||
continue;
|
||
}
|
||
};
|
||
let action = match machine.on_event(OstpEvent::Inbound(packet.clone())) {
|
||
Ok(a) => a,
|
||
Err(e) => {
|
||
failed_trials.push(format!("key {}: crypto err: {}", candidate_key, e));
|
||
continue;
|
||
}
|
||
};
|
||
|
||
if let ProtocolAction::HandshakePayload(payload, response_opt) = action {
|
||
if payload.len() >= 12 {
|
||
let mut ts_bytes = [0_u8; 8];
|
||
ts_bytes.copy_from_slice(&payload[..8]);
|
||
let ts = u64::from_be_bytes(ts_bytes);
|
||
|
||
let mut sid_bytes = [0_u8; 4];
|
||
sid_bytes.copy_from_slice(&payload[8..12]);
|
||
let sid_from_payload = u32::from_be_bytes(sid_bytes);
|
||
|
||
if sid_from_payload != candidate_session_id {
|
||
failed_trials.push(format!("key {}: sid mismatch", candidate_key));
|
||
continue;
|
||
}
|
||
|
||
let key_bytes = &payload[12..];
|
||
if let Ok(key_from_payload) = std::str::from_utf8(key_bytes) {
|
||
// The key embedded in the payload must match the candidate key we decoded with
|
||
if key_from_payload != candidate_key {
|
||
failed_trials.push(format!("key {}: embedded key mismatch", candidate_key));
|
||
continue;
|
||
}
|
||
|
||
let now = std::time::SystemTime::now()
|
||
.duration_since(std::time::UNIX_EPOCH)
|
||
.unwrap_or_default()
|
||
.as_secs();
|
||
|
||
let drift = (now as i64 - ts as i64).abs();
|
||
if drift > 300 {
|
||
let reason = format!("timestamp drift {}s exceeds 300s limit", drift);
|
||
tracing::warn!("Handshake rejected for {}: {}", peer, reason);
|
||
failed_trials.push(format!("key {}: {}", candidate_key, reason));
|
||
continue;
|
||
}
|
||
|
||
if self.replay_cache.contains_key(&payload.to_vec()) {
|
||
tracing::debug!("Replay detected from {}, ignoring", peer);
|
||
return Ok(DispatchOutcome::Ignored);
|
||
}
|
||
|
||
if self.replay_cache.len() >= 50_000 {
|
||
tracing::warn!("Replay cache full (50000 entries), rejecting handshake from {}", peer);
|
||
return Ok(DispatchOutcome::Unauthorized("replay cache full".to_string()));
|
||
}
|
||
|
||
self.replay_cache.insert(payload.to_vec(), ts);
|
||
|
||
machine.set_session_keys(candidate_session_id, secrets.obfuscation_key);
|
||
|
||
// Track per-user connection count
|
||
let user_stats = self.get_or_create_user_stats(&candidate_key);
|
||
user_stats.connections.fetch_add(1, Ordering::Relaxed);
|
||
|
||
// Check traffic limit before accepting
|
||
if user_stats.is_over_limit() {
|
||
tracing::warn!("User {} exceeded traffic limit, rejecting handshake from {}", candidate_key, peer);
|
||
return Ok(DispatchOutcome::Unauthorized("user over traffic limit".to_string()));
|
||
}
|
||
|
||
self.peer_machines.insert(candidate_session_id, PeerState {
|
||
machine,
|
||
last_addr: peer,
|
||
obfuscation_key: secrets.obfuscation_key,
|
||
last_seen: std::time::Instant::now(),
|
||
access_key: candidate_key.clone(),
|
||
});
|
||
self.addr_to_session.insert(peer, candidate_session_id);
|
||
|
||
tracing::debug!("New session authenticated: sid={} peer={} (active_sessions={}, replay_cache={})",
|
||
candidate_session_id, peer, self.peer_machines.len(), self.replay_cache.len()
|
||
);
|
||
|
||
return Ok(DispatchOutcome::Accepted {
|
||
responses: response_opt.into_iter().collect(),
|
||
app_payloads: Vec::new(),
|
||
peer_addr: peer,
|
||
});
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
let reason = if failed_trials.is_empty() {
|
||
"no valid handshake payload found".to_string()
|
||
} else {
|
||
format!("all key trials failed: {}", failed_trials.join(", "))
|
||
};
|
||
|
||
Ok(DispatchOutcome::Unauthorized(reason))
|
||
}
|
||
|
||
pub fn outbound_to_session(&mut self, session_id: u32, stream_id: u16, payload: Bytes) -> Result<Vec<(Bytes, SocketAddr)>> {
|
||
let peer_state = if let Some(existing) = self.peer_machines.get_mut(&session_id) {
|
||
existing
|
||
} else {
|
||
return Ok(Vec::new());
|
||
};
|
||
|
||
let addr = peer_state.last_addr;
|
||
let key = peer_state.access_key.clone();
|
||
let action = peer_state.machine.on_event(OstpEvent::Outbound(stream_id, payload))?;
|
||
|
||
let mut frames = Vec::new();
|
||
let mut queue = vec![action];
|
||
while let Some(current) = queue.pop() {
|
||
match current {
|
||
ProtocolAction::Multiple(list) => {
|
||
for item in list {
|
||
queue.push(item);
|
||
}
|
||
}
|
||
ProtocolAction::SendDatagram(frame) => {
|
||
track_user_bytes_down(&self.user_stats, &self.access_keys, &key, frame.len() as u64);
|
||
frames.push((frame, addr));
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
Ok(frames)
|
||
}
|
||
|
||
pub fn on_tick(&mut self) -> (Vec<(Bytes, SocketAddr)>, Vec<u32>) {
|
||
// Purge expired handshakes from replay cache (older than 5 min drift allowance)
|
||
let current_sys_time = std::time::SystemTime::now()
|
||
.duration_since(std::time::UNIX_EPOCH)
|
||
.unwrap_or_default()
|
||
.as_secs();
|
||
self.replay_cache.retain(|_, &mut ts| (current_sys_time as i64 - ts as i64).abs() <= 300);
|
||
|
||
let mut frames = Vec::new();
|
||
let mut expired = Vec::new();
|
||
let now = std::time::Instant::now();
|
||
let timeout_dur = std::time::Duration::from_secs(600); // 10-minute session timeout (mobile NAT mappings can live 5–10 min)
|
||
|
||
// Gather expired or invalid sessions
|
||
for (&sid, peer_state) in &self.peer_machines {
|
||
let key_valid = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).contains_key(&peer_state.access_key);
|
||
let user_stats = self.get_or_create_user_stats(&peer_state.access_key);
|
||
if now.duration_since(peer_state.last_seen) > timeout_dur || !key_valid || user_stats.is_over_limit() {
|
||
expired.push(sid);
|
||
}
|
||
}
|
||
|
||
// Clear expired/invalid sessions from internal state
|
||
for sid in &expired {
|
||
let peer_state_opt = self.peer_machines.get(sid);
|
||
let reason = if let Some(ps) = peer_state_opt {
|
||
let key_valid = self.access_keys.read().unwrap_or_else(|e| e.into_inner()).contains_key(&ps.access_key);
|
||
let user_stats = self.get_or_create_user_stats(&ps.access_key);
|
||
if now.duration_since(ps.last_seen) > timeout_dur {
|
||
"inactive >10min"
|
||
} else if !key_valid {
|
||
"key deleted"
|
||
} else if user_stats.is_over_limit() {
|
||
"traffic limit exceeded"
|
||
} else {
|
||
"unknown"
|
||
}
|
||
} else {
|
||
"unknown"
|
||
};
|
||
tracing::info!("Session {} closed ({}), releasing", sid, reason);
|
||
self.drop_session(*sid);
|
||
}
|
||
|
||
// Drive ticks for remaining active sessions
|
||
for peer_state in self.peer_machines.values_mut() {
|
||
let action = match peer_state.machine.on_event(OstpEvent::Tick) {
|
||
Ok(a) => a,
|
||
Err(e) => {
|
||
tracing::warn!("Tick error for session: {}", e);
|
||
continue;
|
||
}
|
||
};
|
||
|
||
let mut queue = vec![action];
|
||
while let Some(current) = queue.pop() {
|
||
match current {
|
||
ProtocolAction::Multiple(list) => {
|
||
for item in list {
|
||
queue.push(item);
|
||
}
|
||
}
|
||
ProtocolAction::SendDatagram(frame) => {
|
||
frames.push((frame, peer_state.last_addr));
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
}
|
||
|
||
(frames, expired)
|
||
}
|
||
|
||
pub fn drop_session(&mut self, session_id: u32) {
|
||
if let Some(state) = self.peer_machines.remove(&session_id) {
|
||
self.addr_to_session.remove(&state.last_addr);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Free functions to avoid borrow-checker conflicts when tracking stats
|
||
// while holding a mutable reference to peer_machines.
|
||
|
||
fn get_or_create_stats(
|
||
user_stats: &Arc<RwLock<HashMap<String, Arc<UserStats>>>>,
|
||
access_keys: &Arc<RwLock<HashMap<String, crate::api::UserMeta>>>,
|
||
key: &str,
|
||
) -> Arc<UserStats> {
|
||
{
|
||
let stats = user_stats.read().unwrap_or_else(|e| e.into_inner());
|
||
if let Some(existing) = stats.get(key) {
|
||
return existing.clone();
|
||
}
|
||
}
|
||
|
||
let limit_bytes = access_keys.read().unwrap_or_else(|e| e.into_inner()).get(key).and_then(|m| m.limit_bytes);
|
||
|
||
let mut stats = user_stats.write().unwrap_or_else(|e| e.into_inner());
|
||
stats.entry(key.to_string())
|
||
.or_insert_with(|| Arc::new(UserStats::new(limit_bytes)))
|
||
.clone()
|
||
}
|
||
|
||
fn track_user_bytes_up(
|
||
user_stats: &Arc<RwLock<HashMap<String, Arc<UserStats>>>>,
|
||
access_keys: &Arc<RwLock<HashMap<String, crate::api::UserMeta>>>,
|
||
key: &str,
|
||
bytes: u64,
|
||
) {
|
||
let stats = get_or_create_stats(user_stats, access_keys, key);
|
||
stats.bytes_up.fetch_add(bytes, Ordering::Relaxed);
|
||
}
|
||
|
||
fn track_user_bytes_down(
|
||
user_stats: &Arc<RwLock<HashMap<String, Arc<UserStats>>>>,
|
||
access_keys: &Arc<RwLock<HashMap<String, crate::api::UserMeta>>>,
|
||
key: &str,
|
||
bytes: u64,
|
||
) {
|
||
let stats = get_or_create_stats(user_stats, access_keys, key);
|
||
stats.bytes_down.fetch_add(bytes, Ordering::Relaxed);
|
||
}
|