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 find(&self, key_or_fp: &str) -> Option<(String, String)> {
175 self.map
176 .iter()
177 .find(|(k, _)| {
178 k.as_str() == key_or_fp || k.chars().take(8).collect::<String>() == key_or_fp
179 })
180 .map(|(k, v)| (k.clone(), v.clone()))
181 }
182
183 pub fn entries(&self) -> impl Iterator<Item = (&String, &String)> {
184 self.map.iter()
185 }
186
187 pub fn is_empty(&self) -> bool {
188 self.map.is_empty()
189 }
190}
191
192pub struct Dedupe {
199 order: VecDeque<String>,
200 seen: std::collections::HashSet<String>,
201 cap: usize,
202}
203
204impl Dedupe {
205 pub fn new(cap: usize) -> Self {
206 Self {
207 order: VecDeque::with_capacity(cap),
208 seen: std::collections::HashSet::with_capacity(cap),
209 cap: cap.max(1),
210 }
211 }
212
213 pub fn insert(&mut self, id: &str) -> bool {
215 if self.seen.contains(id) {
216 return false;
217 }
218 if self.order.len() >= self.cap
219 && let Some(old) = self.order.pop_front()
220 {
221 self.seen.remove(&old);
222 }
223 self.order.push_back(id.to_string());
224 self.seen.insert(id.to_string());
225 true
226 }
227}
228
229#[cfg(test)]
230mod tests {
231 use super::*;
232 use crate::identity::AgentKey;
233 use crate::policy::Policy;
234
235 fn policy_for(key: &AgentKey) -> Policy {
236 let raw = format!(r#"{{ "alice": {{ "key": "{}" }} }}"#, key.id().to_b64());
237 Policy::parse(&raw).unwrap()
238 }
239
240 #[test]
241 fn admitted_peer_is_dispatched_inline() {
242 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
243 let policy = policy_for(&alice);
244 let msg = alice.sign(me.id(), "run the deploy", 1_000, "m1");
245 let mut seen = Dedupe::new(16);
246 assert_eq!(
247 decide(&msg, me.id(), &policy, 1_000, &mut seen),
248 Ok(Dispatch::Inline {
249 petname: "alice".into(),
250 text: "run the deploy".into(),
251 task_id: None,
252 status: None,
253 in_reply_to: None,
254 })
255 );
256 }
257
258 #[test]
259 fn stranger_is_rejected_even_with_valid_signature() {
260 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
261 let policy = policy_for(&AgentKey::generate().unwrap()); let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
263 let mut seen = Dedupe::new(16);
264 assert_eq!(
265 decide(&msg, me.id(), &policy, 1_000, &mut seen),
266 Err(Reject::NotAllowlisted)
267 );
268 }
269
270 #[test]
271 fn own_key_is_trusted_without_a_peers_entry() {
272 let me = AgentKey::generate().unwrap();
275 let policy = Policy::default(); let msg = me.sign(me.id(), "from my other session", 1_000, "m1");
277 let mut seen = Dedupe::new(16);
278 assert_eq!(
279 decide(&msg, me.id(), &policy, 1_000, &mut seen),
280 Ok(Dispatch::Inline {
281 petname: "self".into(),
282 text: "from my other session".into(),
283 task_id: None,
284 status: None,
285 in_reply_to: None,
286 })
287 );
288 }
289
290 #[test]
291 fn forged_sender_is_bad_signature() {
292 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
293 let policy = policy_for(&alice);
294 let eve = AgentKey::generate().unwrap();
296 let mut msg = eve.sign(me.id(), "hi", 1_000, "m1");
297 msg.from = alice.id().to_b64();
298 let mut seen = Dedupe::new(16);
299 assert_eq!(
300 decide(&msg, me.id(), &policy, 1_000, &mut seen),
301 Err(Reject::BadSignature)
302 );
303 }
304
305 #[test]
306 fn message_for_someone_else_is_rejected() {
307 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
308 let other = AgentKey::generate().unwrap();
309 let policy = policy_for(&alice);
310 let msg = alice.sign(other.id(), "hi", 1_000, "m1"); let mut seen = Dedupe::new(16);
312 assert_eq!(
313 decide(&msg, me.id(), &policy, 1_000, &mut seen),
314 Err(Reject::WrongRecipient)
315 );
316 }
317
318 #[test]
319 fn stale_message_is_rejected() {
320 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
321 let policy = policy_for(&alice);
322 let msg = alice.sign(me.id(), "hi", 1_000, "m1");
323 let mut seen = Dedupe::new(16);
324 let long_after = 1_000 + MAX_PAST_MS + 1;
326 assert_eq!(
327 decide(&msg, me.id(), &policy, long_after, &mut seen),
328 Err(Reject::Stale)
329 );
330 let future_msg = alice.sign(me.id(), "hi", 1_000 + MAX_FUTURE_MS + 1, "m2");
332 assert_eq!(
333 decide(&future_msg, me.id(), &policy, 1_000, &mut seen),
334 Err(Reject::Stale)
335 );
336 let ok_msg = alice.sign(me.id(), "hi", 1_000, "m3");
338 assert!(decide(&ok_msg, me.id(), &policy, 1_000 + 300_000, &mut seen).is_ok());
339 }
340
341 #[test]
342 fn replay_is_rejected_the_second_time() {
343 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
344 let policy = policy_for(&alice);
345 let msg = alice.sign(me.id(), "hi", 1_000, "m1");
346 let mut seen = Dedupe::new(16);
347 assert!(decide(&msg, me.id(), &policy, 1_000, &mut seen).is_ok());
348 assert_eq!(
349 decide(&msg, me.id(), &policy, 1_000, &mut seen),
350 Err(Reject::Replay)
351 );
352 }
353
354 #[test]
355 fn dedupe_forgets_oldest_beyond_cap() {
356 let mut d = Dedupe::new(2);
357 assert!(d.insert("a"));
358 assert!(d.insert("b"));
359 assert!(d.insert("c")); assert!(d.insert("a"), "a was evicted, so it is fresh again");
361 assert!(!d.insert("c"), "c is still within the window");
362 }
363
364 #[test]
365 fn non_peer_knock_is_surfaced_not_dropped() {
366 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
367 let policy = Policy::default(); let msg = stranger.sign_as(
369 me.id(),
370 "stranger-laptop",
371 1_000,
372 "k1",
373 MessageKind::PairRequest,
374 );
375 let mut seen = Dedupe::new(16);
376 assert_eq!(
377 decide(&msg, me.id(), &policy, 1_000, &mut seen),
378 Ok(Dispatch::PairRequest {
379 from_key: stranger.id().to_b64(),
380 name: "stranger-laptop".into()
381 })
382 );
383 }
384
385 #[test]
386 fn non_peer_plain_message_is_still_denied() {
387 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
388 let policy = Policy::default();
389 let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
390 let mut seen = Dedupe::new(16);
391 assert_eq!(
392 decide(&msg, me.id(), &policy, 1_000, &mut seen),
393 Err(Reject::NotAllowlisted)
394 );
395 }
396
397 #[test]
398 fn pairing_kind_from_existing_peer_is_unexpected() {
399 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
400 let policy = policy_for(&alice);
401 let msg = alice.sign_as(me.id(), "x", 1_000, "k1", MessageKind::PairRequest);
402 let mut seen = Dedupe::new(16);
403 assert_eq!(
404 decide(&msg, me.id(), &policy, 1_000, &mut seen),
405 Err(Reject::Unexpected)
406 );
407 }
408
409 #[test]
410 fn pair_table_put_take_and_find() {
411 let mut t = PairTable::new(8);
412 t.put("aaaabbbbcccc".into(), "desktop".into());
413 assert_eq!(
414 t.find("aaaabbbbcccc"), Some(("aaaabbbbcccc".into(), "desktop".into()))
416 );
417 assert_eq!(
418 t.find("aaaabbbb"), Some(("aaaabbbbcccc".into(), "desktop".into()))
420 );
421 assert_eq!(t.take("aaaabbbbcccc"), Some("desktop".into()));
422 assert!(t.is_empty());
423 }
424}