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