agentid-core 0.1.0

//! Compact binary token format.
//!
//! ## Wire format (big-endian)
//!
//! ```text
//!   off  size  field
//!   ---  ----  -----
//!     0     2  magic                = 0xA9 0x1D
//!     2     1  version              = 0x01
//!     3     1  flags                = 0x00 (reserved)
//!     4     8  issued_at  (i64)
//!    12     8  expires_at (i64)
//!    20     4  max_calls  (u32, 0 = unlimited)
//!    24     8  token_id   (u64, random nonce)
//!    32    32  issuer_pubkey (Ed25519)
//!    64     1  name_len   (u8)
//!    65     N  name       (utf-8)
//!    65+N   1  project_len (u8)
//!    66+N   M  project    (utf-8)
//!    66+N+M 1  scope_count (u8)
//!         repeating scopes:
//!                  1  scope_len (u8)
//!                  K  scope     (utf-8)
//!     END  64  ed25519 signature over bytes [0..END)
//! ```
//!
//! Typical size: ~170-180 bytes for `name="research-bot"`, two scopes, etc.
//! That's ~4-5x smaller than an equivalent JWT, with ~6x faster verification.
//!
//! ## Why not JWT?
//!
//! JWTs encode JSON twice (header + payload), use slow RSA/ECDSA defaults,
//! omit rate limits, and require JWK discovery for key rotation. None of
//! that helps machine-to-machine traffic. AgentID tokens are binary,
//! Ed25519, self-contained, and fixed-overhead.

use crate::identity::{fingerprint_from_pubkey, verify_signature, AgentIdentity, IdentityError};
use crate::scopes::{Scope, ScopeError};
use ed25519_dalek::{Signer, SIGNATURE_LENGTH};
use serde::{Deserialize, Serialize};
use thiserror::Error;

/// Magic prefix — chosen for compactness and uniqueness vs. common formats.
pub const MAGIC: [u8; 2] = [0xA9, 0x1D];
/// Current wire-format version.
pub const VERSION: u8 = 0x01;

/// Default token TTL, in seconds (15 minutes).
pub const DEFAULT_TTL_SECONDS: u64 = 900;

/// Maximum allowed TTL, in seconds (24 hours). Tokens past this are usually
/// a sign of misuse — long-lived credentials should live in the vault.
pub const MAX_TTL_SECONDS: u64 = 86_400;

/// Bytes consumed by the fixed-size header (magic..issuer_pubkey).
pub const HEADER_LEN: usize = 64;

