agent_rooms/

protocol.rs

//! Pure protocol logic — signed-payload builders, freshness check, turn
//! rotation, replay dedup. Mirrors the non-DB parts of
//! `parley/services/{rooms,messages,participants,dedup}.py`.

use std::collections::{HashMap, HashSet};
use std::sync::Mutex;
use std::time::{Duration, SystemTime};

use chrono::{DateTime, Utc};
use once_cell::sync::Lazy;
use serde_json::{json, Map, Value};
use sha2::{Digest, Sha256};

use crate::canonical::canonical_json;
use crate::error::Error;
use crate::keys::{verify, PUBKEY_LEN};
use crate::models::Participant;

/// SPEC §6.6: `created_at` must be within ±60 s of server `now`.
pub const CLOCK_SKEW_SECS: i64 = 60;

/// SPEC C9: message body max bytes.
pub const MAX_BODY_BYTES: usize = 16 * 1024;

/// SPEC C5: room topic minimum character count.
pub const TOPIC_MIN_CHARS: usize = 1;
/// SPEC C5: room topic maximum character count.
pub const TOPIC_MAX_CHARS: usize = 256;

/// SPEC C6: minimum allowed `max_turns`.
pub const MAX_TURNS_MIN: i64 = 1;
/// SPEC C6: maximum allowed `max_turns`.
pub const MAX_TURNS_MAX: i64 = 1000;

/// Minimum allowed `ttl_hours` on room creation.
pub const TTL_HOURS_MIN: i64 = 1;
/// Maximum allowed `ttl_hours` on room creation.
pub const TTL_HOURS_MAX: i64 = 720;

// ---------------------------------------------------------------------------
// Signed-payload builders — Appendix A
// ---------------------------------------------------------------------------

/// Build the canonical signed bytes for `POST /v1/rooms` (C11).
/// Sorted keys: `created_at, invite_pubkeys, max_turns, topic, ttl_hours`.
pub fn create_room_payload(
    topic: &str,
    invite_pubkeys: &[String],
    max_turns: i64,
    ttl_hours: i64,
    created_at: &DateTime<Utc>,
) -> Vec<u8> {
    let v = json!({
        "topic": topic,
        "invite_pubkeys": invite_pubkeys,
        "max_turns": max_turns,
        "ttl_hours": ttl_hours,
        "created_at": iso8601_python(created_at),
    });
    canonical_json(&v)
}

/// Build the canonical signed bytes for `POST /v1/rooms/{id}/accept` (C13).
/// Sorted keys: `agent_pubkey, created_at, room_id`.
pub fn accept_payload(room_id: &str, agent_pubkey: &str, created_at: &DateTime<Utc>) -> Vec<u8> {
    let v = json!({
        "room_id": room_id,
        "agent_pubkey": agent_pubkey,
        "created_at": iso8601_python(created_at),
    });
    canonical_json(&v)
}

/// Build the canonical signed bytes for `POST /v1/rooms/{id}/close` (C14).
/// Sorted keys: `created_at, room_id, summary`. `summary` of `None` is
/// emitted as literal JSON `null`.
pub fn close_payload(room_id: &str, summary: Option<&str>, created_at: &DateTime<Utc>) -> Vec<u8> {
    let mut obj = Map::new();
    obj.insert("room_id".to_string(), Value::String(room_id.to_string()));
    obj.insert(
        "summary".to_string(),
        match summary {
            Some(s) => Value::String(s.to_string()),
            None => Value::Null,
        },
    );
    obj.insert(
        "created_at".to_string(),
        Value::String(iso8601_python(created_at)),
    );
    canonical_json(&Value::Object(obj))
}

/// Build the canonical signed bytes for `POST /v1/rooms/{id}/messages` (C16).
/// Sorted keys: `author_pubkey, body, created_at, room_id, turn_n`.
pub fn post_message_payload(
    room_id: &str,
    author_pubkey: &str,
    turn_n: i64,
    body: &str,
    created_at: &DateTime<Utc>,
) -> Vec<u8> {
    let v = json!({
        "room_id": room_id,
        "turn_n": turn_n,
        "author_pubkey": author_pubkey,
        "body": body,
        "created_at": iso8601_python(created_at),
    });
    canonical_json(&v)
}

/// Format a UTC timestamp the way Python's `datetime.isoformat()` does for
/// tz-aware UTC: `YYYY-MM-DDThh:mm:ss[.ffffff]+00:00`. Microsecond precision
/// is preserved when non-zero. The `+00:00` suffix (not `Z`) is required by
/// SPEC §4 clause 7.
pub fn iso8601_python(dt: &DateTime<Utc>) -> String {
    let micros = dt.timestamp_subsec_micros();
    if micros == 0 {
        dt.format("%Y-%m-%dT%H:%M:%S+00:00").to_string()
    } else {
        format!("{}.{:06}+00:00", dt.format("%Y-%m-%dT%H:%M:%S"), micros)
    }
}

// ---------------------------------------------------------------------------
// Freshness — SPEC §6.6 / C20
// ---------------------------------------------------------------------------

