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);
79 if peer.is_none() && msg.kind == MessageKind::Message {
80 return Err(Reject::NotAllowlisted);
81 }
82
83 check_freshness(msg.ts, now, MAX_SKEW_MS).map_err(|_| Reject::Stale)?;
85
86 if !seen.insert(&msg.msg_id) {
89 return Err(Reject::Replay);
90 }
91
92 match (peer, msg.kind) {
93 (Some(peer), MessageKind::Message) => Ok(Dispatch::Inline {
94 petname: peer.petname.clone(),
95 text: msg.text.clone(),
96 task_id: msg.task_id.clone(),
97 status: msg.status,
98 in_reply_to: msg.in_reply_to.clone(),
99 }),
100 (None, MessageKind::PairRequest) => Ok(Dispatch::PairRequest {
102 from_key: from.to_b64(),
103 name: msg.text.clone(),
104 }),
105 (None, MessageKind::PairAccept) => Ok(Dispatch::PairAccept {
106 from_key: from.to_b64(),
107 name: msg.text.clone(),
108 }),
109 _ => Err(Reject::Unexpected),
111 }
112}
113
114pub struct PairTable {
118 order: VecDeque<String>,
119 map: HashMap<String, String>,
120 cap: usize,
121}
122
123impl PairTable {
124 pub fn new(cap: usize) -> Self {
125 Self {
126 order: VecDeque::new(),
127 map: HashMap::new(),
128 cap: cap.max(1),
129 }
130 }
131
132 pub fn put(&mut self, key: String, value: String) {
134 if !self.map.contains_key(&key) {
135 if self.order.len() >= self.cap
136 && let Some(old) = self.order.pop_front()
137 {
138 self.map.remove(&old);
139 }
140 self.order.push_back(key.clone());
141 }
142 self.map.insert(key, value);
143 }
144
145 pub fn take(&mut self, key: &str) -> Option<String> {
147 let v = self.map.remove(key)?;
148 self.order.retain(|k| k != key);
149 Some(v)
150 }
151
152 pub fn get(&self, key: &str) -> Option<&String> {
153 self.map.get(key)
154 }
155
156 pub fn find(&self, key_or_fp: &str) -> Option<(String, String)> {
158 self.map
159 .iter()
160 .find(|(k, _)| {
161 k.as_str() == key_or_fp || k.chars().take(8).collect::<String>() == key_or_fp
162 })
163 .map(|(k, v)| (k.clone(), v.clone()))
164 }
165
166 pub fn entries(&self) -> impl Iterator<Item = (&String, &String)> {
167 self.map.iter()
168 }
169
170 pub fn is_empty(&self) -> bool {
171 self.map.is_empty()
172 }
173}
174
175pub struct Dedupe {
179 order: VecDeque<String>,
180 seen: std::collections::HashSet<String>,
181 cap: usize,
182}
183
184impl Dedupe {
185 pub fn new(cap: usize) -> Self {
186 Self {
187 order: VecDeque::with_capacity(cap),
188 seen: std::collections::HashSet::with_capacity(cap),
189 cap: cap.max(1),
190 }
191 }
192
193 pub fn insert(&mut self, id: &str) -> bool {
195 if self.seen.contains(id) {
196 return false;
197 }
198 if self.order.len() >= self.cap
199 && let Some(old) = self.order.pop_front()
200 {
201 self.seen.remove(&old);
202 }
203 self.order.push_back(id.to_string());
204 self.seen.insert(id.to_string());
205 true
206 }
207}
208
209#[cfg(test)]
210mod tests {
211 use super::*;
212 use crate::identity::AgentKey;
213 use crate::policy::Policy;
214
215 fn policy_for(key: &AgentKey) -> Policy {
216 let raw = format!(r#"{{ "alice": {{ "key": "{}" }} }}"#, key.id().to_b64());
217 Policy::parse(&raw).unwrap()
218 }
219
220 #[test]
221 fn admitted_peer_is_dispatched_inline() {
222 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
223 let policy = policy_for(&alice);
224 let msg = alice.sign(me.id(), "run the deploy", 1_000, "m1");
225 let mut seen = Dedupe::new(16);
226 assert_eq!(
227 decide(&msg, me.id(), &policy, 1_000, &mut seen),
228 Ok(Dispatch::Inline {
229 petname: "alice".into(),
230 text: "run the deploy".into(),
231 task_id: None,
232 status: None,
233 in_reply_to: None,
234 })
235 );
236 }
237
238 #[test]
239 fn stranger_is_rejected_even_with_valid_signature() {
240 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
241 let policy = policy_for(&AgentKey::generate().unwrap()); let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
243 let mut seen = Dedupe::new(16);
244 assert_eq!(
245 decide(&msg, me.id(), &policy, 1_000, &mut seen),
246 Err(Reject::NotAllowlisted)
247 );
248 }
249
250 #[test]
251 fn forged_sender_is_bad_signature() {
252 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
253 let policy = policy_for(&alice);
254 let eve = AgentKey::generate().unwrap();
256 let mut msg = eve.sign(me.id(), "hi", 1_000, "m1");
257 msg.from = alice.id().to_b64();
258 let mut seen = Dedupe::new(16);
259 assert_eq!(
260 decide(&msg, me.id(), &policy, 1_000, &mut seen),
261 Err(Reject::BadSignature)
262 );
263 }
264
265 #[test]
266 fn message_for_someone_else_is_rejected() {
267 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
268 let other = AgentKey::generate().unwrap();
269 let policy = policy_for(&alice);
270 let msg = alice.sign(other.id(), "hi", 1_000, "m1"); let mut seen = Dedupe::new(16);
272 assert_eq!(
273 decide(&msg, me.id(), &policy, 1_000, &mut seen),
274 Err(Reject::WrongRecipient)
275 );
276 }
277
278 #[test]
279 fn stale_message_is_rejected() {
280 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
281 let policy = policy_for(&alice);
282 let msg = alice.sign(me.id(), "hi", 1_000, "m1");
283 let mut seen = Dedupe::new(16);
284 let far_future = 1_000 + MAX_SKEW_MS + 1;
285 assert_eq!(
286 decide(&msg, me.id(), &policy, far_future, &mut seen),
287 Err(Reject::Stale)
288 );
289 }
290
291 #[test]
292 fn replay_is_rejected_the_second_time() {
293 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
294 let policy = policy_for(&alice);
295 let msg = alice.sign(me.id(), "hi", 1_000, "m1");
296 let mut seen = Dedupe::new(16);
297 assert!(decide(&msg, me.id(), &policy, 1_000, &mut seen).is_ok());
298 assert_eq!(
299 decide(&msg, me.id(), &policy, 1_000, &mut seen),
300 Err(Reject::Replay)
301 );
302 }
303
304 #[test]
305 fn dedupe_forgets_oldest_beyond_cap() {
306 let mut d = Dedupe::new(2);
307 assert!(d.insert("a"));
308 assert!(d.insert("b"));
309 assert!(d.insert("c")); assert!(d.insert("a"), "a was evicted, so it is fresh again");
311 assert!(!d.insert("c"), "c is still within the window");
312 }
313
314 #[test]
315 fn non_peer_knock_is_surfaced_not_dropped() {
316 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
317 let policy = Policy::default(); let msg = stranger.sign_as(
319 me.id(),
320 "stranger-laptop",
321 1_000,
322 "k1",
323 MessageKind::PairRequest,
324 );
325 let mut seen = Dedupe::new(16);
326 assert_eq!(
327 decide(&msg, me.id(), &policy, 1_000, &mut seen),
328 Ok(Dispatch::PairRequest {
329 from_key: stranger.id().to_b64(),
330 name: "stranger-laptop".into()
331 })
332 );
333 }
334
335 #[test]
336 fn non_peer_plain_message_is_still_denied() {
337 let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
338 let policy = Policy::default();
339 let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
340 let mut seen = Dedupe::new(16);
341 assert_eq!(
342 decide(&msg, me.id(), &policy, 1_000, &mut seen),
343 Err(Reject::NotAllowlisted)
344 );
345 }
346
347 #[test]
348 fn pairing_kind_from_existing_peer_is_unexpected() {
349 let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
350 let policy = policy_for(&alice);
351 let msg = alice.sign_as(me.id(), "x", 1_000, "k1", MessageKind::PairRequest);
352 let mut seen = Dedupe::new(16);
353 assert_eq!(
354 decide(&msg, me.id(), &policy, 1_000, &mut seen),
355 Err(Reject::Unexpected)
356 );
357 }
358
359 #[test]
360 fn pair_table_put_take_and_find() {
361 let mut t = PairTable::new(8);
362 t.put("aaaabbbbcccc".into(), "desktop".into());
363 assert_eq!(t.get("aaaabbbbcccc"), Some(&"desktop".to_string()));
364 assert_eq!(
365 t.find("aaaabbbb"), Some(("aaaabbbbcccc".into(), "desktop".into()))
367 );
368 assert_eq!(t.take("aaaabbbbcccc"), Some("desktop".into()));
369 assert!(t.is_empty());
370 }
371}