use std::collections::{HashMap, VecDeque};
use crate::identity::{AgentId, MessageKind, SignedMessage, TaskStatus, check_freshness};
use crate::policy::Policy;
pub const MAX_SKEW_MS: u64 = 60_000;
#[derive(Debug, PartialEq, Eq)]
pub enum Dispatch {
Inline {
petname: String,
text: String,
task_id: Option<String>,
status: Option<TaskStatus>,
in_reply_to: Option<String>,
},
PairRequest { from_key: String, name: String },
PairAccept { from_key: String, name: String },
}
#[derive(Debug, PartialEq, Eq)]
pub enum Reject {
BadSignature,
NotAllowlisted,
WrongRecipient,
Stale,
Replay,
Unexpected,
}
pub type Verdict = Result<Dispatch, Reject>;
pub fn decide(
msg: &SignedMessage,
me: AgentId,
policy: &Policy,
now: u64,
seen: &mut Dedupe,
) -> Verdict {
let from = msg.verify().map_err(|_| Reject::BadSignature)?;
let to = AgentId::from_b64(&msg.to).map_err(|_| Reject::WrongRecipient)?;
if to != me {
return Err(Reject::WrongRecipient);
}
let peer = policy.peer(from);
if peer.is_none() && msg.kind == MessageKind::Message {
return Err(Reject::NotAllowlisted);
}
check_freshness(msg.ts, now, MAX_SKEW_MS).map_err(|_| Reject::Stale)?;
if !seen.insert(&msg.msg_id) {
return Err(Reject::Replay);
}
match (peer, msg.kind) {
(Some(peer), MessageKind::Message) => Ok(Dispatch::Inline {
petname: peer.petname.clone(),
text: msg.text.clone(),
task_id: msg.task_id.clone(),
status: msg.status,
in_reply_to: msg.in_reply_to.clone(),
}),
(None, MessageKind::PairRequest) => Ok(Dispatch::PairRequest {
from_key: from.to_b64(),
name: msg.text.clone(),
}),
(None, MessageKind::PairAccept) => Ok(Dispatch::PairAccept {
from_key: from.to_b64(),
name: msg.text.clone(),
}),
_ => Err(Reject::Unexpected),
}
}
pub struct PairTable {
order: VecDeque<String>,
map: HashMap<String, String>,
cap: usize,
}
impl PairTable {
pub fn new(cap: usize) -> Self {
Self {
order: VecDeque::new(),
map: HashMap::new(),
cap: cap.max(1),
}
}
pub fn put(&mut self, key: String, value: String) {
if !self.map.contains_key(&key) {
if self.order.len() >= self.cap
&& let Some(old) = self.order.pop_front()
{
self.map.remove(&old);
}
self.order.push_back(key.clone());
}
self.map.insert(key, value);
}
pub fn take(&mut self, key: &str) -> Option<String> {
let v = self.map.remove(key)?;
self.order.retain(|k| k != key);
Some(v)
}
pub fn get(&self, key: &str) -> Option<&String> {
self.map.get(key)
}
pub fn find(&self, key_or_fp: &str) -> Option<(String, String)> {
self.map
.iter()
.find(|(k, _)| {
k.as_str() == key_or_fp || k.chars().take(8).collect::<String>() == key_or_fp
})
.map(|(k, v)| (k.clone(), v.clone()))
}
pub fn entries(&self) -> impl Iterator<Item = (&String, &String)> {
self.map.iter()
}
pub fn is_empty(&self) -> bool {
self.map.is_empty()
}
}
pub struct Dedupe {
order: VecDeque<String>,
seen: std::collections::HashSet<String>,
cap: usize,
}
impl Dedupe {
pub fn new(cap: usize) -> Self {
Self {
order: VecDeque::with_capacity(cap),
seen: std::collections::HashSet::with_capacity(cap),
cap: cap.max(1),
}
}
pub fn insert(&mut self, id: &str) -> bool {
if self.seen.contains(id) {
return false;
}
if self.order.len() >= self.cap
&& let Some(old) = self.order.pop_front()
{
self.seen.remove(&old);
}
self.order.push_back(id.to_string());
self.seen.insert(id.to_string());
true
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::identity::AgentKey;
use crate::policy::Policy;
fn policy_for(key: &AgentKey) -> Policy {
let raw = format!(r#"{{ "alice": {{ "key": "{}" }} }}"#, key.id().to_b64());
Policy::parse(&raw).unwrap()
}
#[test]
fn admitted_peer_is_dispatched_inline() {
let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
let policy = policy_for(&alice);
let msg = alice.sign(me.id(), "run the deploy", 1_000, "m1");
let mut seen = Dedupe::new(16);
assert_eq!(
decide(&msg, me.id(), &policy, 1_000, &mut seen),
Ok(Dispatch::Inline {
petname: "alice".into(),
text: "run the deploy".into(),
task_id: None,
status: None,
in_reply_to: None,
})
);
}
#[test]
fn stranger_is_rejected_even_with_valid_signature() {
let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
let policy = policy_for(&AgentKey::generate().unwrap()); let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
let mut seen = Dedupe::new(16);
assert_eq!(
decide(&msg, me.id(), &policy, 1_000, &mut seen),
Err(Reject::NotAllowlisted)
);
}
#[test]
fn forged_sender_is_bad_signature() {
let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
let policy = policy_for(&alice);
let eve = AgentKey::generate().unwrap();
let mut msg = eve.sign(me.id(), "hi", 1_000, "m1");
msg.from = alice.id().to_b64();
let mut seen = Dedupe::new(16);
assert_eq!(
decide(&msg, me.id(), &policy, 1_000, &mut seen),
Err(Reject::BadSignature)
);
}
#[test]
fn message_for_someone_else_is_rejected() {
let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
let other = AgentKey::generate().unwrap();
let policy = policy_for(&alice);
let msg = alice.sign(other.id(), "hi", 1_000, "m1"); let mut seen = Dedupe::new(16);
assert_eq!(
decide(&msg, me.id(), &policy, 1_000, &mut seen),
Err(Reject::WrongRecipient)
);
}
#[test]
fn stale_message_is_rejected() {
let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
let policy = policy_for(&alice);
let msg = alice.sign(me.id(), "hi", 1_000, "m1");
let mut seen = Dedupe::new(16);
let far_future = 1_000 + MAX_SKEW_MS + 1;
assert_eq!(
decide(&msg, me.id(), &policy, far_future, &mut seen),
Err(Reject::Stale)
);
}
#[test]
fn replay_is_rejected_the_second_time() {
let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
let policy = policy_for(&alice);
let msg = alice.sign(me.id(), "hi", 1_000, "m1");
let mut seen = Dedupe::new(16);
assert!(decide(&msg, me.id(), &policy, 1_000, &mut seen).is_ok());
assert_eq!(
decide(&msg, me.id(), &policy, 1_000, &mut seen),
Err(Reject::Replay)
);
}
#[test]
fn dedupe_forgets_oldest_beyond_cap() {
let mut d = Dedupe::new(2);
assert!(d.insert("a"));
assert!(d.insert("b"));
assert!(d.insert("c")); assert!(d.insert("a"), "a was evicted, so it is fresh again");
assert!(!d.insert("c"), "c is still within the window");
}
#[test]
fn non_peer_knock_is_surfaced_not_dropped() {
let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
let policy = Policy::default(); let msg = stranger.sign_as(
me.id(),
"stranger-laptop",
1_000,
"k1",
MessageKind::PairRequest,
);
let mut seen = Dedupe::new(16);
assert_eq!(
decide(&msg, me.id(), &policy, 1_000, &mut seen),
Ok(Dispatch::PairRequest {
from_key: stranger.id().to_b64(),
name: "stranger-laptop".into()
})
);
}
#[test]
fn non_peer_plain_message_is_still_denied() {
let (stranger, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
let policy = Policy::default();
let msg = stranger.sign(me.id(), "hi", 1_000, "m1");
let mut seen = Dedupe::new(16);
assert_eq!(
decide(&msg, me.id(), &policy, 1_000, &mut seen),
Err(Reject::NotAllowlisted)
);
}
#[test]
fn pairing_kind_from_existing_peer_is_unexpected() {
let (alice, me) = (AgentKey::generate().unwrap(), AgentKey::generate().unwrap());
let policy = policy_for(&alice);
let msg = alice.sign_as(me.id(), "x", 1_000, "k1", MessageKind::PairRequest);
let mut seen = Dedupe::new(16);
assert_eq!(
decide(&msg, me.id(), &policy, 1_000, &mut seen),
Err(Reject::Unexpected)
);
}
#[test]
fn pair_table_put_take_and_find() {
let mut t = PairTable::new(8);
t.put("aaaabbbbcccc".into(), "desktop".into());
assert_eq!(t.get("aaaabbbbcccc"), Some(&"desktop".to_string()));
assert_eq!(
t.find("aaaabbbb"), Some(("aaaabbbbcccc".into(), "desktop".into()))
);
assert_eq!(t.take("aaaabbbbcccc"), Some("desktop".into()));
assert!(t.is_empty());
}
}