1use std::collections::{HashMap, VecDeque};
12
13use crate::identity::{AgentId, MessageKind, SignedMessage, TaskStatus, check_freshness};
14use crate::policy::Policy;
15
16pub const MAX_SKEW_MS: u64 = 60_000;
19
20#[derive(Debug, PartialEq, Eq)]
22pub enum Dispatch {
23 Inline {
27 petname: String,
28 text: String,
29 task_id: Option<String>,
30 status: Option<TaskStatus>,
31 in_reply_to: Option<String>,
32 },
33 PairRequest { from_key: String, name: String },
36 PairAccept { from_key: String, name: String },
39}
40
41#[derive(Debug, PartialEq, Eq)]
43pub enum Reject {
44 BadSignature,
45 NotAllowlisted,
46 WrongRecipient,
47 Stale,
48 Replay,
49 Unexpected,
52}
53
54pub type Verdict = Result<Dispatch, Reject>;
55
56pub fn decide(
59 msg: &SignedMessage,
60 me: AgentId,
61 policy: &Policy,
62 now: u64,
63 seen: &mut Dedupe,
64) -> Verdict {
65 let from = msg.verify().map_err(|_| Reject::BadSignature)?;
68
69 let to = AgentId::from_b64(&msg.to).map_err(|_| Reject::WrongRecipient)?;
71 if to != me {
72 return Err(Reject::WrongRecipient);
73 }
74
75 let peer = policy.peer(from);
83 let is_self = from == me;
84 if peer.is_none() && !is_self && msg.kind == MessageKind::Message {
85 return Err(Reject::NotAllowlisted);
86 }
87
88 check_freshness(msg.ts, now, MAX_SKEW_MS).map_err(|_| Reject::Stale)?;
90
91 if !seen.insert(&msg.msg_id) {
94 return Err(Reject::Replay);
95 }
96
97 match (peer, msg.kind) {
98 (Some(peer), MessageKind::Message) => Ok(Dispatch::Inline {
99 petname: peer.petname.clone(),
100 text: msg.text.clone(),
101 task_id: msg.task_id.clone(),
102 status: msg.status,
103 in_reply_to: msg.in_reply_to.clone(),
104 }),
105 (None, MessageKind::Message) if is_self => Ok(Dispatch::Inline {
107 petname: "self".to_string(),
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::PairRequest) => Ok(Dispatch::PairRequest {
115 from_key: from.to_b64(),
116 name: msg.text.clone(),
117 }),
118 (None, MessageKind::PairAccept) => Ok(Dispatch::PairAccept {
119 from_key: from.to_b64(),
120 name: msg.text.clone(),
121 }),
122 _ => Err(Reject::Unexpected),
124 }
125}
126
127pub struct PairTable {
131 order: VecDeque<String>,
132 map: HashMap<String, String>,
133 cap: usize,
134}
135
136impl PairTable {
137 pub fn new(cap: usize) -> Self {
138 Self {
139 order: VecDeque::new(),
140 map: HashMap::new(),
141 cap: cap.max(1),
142 }
143 }
144
145 pub fn put(&mut self, key: String, value: String) {
147 if !self.map.contains_key(&key) {
148 if self.order.len() >= self.cap
149 && let Some(old) = self.order.pop_front()
150 {
151 self.map.remove(&old);
152 }
153 self.order.push_back(key.clone());
154 }
155 self.map.insert(key, value);
156 }
157
158 pub fn take(&mut self, key: &str) -> Option<String> {
160 let v = self.map.remove(key)?;
161 self.order.retain(|k| k != key);
162 Some(v)
163 }
164
165 pub fn get(&self, key: &str) -> Option<&String> {
166 self.map.get(key)
167 }
168
169 pub fn find(&self, key_or_fp: &str) -> Option<(String, String)> {
171 self.map
172 .iter()
173 .find(|(k, _)| {
174 k.as_str() == key_or_fp || k.chars().take(8).collect::<String>() == key_or_fp
175 })
176 .map(|(k, v)| (k.clone(), v.clone()))
177 }
178
179 pub fn entries(&self) -> impl Iterator<Item = (&String, &String)> {
180 self.map.iter()
181 }
182
183 pub fn is_empty(&self) -> bool {
184 self.map.is_empty()
185 }
186}
187
188pub struct Dedupe {
192 order: VecDeque<String>,
193 seen: std::collections::HashSet<String>,
194 cap: usize,
195}
196
197impl Dedupe {
198 pub fn new(cap: usize) -> Self {
199 Self {
200 order: VecDeque::with_capacity(cap),
201 seen: std::collections::HashSet::with_capacity(cap),
202 cap: cap.max(1),
203 }
204 }
205
206 pub fn insert(&mut self, id: &str) -> bool {
208 if self.seen.contains(id) {
209 return false;
210 }
211 if self.order.len() >= self.cap
212 && let Some(old) = self.order.pop_front()
213 {
214 self.seen.remove(&old);
215 }
216 self.order.push_back(id.to_string());
217 self.seen.insert(id.to_string());
218 true
219 }
220}
221
222#[cfg(test)]
223mod tests {
224 use super::*;
225 use crate::identity::AgentKey;
226 use crate::policy::Policy;
227
228 fn policy_for(key: &AgentKey) -> Policy {
229 let raw = format!(r#"{{ "alice": {{ "key": "{}" }} }}"#, key.id().to_b64());
230 Policy::parse(&raw).unwrap()
231 }
232
233 #[test]
234 fn admitted_peer_is_dispatched_inline() {
235 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
236 let policy = policy_for(&alice);
237 let msg = alice.sign(me.id(), "run the deploy", 1_000, "m1");
238 let mut seen = Dedupe::new(16);
239 assert_eq!(
240 decide(&msg, me.id(), &policy, 1_000, &mut seen),
241 Ok(Dispatch::Inline {
242 petname: "alice".into(),
243 text: "run the deploy".into(),
244 task_id: None,
245 status: None,
246 in_reply_to: None,
247 })
248 );
249 }
250
251 #[test]
252 fn stranger_is_rejected_even_with_valid_signature() {
253 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
254 let policy = policy_for(&AgentKey::generate().unwrap()); let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
256 let mut seen = Dedupe::new(16);
257 assert_eq!(
258 decide(&msg, me.id(), &policy, 1_000, &mut seen),
259 Err(Reject::NotAllowlisted)
260 );
261 }
262
263 #[test]
264 fn own_key_is_trusted_without_a_peers_entry() {
265 let me = AgentKey::generate().unwrap();
268 let policy = Policy::default(); let msg = me.sign(me.id(), "from my other session", 1_000, "m1");
270 let mut seen = Dedupe::new(16);
271 assert_eq!(
272 decide(&msg, me.id(), &policy, 1_000, &mut seen),
273 Ok(Dispatch::Inline {
274 petname: "self".into(),
275 text: "from my other session".into(),
276 task_id: None,
277 status: None,
278 in_reply_to: None,
279 })
280 );
281 }
282
283 #[test]
284 fn forged_sender_is_bad_signature() {
285 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
286 let policy = policy_for(&alice);
287 let eve = AgentKey::generate().unwrap();
289 let mut msg = eve.sign(me.id(), "hi", 1_000, "m1");
290 msg.from = alice.id().to_b64();
291 let mut seen = Dedupe::new(16);
292 assert_eq!(
293 decide(&msg, me.id(), &policy, 1_000, &mut seen),
294 Err(Reject::BadSignature)
295 );
296 }
297
298 #[test]
299 fn message_for_someone_else_is_rejected() {
300 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
301 let other = AgentKey::generate().unwrap();
302 let policy = policy_for(&alice);
303 let msg = alice.sign(other.id(), "hi", 1_000, "m1"); let mut seen = Dedupe::new(16);
305 assert_eq!(
306 decide(&msg, me.id(), &policy, 1_000, &mut seen),
307 Err(Reject::WrongRecipient)
308 );
309 }
310
311 #[test]
312 fn stale_message_is_rejected() {
313 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
314 let policy = policy_for(&alice);
315 let msg = alice.sign(me.id(), "hi", 1_000, "m1");
316 let mut seen = Dedupe::new(16);
317 let far_future = 1_000 + MAX_SKEW_MS + 1;
318 assert_eq!(
319 decide(&msg, me.id(), &policy, far_future, &mut seen),
320 Err(Reject::Stale)
321 );
322 }
323
324 #[test]
325 fn replay_is_rejected_the_second_time() {
326 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
327 let policy = policy_for(&alice);
328 let msg = alice.sign(me.id(), "hi", 1_000, "m1");
329 let mut seen = Dedupe::new(16);
330 assert!(decide(&msg, me.id(), &policy, 1_000, &mut seen).is_ok());
331 assert_eq!(
332 decide(&msg, me.id(), &policy, 1_000, &mut seen),
333 Err(Reject::Replay)
334 );
335 }
336
337 #[test]
338 fn dedupe_forgets_oldest_beyond_cap() {
339 let mut d = Dedupe::new(2);
340 assert!(d.insert("a"));
341 assert!(d.insert("b"));
342 assert!(d.insert("c")); assert!(d.insert("a"), "a was evicted, so it is fresh again");
344 assert!(!d.insert("c"), "c is still within the window");
345 }
346
347 #[test]
348 fn non_peer_knock_is_surfaced_not_dropped() {
349 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
350 let policy = Policy::default(); let msg = stranger.sign_as(
352 me.id(),
353 "stranger-laptop",
354 1_000,
355 "k1",
356 MessageKind::PairRequest,
357 );
358 let mut seen = Dedupe::new(16);
359 assert_eq!(
360 decide(&msg, me.id(), &policy, 1_000, &mut seen),
361 Ok(Dispatch::PairRequest {
362 from_key: stranger.id().to_b64(),
363 name: "stranger-laptop".into()
364 })
365 );
366 }
367
368 #[test]
369 fn non_peer_plain_message_is_still_denied() {
370 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
371 let policy = Policy::default();
372 let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
373 let mut seen = Dedupe::new(16);
374 assert_eq!(
375 decide(&msg, me.id(), &policy, 1_000, &mut seen),
376 Err(Reject::NotAllowlisted)
377 );
378 }
379
380 #[test]
381 fn pairing_kind_from_existing_peer_is_unexpected() {
382 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
383 let policy = policy_for(&alice);
384 let msg = alice.sign_as(me.id(), "x", 1_000, "k1", MessageKind::PairRequest);
385 let mut seen = Dedupe::new(16);
386 assert_eq!(
387 decide(&msg, me.id(), &policy, 1_000, &mut seen),
388 Err(Reject::Unexpected)
389 );
390 }
391
392 #[test]
393 fn pair_table_put_take_and_find() {
394 let mut t = PairTable::new(8);
395 t.put("aaaabbbbcccc".into(), "desktop".into());
396 assert_eq!(t.get("aaaabbbbcccc"), Some(&"desktop".to_string()));
397 assert_eq!(
398 t.find("aaaabbbb"), Some(("aaaabbbbcccc".into(), "desktop".into()))
400 );
401 assert_eq!(t.take("aaaabbbbcccc"), Some("desktop".into()));
402 assert!(t.is_empty());
403 }
404}