macula_rust/station_link/
admission.rs1use std::collections::HashMap;
16use std::hash::Hash;
17use std::sync::{Arc, Mutex, MutexGuard};
18
19use crate::frame::VerifiedRequest;
20
21use super::LinkError;
22
23const DEADLINE_PAST_TOLERANCE_MS: i64 = 5 * 60_000;
24const DEADLINE_AHEAD_MAX_MS: i64 = 10 * 60_000;
25const KEPT_PAST_DEADLINE_MS: i64 = 5 * 60_000;
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub struct AdmissionLimits {
32 pub caller_quota: usize,
33 pub share: usize,
34 pub cap: usize,
35 pub reply_bytes: usize,
36 pub reply_bytes_total: usize,
37 pub sessions_per_caller: usize,
38 pub sessions: usize,
39 pub inbox_bytes_per_caller: usize,
40 pub inbox_bytes: usize,
41}
42
43impl Default for AdmissionLimits {
44 fn default() -> Self {
48 AdmissionLimits {
49 caller_quota: 256,
50 share: 1024,
51 cap: 1024,
52 reply_bytes: 256 * 1024,
53 reply_bytes_total: 16 * 1024 * 1024,
54 sessions_per_caller: 16,
55 sessions: 1000,
56 inbox_bytes_per_caller: 16 * 1024 * 1024,
57 inbox_bytes: 256 * 1024 * 1024,
58 }
59 }
60}
61
62impl AdmissionLimits {
63 pub fn validate(&self) -> Result<(), LinkError> {
66 let l = self;
67 let valid = l.caller_quota > 0
68 && l.share > 0
69 && l.cap > 0
70 && l.reply_bytes > 0
71 && l.reply_bytes_total > 0
72 && l.caller_quota <= l.share
73 && l.reply_bytes <= l.reply_bytes_total
74 && l.sessions_per_caller > 0
75 && l.sessions >= l.sessions_per_caller
76 && l.inbox_bytes_per_caller > 0
77 && l.inbox_bytes >= l.inbox_bytes_per_caller;
78 if valid {
79 Ok(())
80 } else {
81 Err(LinkError::InvalidConfig(
82 "admission limits must be positive, each per-caller bound within its total".into(),
83 ))
84 }
85 }
86}
87
88#[derive(Debug, Clone, PartialEq, Eq)]
91pub(super) enum Verdict {
92 Refused(&'static str),
93 Copy(Option<Vec<u8>>),
94 New,
95}
96
97struct Entry {
98 hash: [u8; 48],
99 expires_at: i64,
100 share: String,
101 answered: bool,
102 reply: Option<Vec<u8>>,
104}
105
106impl Entry {
107 fn judge_copy(&self, request_hash: &[u8; 48]) -> Verdict {
111 if self.hash != *request_hash {
112 return Verdict::Refused("request_id_reused");
113 }
114 if self.answered && self.reply.is_none() {
115 return Verdict::Refused("reply_not_kept");
116 }
117 Verdict::Copy(self.reply.clone())
118 }
119}
120
121#[derive(Default)]
122struct Held {
123 entries: HashMap<([u8; 32], [u8; 16]), Entry>,
124 callers: HashMap<[u8; 32], usize>,
125 shares: HashMap<String, usize>,
126 reply_bytes: HashMap<[u8; 32], usize>,
127 reply_total: usize,
128 sessions: HashMap<[u8; 32], usize>,
129 sessions_total: usize,
130 inbox: HashMap<[u8; 32], usize>,
131 inbox_total: usize,
132}
133
134pub struct Admission {
136 limits: AdmissionLimits,
137 held: Mutex<Held>,
138}
139
140impl Admission {
141 pub fn new(limits: AdmissionLimits) -> Admission {
144 Admission {
145 limits,
146 held: Mutex::new(Held::default()),
147 }
148 }
149
150 pub fn limits(&self) -> AdmissionLimits {
152 self.limits
153 }
154
155 fn lock(&self) -> MutexGuard<'_, Held> {
156 self.held.lock().unwrap_or_else(|p| p.into_inner())
157 }
158
159 pub(super) fn admit(&self, request: &VerifiedRequest, share: &str, now_ms: i64) -> Verdict {
162 let deadline = request.deadline as i64;
163 if deadline < now_ms - DEADLINE_PAST_TOLERANCE_MS {
164 return Verdict::Refused("expired");
165 }
166 if deadline > now_ms + DEADLINE_AHEAD_MAX_MS {
167 return Verdict::Refused("not_yet_valid");
168 }
169 let mut held = self.lock();
170 held.sweep(now_ms);
171 let key = (request.caller, request.request_id);
172 if let Some(entry) = held.entries.get(&key) {
173 return entry.judge_copy(&request.request_hash);
174 }
175 if held.callers.get(&request.caller).copied().unwrap_or(0) >= self.limits.caller_quota {
176 return Verdict::Refused("caller_quota");
177 }
178 if held.shares.get(share).copied().unwrap_or(0) >= self.limits.share {
179 return Verdict::Refused("share_full");
180 }
181 if held.entries.len() >= self.limits.cap {
182 return Verdict::Refused("admission_full");
183 }
184 held.entries.insert(
185 key,
186 Entry {
187 hash: request.request_hash,
188 expires_at: deadline + KEPT_PAST_DEADLINE_MS,
189 share: share.to_string(),
190 answered: false,
191 reply: None,
192 },
193 );
194 *held.callers.entry(request.caller).or_default() += 1;
195 *held.shares.entry(share.to_string()).or_default() += 1;
196 Verdict::New
197 }
198
199 pub(super) fn store(&self, request: &VerifiedRequest, reply: Vec<u8>) {
202 let mut held = self.lock();
203 let held = &mut *held;
204 let Some(entry) = held.entries.get_mut(&(request.caller, request.request_id)) else {
205 return;
206 };
207 if entry.answered || entry.hash != request.request_hash {
208 return;
209 }
210 entry.answered = true;
211 let caller_bytes = held.reply_bytes.get(&request.caller).copied().unwrap_or(0);
212 if caller_bytes + reply.len() > self.limits.reply_bytes
213 || held.reply_total + reply.len() > self.limits.reply_bytes_total
214 {
215 return;
216 }
217 *held.reply_bytes.entry(request.caller).or_default() += reply.len();
218 held.reply_total += reply.len();
219 entry.reply = Some(reply);
220 }
221
222 pub(super) fn open_session(self: &Arc<Self>, caller: [u8; 32]) -> Option<SessionPlace> {
225 let mut held = self.lock();
226 if held.sessions.get(&caller).copied().unwrap_or(0) >= self.limits.sessions_per_caller
227 || held.sessions_total >= self.limits.sessions
228 {
229 return None;
230 }
231 *held.sessions.entry(caller).or_default() += 1;
232 held.sessions_total += 1;
233 Some(SessionPlace {
234 admission: self.clone(),
235 caller,
236 })
237 }
238
239 pub(super) fn charge_inbox(&self, caller: [u8; 32], n: usize) -> bool {
242 let mut held = self.lock();
243 if held.inbox.get(&caller).copied().unwrap_or(0) + n > self.limits.inbox_bytes_per_caller
244 || held.inbox_total + n > self.limits.inbox_bytes
245 {
246 return false;
247 }
248 *held.inbox.entry(caller).or_default() += n;
249 held.inbox_total += n;
250 true
251 }
252
253 pub(super) fn release_inbox(&self, caller: [u8; 32], n: usize) {
255 let mut held = self.lock();
256 decrement(&mut held.inbox, caller, n);
257 held.inbox_total = held.inbox_total.saturating_sub(n);
258 }
259}
260
261pub(super) struct SessionPlace {
263 admission: Arc<Admission>,
264 caller: [u8; 32],
265}
266
267impl Drop for SessionPlace {
268 fn drop(&mut self) {
269 let mut held = self.admission.lock();
270 decrement(&mut held.sessions, self.caller, 1);
271 held.sessions_total = held.sessions_total.saturating_sub(1);
272 }
273}
274
275impl Held {
276 fn sweep(&mut self, now_ms: i64) {
279 let expired: Vec<_> = self
280 .entries
281 .iter()
282 .filter(|(_, e)| e.expires_at < now_ms)
283 .map(|(k, _)| *k)
284 .collect();
285 for key in expired {
286 self.remove_entry(key);
287 }
288 }
289
290 fn remove_entry(&mut self, key: ([u8; 32], [u8; 16])) {
293 let Some(entry) = self.entries.remove(&key) else {
294 return;
295 };
296 decrement(&mut self.callers, key.0, 1);
297 decrement(&mut self.shares, entry.share, 1);
298 if let Some(reply) = entry.reply {
299 decrement(&mut self.reply_bytes, key.0, reply.len());
300 self.reply_total = self.reply_total.saturating_sub(reply.len());
301 }
302 }
303}
304
305fn decrement<K: Eq + Hash>(m: &mut HashMap<K, usize>, k: K, n: usize) {
307 if let Some(v) = m.get_mut(&k) {
308 *v = v.saturating_sub(n);
309 if *v == 0 {
310 m.remove(&k);
311 }
312 }
313}