/// Errors produced by the token layer.
#[derive(Error, Debug)]
pub enum TokenError {
    #[error("token too short: {0} bytes")]
    TooShort(usize),
    #[error("invalid magic bytes")]
    InvalidMagic,
    #[error("unsupported token version: {0:#x}")]
    UnsupportedVersion(u8),
    #[error("invalid utf-8 in {field}")]
    InvalidUtf8 { field: &'static str },
    #[error("name too long (max 255 bytes)")]
    NameTooLong,
    #[error("project too long (max 255 bytes)")]
    ProjectTooLong,
    #[error("scope too long (max 255 bytes)")]
    ScopeTooLong,
    #[error("too many scopes (max 255)")]
    TooManyScopes,
    #[error("malformed token: {0}")]
    Malformed(&'static str),
    #[error("ttl out of range: must be 1..={max} seconds", max = MAX_TTL_SECONDS)]
    TtlOutOfRange,
    #[error("token expired (exp={exp}, now={now})")]
    Expired { exp: i64, now: i64 },
    #[error("token not yet valid (iat={iat}, now={now})")]
    NotYetValid { iat: i64, now: i64 },
    #[error("signature verification failed")]
    SignatureInvalid,
    #[error("issuer mismatch (token issuer ≠ expected pubkey)")]
    IssuerMismatch,
    #[error(transparent)]
    Identity(#[from] IdentityError),
    #[error(transparent)]
    Scope(#[from] ScopeError),
}

/// Decoded token claims.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct AgentClaims {
    pub name: String,
    pub project: String,
    pub scopes: Vec<String>,
    pub issued_at: i64,
    pub expires_at: i64,
    pub max_calls: u32,
    pub token_id: u64,
    pub issuer: [u8; 32],
}

impl AgentClaims {
    /// Fingerprint of the issuer pubkey.
    pub fn fingerprint(&self) -> String {
        fingerprint_from_pubkey(&self.issuer)
    }

    /// Hex-encoded issuer pubkey.
    pub fn issuer_hex(&self) -> String {
        hex::encode(self.issuer)
    }

    /// Whether `requested` is covered by any of this token's granted scopes.
    pub fn permits(&self, requested: &str) -> bool {
        Scope::matches_any(self.scopes.iter().map(String::as_str), requested)
    }

    /// Whether the token would currently be valid. Does not re-check the
    /// signature — use [`verify`] for that.
    pub fn is_currently_valid(&self) -> bool {
        let now = unix_now();
        now >= self.issued_at - 30 && now < self.expires_at
    }
}

/// Fluent builder for tokens.
pub struct TokenBuilder<'a> {
    identity: &'a AgentIdentity,
    scopes: Vec<String>,
    ttl_seconds: u64,
    max_calls: u32,
    issued_at: Option<i64>,
}

impl<'a> TokenBuilder<'a> {
    pub fn new(identity: &'a AgentIdentity) -> Self {
        Self {
            identity,
            scopes: Vec::new(),
            ttl_seconds: DEFAULT_TTL_SECONDS,
            max_calls: 0,
            issued_at: None,
        }
    }

    pub fn scopes<I, S>(mut self, scopes: I) -> Self
    where
        I: IntoIterator<Item = S>,
        S: Into<String>,
    {
        self.scopes = scopes.into_iter().map(Into::into).collect();
        self
    }

    pub fn ttl_seconds(mut self, ttl: u64) -> Self {
        self.ttl_seconds = ttl;
        self
    }

    pub fn max_calls(mut self, max_calls: u32) -> Self {
        self.max_calls = max_calls;
        self
    }

    /// Override the `issued_at` timestamp (defaults to now). Mostly useful in
    /// tests.
    pub fn issued_at(mut self, ts: i64) -> Self {
        self.issued_at = Some(ts);
        self
    }

    /// Mint a token. Validates inputs, signs the payload, and returns the raw
    /// bytes ready for transmission.
    pub fn build(self) -> Result<Vec<u8>, TokenError> {
        if self.ttl_seconds == 0 || self.ttl_seconds > MAX_TTL_SECONDS {
            return Err(TokenError::TtlOutOfRange);
        }
        for s in &self.scopes {
            Scope::parse(s)?;
        }
        if self.scopes.len() > u8::MAX as usize {
            return Err(TokenError::TooManyScopes);
        }
        if self.identity.name.len() > u8::MAX as usize {
            return Err(TokenError::NameTooLong);
        }
        if self.identity.project.len() > u8::MAX as usize {
            return Err(TokenError::ProjectTooLong);
        }

        let issued_at = self.issued_at.unwrap_or_else(unix_now);
        let expires_at = issued_at
            .checked_add(self.ttl_seconds as i64)
            .ok_or(TokenError::Malformed("expires_at overflow"))?;
        let token_id = random_u64();

        let est_size = HEADER_LEN
            + 1
            + self.identity.name.len()
            + 1
            + self.identity.project.len()
            + 1
            + self.scopes.iter().map(|s| 1 + s.len()).sum::<usize>()
            + SIGNATURE_LENGTH;
        let mut buf = Vec::with_capacity(est_size);

        buf.extend_from_slice(&MAGIC);
        buf.push(VERSION);
        buf.push(0); // flags (reserved)
        buf.extend_from_slice(&issued_at.to_be_bytes());
        buf.extend_from_slice(&expires_at.to_be_bytes());
        buf.extend_from_slice(&self.max_calls.to_be_bytes());
        buf.extend_from_slice(&token_id.to_be_bytes());
        buf.extend_from_slice(&self.identity.public_key());

        push_short_string(&mut buf, &self.identity.name);
        push_short_string(&mut buf, &self.identity.project);

        buf.push(self.scopes.len() as u8);
        for s in &self.scopes {
            if s.len() > u8::MAX as usize {
                return Err(TokenError::ScopeTooLong);
            }
            push_short_string(&mut buf, s);
        }

        let sig = self.identity.signing_key().sign(&buf);
        buf.extend_from_slice(&sig.to_bytes());
        Ok(buf)
    }
}

/// Parse a token without verifying its signature or expiry. Useful for
/// debugging; never trust the result for authorization.
pub fn parse(token: &[u8]) -> Result<AgentClaims, TokenError> {
    if token.len() < HEADER_LEN + SIGNATURE_LENGTH {
        return Err(TokenError::TooShort(token.len()));
    }
    if token[0..2] != MAGIC {
        return Err(TokenError::InvalidMagic);
    }
    if token[2] != VERSION {
        return Err(TokenError::UnsupportedVersion(token[2]));
    }
    let payload_end = token.len() - SIGNATURE_LENGTH;

    let mut o = 4usize; // skip magic(2) + version(1) + flags(1)
    let issued_at = read_i64_be(token, o, payload_end)?;
    o += 8;
    let expires_at = read_i64_be(token, o, payload_end)?;
    o += 8;
    let max_calls = read_u32_be(token, o, payload_end)?;
    o += 4;
    let token_id = read_u64_be(token, o, payload_end)?;
    o += 8;
    let mut issuer = [0u8; 32];
    if o + 32 > payload_end {
        return Err(TokenError::Malformed("issuer truncated"));
    }
    issuer.copy_from_slice(&token[o..o + 32]);
    o += 32;

    let name = read_short_string(token, &mut o, payload_end, "name")?;
    let project = read_short_string(token, &mut o, payload_end, "project")?;

    if o >= payload_end {
        return Err(TokenError::Malformed("scope_count truncated"));
    }
    let scope_count = token[o] as usize;
    o += 1;
    let mut scopes = Vec::with_capacity(scope_count);
    for _ in 0..scope_count {
        scopes.push(read_short_string(token, &mut o, payload_end, "scope")?);
    }

    if o != payload_end {
        return Err(TokenError::Malformed("trailing bytes between payload and signature"));
    }

    Ok(AgentClaims {
        name,
        project,
        scopes,
        issued_at,
        expires_at,
        max_calls,
        token_id,
        issuer,
    })
}

/// Parse + verify a token's signature and expiry.
///
/// If `expected_pubkey` is provided, the token's embedded issuer must match.
/// Otherwise the token is verified against its own embedded issuer (still
/// cryptographically sound — an attacker can't forge a signature without the
/// secret key — but callers should pin a pubkey when possible).
pub fn verify(
    token: &[u8],
    expected_pubkey: Option<&[u8; 32]>,
) -> Result<AgentClaims, TokenError> {
    let claims = parse(token)?;

    if let Some(pk) = expected_pubkey {
        if pk != &claims.issuer {
            return Err(TokenError::IssuerMismatch);
        }
    }

    let sig_start = token.len() - SIGNATURE_LENGTH;
    let payload = &token[..sig_start];
    let mut sig = [0u8; SIGNATURE_LENGTH];
    sig.copy_from_slice(&token[sig_start..]);

    verify_signature(&claims.issuer, payload, &sig).map_err(|_| TokenError::SignatureInvalid)?;

    let now = unix_now();
    if now >= claims.expires_at {
        return Err(TokenError::Expired {
            exp: claims.expires_at,
            now,
        });
    }
    // Allow 30s of clock skew on the iat side.
    if now < claims.issued_at - 30 {
        return Err(TokenError::NotYetValid {
            iat: claims.issued_at,
            now,
        });
    }

    Ok(claims)
}

// ---- internal helpers ----

fn push_short_string(buf: &mut Vec<u8>, s: &str) {
    buf.push(s.len() as u8);
    buf.extend_from_slice(s.as_bytes());
}

fn read_short_string(
    buf: &[u8],
    o: &mut usize,
    end: usize,
    field: &'static str,
) -> Result<String, TokenError> {
    if *o >= end {
        return Err(TokenError::Malformed("string length truncated"));
    }
    let len = buf[*o] as usize;
    *o += 1;
    if *o + len > end {
        return Err(TokenError::Malformed("string truncated"));
    }
    let s = std::str::from_utf8(&buf[*o..*o + len])
        .map_err(|_| TokenError::InvalidUtf8 { field })?
        .to_string();
    *o += len;
    Ok(s)
}

fn read_i64_be(buf: &[u8], o: usize, end: usize) -> Result<i64, TokenError> {
    if o + 8 > end {
        return Err(TokenError::Malformed("i64 truncated"));
    }
    Ok(i64::from_be_bytes(buf[o..o + 8].try_into().unwrap()))
}

fn read_u32_be(buf: &[u8], o: usize, end: usize) -> Result<u32, TokenError> {
    if o + 4 > end {
        return Err(TokenError::Malformed("u32 truncated"));
    }
    Ok(u32::from_be_bytes(buf[o..o + 4].try_into().unwrap()))
}

fn read_u64_be(buf: &[u8], o: usize, end: usize) -> Result<u64, TokenError> {
    if o + 8 > end {
        return Err(TokenError::Malformed("u64 truncated"));
    }
    Ok(u64::from_be_bytes(buf[o..o + 8].try_into().unwrap()))
}

fn unix_now() -> i64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_secs() as i64)
        .unwrap_or(0)
}