/// Returns true iff `ts` is within ±`CLOCK_SKEW_SECS` of `now`.
pub fn is_timestamp_fresh(ts: &DateTime<Utc>, now: &DateTime<Utc>) -> bool {
    (ts.timestamp() - now.timestamp()).abs() <= CLOCK_SKEW_SECS
}

/// Same, against the system clock.
pub fn check_freshness(ts: &DateTime<Utc>) -> Result<(), Error> {
    if is_timestamp_fresh(ts, &Utc::now()) {
        Ok(())
    } else {
        Err(Error::StaleTimestamp)
    }
}

// ---------------------------------------------------------------------------
// Signature verification
// ---------------------------------------------------------------------------

/// Verify an Ed25519 signature against the canonical bytes of a signed
/// payload. Returns `Err(BadSignature)` on any failure.
pub fn verify_signed(
    pubkey: &[u8; PUBKEY_LEN],
    canonical_bytes: &[u8],
    signature: &[u8],
) -> Result<(), Error> {
    if verify(pubkey, canonical_bytes, signature) {
        Ok(())
    } else {
        Err(Error::BadSignature)
    }
}

// ---------------------------------------------------------------------------
// Turn rotation — SPEC §8.2 / C24
// ---------------------------------------------------------------------------

/// Round-robin over **accepted** participants ordered by `invited_at`.
///
/// - If no participant is accepted, returns `None`.
/// - If `current` is not in the accepted set, returns the first accepted.
/// - Otherwise returns `A[(i+1) mod |A|]`.
pub fn next_turn_owner<'a>(
    participants: &'a [Participant],
    current: Option<&str>,
) -> Option<&'a str> {
    let mut accepted: Vec<&Participant> = participants
        .iter()
        .filter(|p| p.accepted_at.is_some())
        .collect();
    accepted.sort_by_key(|p| p.invited_at);
    if accepted.is_empty() {
        return None;
    }
    let pks: Vec<&str> = accepted.iter().map(|p| p.agent_pubkey.as_str()).collect();
    let idx = match current.and_then(|c| pks.iter().position(|p| *p == c)) {
        Some(i) => (i + 1) % pks.len(),
        None => 0,
    };
    Some(pks[idx])
}

// ---------------------------------------------------------------------------
// Invite-list deduplication — SPEC §6.1
// ---------------------------------------------------------------------------

/// Drop duplicates from `invite_pubkeys`, preserving first-occurrence order,
/// and drop any entry equal to the creator. Mirrors the behaviour of
/// `parley/services/rooms.py:create_room`.
pub fn dedup_invites(creator_pubkey: &str, invite_pubkeys: &[String]) -> Vec<String> {
    let mut seen: HashSet<&str> = HashSet::new();
    seen.insert(creator_pubkey);
    let mut out = Vec::with_capacity(invite_pubkeys.len());
    for pk in invite_pubkeys {
        if seen.insert(pk.as_str()) {
            out.push(pk.clone());
        }
    }
    out
}

// ---------------------------------------------------------------------------
// Post-message preconditions — SPEC §6.6
// ---------------------------------------------------------------------------

/// Returns Ok iff every precondition for a `post_message` write passes.
/// Order matches SPEC §6.6. Caller passes the relevant room/participant
/// state extracted from their own store.
#[allow(clippy::too_many_arguments)]
pub fn check_post_message(
    body: &str,
    room_status: &str,
    room_ttl_until: &DateTime<Utc>,
    caller_is_accepted_participant: bool,
    room_turn_owner: Option<&str>,
    caller_pubkey: &str,
    expected_next_turn: i64,
    got_turn: i64,
    created_at: &DateTime<Utc>,
    now: &DateTime<Utc>,
) -> Result<(), Error> {
    if body.len() > MAX_BODY_BYTES {
        return Err(Error::BodyTooLarge);
    }
    if room_status == "closed" || now >= room_ttl_until {
        return Err(Error::RoomClosed);
    }
    if !caller_is_accepted_participant {
        return Err(Error::NotAParticipant);
    }
    match room_turn_owner {
        Some(owner) if owner == caller_pubkey => {}
        _ => return Err(Error::NotTurnOwner),
    }
    if got_turn != expected_next_turn {
        return Err(Error::TurnConflict {
            expected: expected_next_turn,
            got: got_turn,
        });
    }
    if !is_timestamp_fresh(created_at, now) {
        return Err(Error::StaleTimestamp);
    }
    Ok(())
}

// ---------------------------------------------------------------------------
// Replay dedup — SPEC §10.1 / C28
// ---------------------------------------------------------------------------

/// In-memory replay-detection store. Single-process only; SPEC §10.2 notes
/// that a multi-worker deployment needs a shared backing store.
///
/// `check_and_mark` returns `Ok(())` if the bytes are new in the window, and
/// `Err(ReplayDetected)` if they've been seen.
pub struct DedupStore {
    inner: Mutex<HashMap<String, SystemTime>>,
    window: Duration,
}

