1use std::collections::{HashMap, VecDeque};
12
13use crate::identity::{AgentId, MessageKind, SignedMessage, TaskStatus, check_freshness};
14use crate::policy::Policy;
15
16pub const MAX_FUTURE_MS: u64 = 60_000;
19
20pub const MAX_PAST_MS: u64 = 86_400_000; #[derive(Debug, PartialEq, Eq)]
29pub enum Dispatch {
30 Inline {
34 petname: String,
35 text: String,
36 task_id: Option<String>,
37 status: Option<TaskStatus>,
38 in_reply_to: Option<String>,
39 },
40 PairRequest { from_key: String, name: String },
43 PairAccept { from_key: String, name: String },
46}
47
48#[derive(Debug, PartialEq, Eq)]
50pub enum Reject {
51 BadSignature,
52 NotAllowlisted,
53 WrongRecipient,
54 Stale,
55 Replay,
56 Unexpected,
59}
60
61pub type Verdict = Result<Dispatch, Reject>;
62
63pub fn decide(
66 msg: &SignedMessage,
67 me: AgentId,
68 policy: &Policy,
69 now: u64,
70 seen: &mut Dedupe,
71) -> Verdict {
72 let from = msg.verify().map_err(|_| Reject::BadSignature)?;
75
76 let to = AgentId::from_b64(&msg.to).map_err(|_| Reject::WrongRecipient)?;
78 if to != me {
79 return Err(Reject::WrongRecipient);
80 }
81
82 let peer = policy.peer(from);
90 let is_self = from == me;
91 if peer.is_none() && !is_self && msg.kind == MessageKind::Message {
92 return Err(Reject::NotAllowlisted);
93 }
94
95 check_freshness(msg.ts, now, MAX_FUTURE_MS, MAX_PAST_MS).map_err(|_| Reject::Stale)?;
98
99 if !seen.insert(&msg.msg_id) {
102 return Err(Reject::Replay);
103 }
104
105 match (peer, msg.kind) {
106 (Some(peer), MessageKind::Message) => Ok(Dispatch::Inline {
107 petname: peer.petname.clone(),
108 text: msg.text.clone(),
109 task_id: msg.task_id.clone(),
110 status: msg.status,
111 in_reply_to: msg.in_reply_to.clone(),
112 }),
113 (None, MessageKind::Message) if is_self => Ok(Dispatch::Inline {
115 petname: "self".to_string(),
116 text: msg.text.clone(),
117 task_id: msg.task_id.clone(),
118 status: msg.status,
119 in_reply_to: msg.in_reply_to.clone(),
120 }),
121 (None, MessageKind::PairRequest) => Ok(Dispatch::PairRequest {
123 from_key: from.to_b64(),
124 name: msg.text.clone(),
125 }),
126 (None, MessageKind::PairAccept) => Ok(Dispatch::PairAccept {
127 from_key: from.to_b64(),
128 name: msg.text.clone(),
129 }),
130 _ => Err(Reject::Unexpected),
132 }
133}
134
135pub struct PairTable {
139 order: VecDeque<String>,
140 map: HashMap<String, String>,
141 cap: usize,
142}
143
144impl PairTable {
145 pub fn new(cap: usize) -> Self {
146 Self {
147 order: VecDeque::new(),
148 map: HashMap::new(),
149 cap: cap.max(1),
150 }
151 }
152
153 pub fn put(&mut self, key: String, value: String) {
155 if !self.map.contains_key(&key) {
156 if self.order.len() >= self.cap
157 && let Some(old) = self.order.pop_front()
158 {
159 self.map.remove(&old);
160 }
161 self.order.push_back(key.clone());
162 }
163 self.map.insert(key, value);
164 }
165
166 pub fn take(&mut self, key: &str) -> Option<String> {
168 let v = self.map.remove(key)?;
169 self.order.retain(|k| k != key);
170 Some(v)
171 }
172
173 pub fn get(&self, key: &str) -> Option<&String> {
174 self.map.get(key)
175 }
176
177 pub fn find(&self, key_or_fp: &str) -> Option<(String, String)> {
179 self.map
180 .iter()
181 .find(|(k, _)| {
182 k.as_str() == key_or_fp || k.chars().take(8).collect::<String>() == key_or_fp
183 })
184 .map(|(k, v)| (k.clone(), v.clone()))
185 }
186
187 pub fn entries(&self) -> impl Iterator<Item = (&String, &String)> {
188 self.map.iter()
189 }
190
191 pub fn is_empty(&self) -> bool {
192 self.map.is_empty()
193 }
194}
195
196pub struct Dedupe {
203 order: VecDeque<String>,
204 seen: std::collections::HashSet<String>,
205 cap: usize,
206}
207
208impl Dedupe {
209 pub fn new(cap: usize) -> Self {
210 Self {
211 order: VecDeque::with_capacity(cap),
212 seen: std::collections::HashSet::with_capacity(cap),
213 cap: cap.max(1),
214 }
215 }
216
217 pub fn insert(&mut self, id: &str) -> bool {
219 if self.seen.contains(id) {
220 return false;
221 }
222 if self.order.len() >= self.cap
223 && let Some(old) = self.order.pop_front()
224 {
225 self.seen.remove(&old);
226 }
227 self.order.push_back(id.to_string());
228 self.seen.insert(id.to_string());
229 true
230 }
231}
232
233#[cfg(test)]
234mod tests {
235 use super::*;
236 use crate::identity::AgentKey;
237 use crate::policy::Policy;
238
239 fn policy_for(key: &AgentKey) -> Policy {
240 let raw = format!(r#"{{ "alice": {{ "key": "{}" }} }}"#, key.id().to_b64());
241 Policy::parse(&raw).unwrap()
242 }
243
244 #[test]
245 fn admitted_peer_is_dispatched_inline() {
246 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
247 let policy = policy_for(&alice);
248 let msg = alice.sign(me.id(), "run the deploy", 1_000, "m1");
249 let mut seen = Dedupe::new(16);
250 assert_eq!(
251 decide(&msg, me.id(), &policy, 1_000, &mut seen),
252 Ok(Dispatch::Inline {
253 petname: "alice".into(),
254 text: "run the deploy".into(),
255 task_id: None,
256 status: None,
257 in_reply_to: None,
258 })
259 );
260 }
261
262 #[test]
263 fn stranger_is_rejected_even_with_valid_signature() {
264 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
265 let policy = policy_for(&AgentKey::generate().unwrap()); let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
267 let mut seen = Dedupe::new(16);
268 assert_eq!(
269 decide(&msg, me.id(), &policy, 1_000, &mut seen),
270 Err(Reject::NotAllowlisted)
271 );
272 }
273
274 #[test]
275 fn own_key_is_trusted_without_a_peers_entry() {
276 let me = AgentKey::generate().unwrap();
279 let policy = Policy::default(); let msg = me.sign(me.id(), "from my other session", 1_000, "m1");
281 let mut seen = Dedupe::new(16);
282 assert_eq!(
283 decide(&msg, me.id(), &policy, 1_000, &mut seen),
284 Ok(Dispatch::Inline {
285 petname: "self".into(),
286 text: "from my other session".into(),
287 task_id: None,
288 status: None,
289 in_reply_to: None,
290 })
291 );
292 }
293
294 #[test]
295 fn forged_sender_is_bad_signature() {
296 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
297 let policy = policy_for(&alice);
298 let eve = AgentKey::generate().unwrap();
300 let mut msg = eve.sign(me.id(), "hi", 1_000, "m1");
301 msg.from = alice.id().to_b64();
302 let mut seen = Dedupe::new(16);
303 assert_eq!(
304 decide(&msg, me.id(), &policy, 1_000, &mut seen),
305 Err(Reject::BadSignature)
306 );
307 }
308
309 #[test]
310 fn message_for_someone_else_is_rejected() {
311 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
312 let other = AgentKey::generate().unwrap();
313 let policy = policy_for(&alice);
314 let msg = alice.sign(other.id(), "hi", 1_000, "m1"); let mut seen = Dedupe::new(16);
316 assert_eq!(
317 decide(&msg, me.id(), &policy, 1_000, &mut seen),
318 Err(Reject::WrongRecipient)
319 );
320 }
321
322 #[test]
323 fn stale_message_is_rejected() {
324 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
325 let policy = policy_for(&alice);
326 let msg = alice.sign(me.id(), "hi", 1_000, "m1");
327 let mut seen = Dedupe::new(16);
328 let long_after = 1_000 + MAX_PAST_MS + 1;
330 assert_eq!(
331 decide(&msg, me.id(), &policy, long_after, &mut seen),
332 Err(Reject::Stale)
333 );
334 let future_msg = alice.sign(me.id(), "hi", 1_000 + MAX_FUTURE_MS + 1, "m2");
336 assert_eq!(
337 decide(&future_msg, me.id(), &policy, 1_000, &mut seen),
338 Err(Reject::Stale)
339 );
340 let ok_msg = alice.sign(me.id(), "hi", 1_000, "m3");
342 assert!(decide(&ok_msg, me.id(), &policy, 1_000 + 300_000, &mut seen).is_ok());
343 }
344
345 #[test]
346 fn replay_is_rejected_the_second_time() {
347 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
348 let policy = policy_for(&alice);
349 let msg = alice.sign(me.id(), "hi", 1_000, "m1");
350 let mut seen = Dedupe::new(16);
351 assert!(decide(&msg, me.id(), &policy, 1_000, &mut seen).is_ok());
352 assert_eq!(
353 decide(&msg, me.id(), &policy, 1_000, &mut seen),
354 Err(Reject::Replay)
355 );
356 }
357
358 #[test]
359 fn dedupe_forgets_oldest_beyond_cap() {
360 let mut d = Dedupe::new(2);
361 assert!(d.insert("a"));
362 assert!(d.insert("b"));
363 assert!(d.insert("c")); assert!(d.insert("a"), "a was evicted, so it is fresh again");
365 assert!(!d.insert("c"), "c is still within the window");
366 }
367
368 #[test]
369 fn non_peer_knock_is_surfaced_not_dropped() {
370 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
371 let policy = Policy::default(); let msg = stranger.sign_as(
373 me.id(),
374 "stranger-laptop",
375 1_000,
376 "k1",
377 MessageKind::PairRequest,
378 );
379 let mut seen = Dedupe::new(16);
380 assert_eq!(
381 decide(&msg, me.id(), &policy, 1_000, &mut seen),
382 Ok(Dispatch::PairRequest {
383 from_key: stranger.id().to_b64(),
384 name: "stranger-laptop".into()
385 })
386 );
387 }
388
389 #[test]
390 fn non_peer_plain_message_is_still_denied() {
391 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
392 let policy = Policy::default();
393 let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
394 let mut seen = Dedupe::new(16);
395 assert_eq!(
396 decide(&msg, me.id(), &policy, 1_000, &mut seen),
397 Err(Reject::NotAllowlisted)
398 );
399 }
400
401 #[test]
402 fn pairing_kind_from_existing_peer_is_unexpected() {
403 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
404 let policy = policy_for(&alice);
405 let msg = alice.sign_as(me.id(), "x", 1_000, "k1", MessageKind::PairRequest);
406 let mut seen = Dedupe::new(16);
407 assert_eq!(
408 decide(&msg, me.id(), &policy, 1_000, &mut seen),
409 Err(Reject::Unexpected)
410 );
411 }
412
413 #[test]
414 fn pair_table_put_take_and_find() {
415 let mut t = PairTable::new(8);
416 t.put("aaaabbbbcccc".into(), "desktop".into());
417 assert_eq!(t.get("aaaabbbbcccc"), Some(&"desktop".to_string()));
418 assert_eq!(
419 t.find("aaaabbbb"), Some(("aaaabbbbcccc".into(), "desktop".into()))
421 );
422 assert_eq!(t.take("aaaabbbbcccc"), Some("desktop".into()));
423 assert!(t.is_empty());
424 }
425}