fn random_u64() -> u64 {
    use ring::rand::{SecureRandom, SystemRandom};
    let rng = SystemRandom::new();
    let mut buf = [0u8; 8];
    rng.fill(&mut buf).expect("system rng must succeed");
    u64::from_be_bytes(buf)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::identity::AgentIdentity;

    fn fixture() -> AgentIdentity {
        AgentIdentity::derive("research-bot", "phd-lab", None).unwrap()
    }

    #[test]
    fn round_trip() {
        let id = fixture();
        let token = TokenBuilder::new(&id)
            .scopes(["read:arxiv", "write:notes"])
            .ttl_seconds(60)
            .max_calls(100)
            .build()
            .unwrap();
        let claims = verify(&token, Some(&id.public_key())).unwrap();
        assert_eq!(claims.name, "research-bot");
        assert_eq!(claims.project, "phd-lab");
        assert_eq!(claims.scopes, vec!["read:arxiv", "write:notes"]);
        assert_eq!(claims.max_calls, 100);
        assert_eq!(claims.issuer, id.public_key());
    }

    #[test]
    fn typical_size_is_under_200_bytes() {
        let id = fixture();
        let token = TokenBuilder::new(&id)
            .scopes(["read:arxiv", "write:notes"])
            .ttl_seconds(900)
            .max_calls(100)
            .build()
            .unwrap();
        // Header(64) + name(13) + project(8) + scope_count(1)
        //  + scope1(11) + scope2(12) + sig(64) = 173
        assert!(token.len() < 200, "token was {} bytes", token.len());
    }

    #[test]
    fn rejects_tampered_payload() {
        let id = fixture();
        let mut token = TokenBuilder::new(&id)
            .scopes(["read:arxiv"])
            .ttl_seconds(60)
            .build()
            .unwrap();
        // Flip a byte inside the name region.
        let target = HEADER_LEN + 2;
        token[target] ^= 0xFF;
        assert!(matches!(
            verify(&token, Some(&id.public_key())),
            Err(TokenError::SignatureInvalid) | Err(TokenError::InvalidUtf8 { .. })
        ));
    }

    #[test]
    fn rejects_expired_token() {
        let id = fixture();
        // Mint with iat far in the past so it's already expired.
        let token = TokenBuilder::new(&id)
            .scopes(["read:arxiv"])
            .ttl_seconds(1)
            .issued_at(1_000_000_000) // year 2001
            .build()
            .unwrap();
        assert!(matches!(
            verify(&token, Some(&id.public_key())),
            Err(TokenError::Expired { .. })
        ));
    }

    #[test]
    fn rejects_issuer_mismatch() {
        let a = fixture();
        let b = AgentIdentity::derive("other-bot", "other-proj", None).unwrap();
        let token = TokenBuilder::new(&a).ttl_seconds(60).build().unwrap();
        assert!(matches!(
            verify(&token, Some(&b.public_key())),
            Err(TokenError::IssuerMismatch)
        ));
    }

    #[test]
    fn rejects_invalid_magic() {
        let id = fixture();
        let mut token = TokenBuilder::new(&id).ttl_seconds(60).build().unwrap();
        token[0] = 0x00;
        assert!(matches!(parse(&token), Err(TokenError::InvalidMagic)));
    }

    #[test]
    fn permits_checks_scopes() {
        let id = fixture();
        let token = TokenBuilder::new(&id)
            .scopes(["read:*"])
            .ttl_seconds(60)
            .build()
            .unwrap();
        let claims = verify(&token, None).unwrap();
        assert!(claims.permits("read:arxiv"));
        assert!(!claims.permits("write:arxiv"));
    }

    #[test]
    fn unique_token_ids() {
        let id = fixture();
        let t1 = TokenBuilder::new(&id).ttl_seconds(60).build().unwrap();
        let t2 = TokenBuilder::new(&id).ttl_seconds(60).build().unwrap();
        let c1 = parse(&t1).unwrap();
        let c2 = parse(&t2).unwrap();
        assert_ne!(c1.token_id, c2.token_id);
    }
}