impl DedupStore {
    /// Construct a new store with the given replay-window duration in seconds.
    pub fn new(window_secs: u64) -> Self {
        Self {
            inner: Mutex::new(HashMap::new()),
            window: Duration::from_secs(window_secs),
        }
    }

    /// Returns `Ok(())` if `canonical_bytes` is new in the window; stores it.
    /// Returns `Err(ReplayDetected)` if the same bytes were already accepted
    /// within `window_secs`. Garbage-collects expired entries on every call.
    pub fn check_and_mark(&self, canonical_bytes: &[u8]) -> Result<(), Error> {
        let now = SystemTime::now();
        let mut hasher = Sha256::new();
        hasher.update(canonical_bytes);
        let h = hex::encode(hasher.finalize());

        let mut map = self.inner.lock().expect("dedup store poisoned");
        // GC expired entries.
        map.retain(|_, exp| *exp >= now);
        if map.contains_key(&h) {
            return Err(Error::ReplayDetected);
        }
        map.insert(h, now + self.window);
        Ok(())
    }

    /// Test helper — clears the store.
    pub fn reset(&self) {
        self.inner.lock().expect("dedup store poisoned").clear();
    }
}

impl Default for DedupStore {
    fn default() -> Self {
        Self::new(CLOCK_SKEW_SECS as u64)
    }
}

/// Shared default dedup store, mirroring the module-level `_store` in
/// `parley/services/dedup.py`. Most embedders will prefer an owned
/// `DedupStore` instance — this exists for parity with the reference impl.
pub static DEFAULT_DEDUP: Lazy<DedupStore> = Lazy::new(DedupStore::default);

// ---------------------------------------------------------------------------
// Helper: parse the wire header pubkey
// ---------------------------------------------------------------------------

/// Validates `X-Agent-Pubkey` header value (SPEC §3 / C3). Bare lowercase
/// hex, exactly 64 chars. Anything else → `InvalidPubkey`.
pub fn parse_header_pubkey(value: &str) -> Result<[u8; PUBKEY_LEN], Error> {
    if value.len() != PUBKEY_LEN * 2
        || !value
            .chars()
            .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())
    {
        return Err(Error::InvalidPubkey);
    }
    crate::keys::pubkey_from_hex(value)
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::TimeZone;

    fn p(agent: &str, invited_secs: i64, accepted: bool) -> Participant {
        Participant {
            room_id: "r".into(),
            agent_pubkey: agent.into(),
            owner_pubkey: agent.into(),
            invited_by_pubkey: "creator".into(),
            invited_at: Utc.timestamp_opt(invited_secs, 0).unwrap(),
            accepted_at: if accepted {
                Some(Utc.timestamp_opt(invited_secs + 1, 0).unwrap())
            } else {
                None
            },
            accept_sig: None,
        }
    }

    #[test]
    fn rotation_round_robin() {
        let ps = vec![p("a", 0, true), p("b", 1, true), p("c", 2, true)];
        assert_eq!(next_turn_owner(&ps, Some("a")), Some("b"));
        assert_eq!(next_turn_owner(&ps, Some("b")), Some("c"));
        assert_eq!(next_turn_owner(&ps, Some("c")), Some("a"));
    }

    #[test]
    fn rotation_skips_pending() {
        let ps = vec![p("a", 0, true), p("b", 1, false), p("c", 2, true)];
        assert_eq!(next_turn_owner(&ps, Some("a")), Some("c"));
        assert_eq!(next_turn_owner(&ps, Some("c")), Some("a"));
    }

    #[test]
    fn dedup_drops_creator() {
        let out = dedup_invites(
            "creator",
            &["creator".into(), "alice".into(), "alice".into()],
        );
        assert_eq!(out, vec!["alice".to_string()]);
    }

    #[test]
    fn freshness_window() {
        let now = Utc.timestamp_opt(1_000_000, 0).unwrap();
        let inside = Utc.timestamp_opt(1_000_000 + 30, 0).unwrap();
        let outside = Utc.timestamp_opt(1_000_000 + 90, 0).unwrap();
        assert!(is_timestamp_fresh(&inside, &now));
        assert!(!is_timestamp_fresh(&outside, &now));
    }

    #[test]
    fn dedup_replay() {
        let store = DedupStore::new(60);
        let payload = b"{\"hello\":\"world\"}";
        assert!(store.check_and_mark(payload).is_ok());
        assert_eq!(store.check_and_mark(payload), Err(Error::ReplayDetected));
        store.reset();
        assert!(store.check_and_mark(payload).is_ok());
    }

    #[test]
    fn header_pubkey_rejects_uppercase() {
        // SPEC C1: bare lowercase hex.
        let lower = "8a88e3dd7409f195fd52db2d3cba5d72ca6709bf1d94121bf3748801b40f6f5c";
        assert!(parse_header_pubkey(lower).is_ok());
        let upper = "8A88E3DD7409F195FD52DB2D3CBA5D72CA6709BF1D94121BF3748801B40F6F5C";
        assert_eq!(parse_header_pubkey(upper), Err(Error::InvalidPubkey));
    }
}