agent_protocols/
identity.rs

1use base64::engine::general_purpose::URL_SAFE_NO_PAD;
2use base64::Engine as _;
3use ed25519_dalek::{Signature, Signer, SigningKey, VerifyingKey};
4use serde::de::Error as DeError;
5use serde::{Deserialize, Deserializer, Serialize, Serializer};
6use serde_json::Value;
7use sha3::{Digest, Sha3_256};
8use std::collections::{BTreeMap, HashMap};
9use std::fmt;
10use std::str::FromStr;
11use std::time::{SystemTime, UNIX_EPOCH};
12
13use crate::error::{Result, SdkError};
14
15pub const AGENT_ID_PREFIX: &str = "did:agent:";
16pub const DEFAULT_LIVE_WRITE_WINDOW_MS: i64 = 300_000;
17pub const DEFAULT_NONCE_TTL_MS: i64 = 300_000;
18pub const DEFAULT_REQUEST_JWT_TTL_SECS: i64 = 300;
19pub const MAX_NONCE_HEADER: &str = "Max-Seen-Nonce";
20
21#[derive(Clone, Eq, PartialEq, Ord, PartialOrd, Hash)]
22pub struct AgentId(String);
23
24impl AgentId {
25    pub fn from_public_key(public_key: &[u8; 32]) -> Self {
26        Self(format!(
27            "{}{}",
28            AGENT_ID_PREFIX,
29            URL_SAFE_NO_PAD.encode(public_key)
30        ))
31    }
32
33    pub fn as_str(&self) -> &str {
34        &self.0
35    }
36
37    pub fn public_key_bytes(&self) -> Result<[u8; 32]> {
38        let encoded = self
39            .0
40            .strip_prefix(AGENT_ID_PREFIX)
41            .ok_or(SdkError::InvalidAgentIdPrefix)?;
42        let bytes = URL_SAFE_NO_PAD.decode(encoded)?;
43        bytes
44            .try_into()
45            .map_err(|bytes: Vec<u8>| SdkError::InvalidPublicKeyLength(bytes.len()))
46    }
47
48    pub fn verifying_key(&self) -> Result<VerifyingKey> {
49        Ok(VerifyingKey::from_bytes(&self.public_key_bytes()?)?)
50    }
51}
52
53impl fmt::Debug for AgentId {
54    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
55        formatter.debug_tuple("AgentId").field(&self.0).finish()
56    }
57}
58
59impl fmt::Display for AgentId {
60    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
61        formatter.write_str(&self.0)
62    }
63}
64
65impl FromStr for AgentId {
66    type Err = SdkError;
67
68    fn from_str(value: &str) -> Result<Self> {
69        if !value.starts_with(AGENT_ID_PREFIX) {
70            return Err(SdkError::InvalidAgentIdPrefix);
71        }
72        let agent_id = Self(value.to_owned());
73        agent_id.public_key_bytes()?;
74        Ok(agent_id)
75    }
76}
77
78impl TryFrom<String> for AgentId {
79    type Error = SdkError;
80
81    fn try_from(value: String) -> Result<Self> {
82        Self::from_str(&value)
83    }
84}
85
86impl Serialize for AgentId {
87    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
88    where
89        S: Serializer,
90    {
91        serializer.serialize_str(&self.0)
92    }
93}
94
95impl<'de> Deserialize<'de> for AgentId {
96    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
97    where
98        D: Deserializer<'de>,
99    {
100        let value = String::deserialize(deserializer)?;
101        Self::from_str(&value).map_err(D::Error::custom)
102    }
103}
104
105pub struct AgentSigner {
106    signing_key: SigningKey,
107}
108
109impl AgentSigner {
110    pub fn generate() -> Self {
111        let mut seed = [0_u8; 32];
112        rand::fill(&mut seed);
113        Self::from_seed(seed)
114    }
115
116    pub fn from_seed(seed: [u8; 32]) -> Self {
117        Self {
118            signing_key: SigningKey::from_bytes(&seed),
119        }
120    }
121
122    pub fn agent_id(&self) -> AgentId {
123        AgentId::from_public_key(&self.signing_key.verifying_key().to_bytes())
124    }
125
126    pub fn verifying_key(&self) -> VerifyingKey {
127        self.signing_key.verifying_key()
128    }
129
130    pub fn sign_event<P>(&self, event: Event<P>) -> Result<Envelope<P>>
131    where
132        P: Serialize,
133    {
134        let hash = event_hash(&event)?;
135        let signature = sign_event(&self.signing_key, &event)?;
136        Ok(Envelope {
137            hash,
138            event,
139            signature,
140        })
141    }
142
143    pub fn sign_request_jwt(&self, claims: &RequestJwtClaims) -> Result<String> {
144        let agent_id = self.agent_id();
145        if claims.iss != agent_id || claims.sub != agent_id {
146            return Err(SdkError::InvalidJwtClaim("iss/sub"));
147        }
148
149        let header = RequestJwtHeader {
150            alg: "EdDSA".to_owned(),
151            typ: "JWT".to_owned(),
152            kid: agent_id,
153        };
154        let header = URL_SAFE_NO_PAD.encode(serde_json::to_vec(&header)?);
155        let payload = URL_SAFE_NO_PAD.encode(serde_json::to_vec(claims)?);
156        let signing_input = format!("{header}.{payload}");
157        let signature = self.signing_key.sign(signing_input.as_bytes());
158        Ok(format!(
159            "{signing_input}.{}",
160            URL_SAFE_NO_PAD.encode(signature.to_bytes())
161        ))
162    }
163}
164
165#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
166pub struct Event<P = Value> {
167    pub protocol: String,
168    #[serde(rename = "type")]
169    pub kind: String,
170    pub actor: AgentId,
171    pub created_at: i64,
172    pub nonce: u64,
173    #[serde(skip_serializing_if = "Option::is_none")]
174    pub room_id: Option<String>,
175    pub payload: P,
176    #[serde(default, flatten)]
177    pub extra: BTreeMap<String, Value>,
178}
179
180impl<P> Event<P> {
181    pub fn new(
182        protocol: impl Into<String>,
183        kind: impl Into<String>,
184        actor: AgentId,
185        created_at: i64,
186        nonce: u64,
187        payload: P,
188    ) -> Self {
189        Self {
190            protocol: protocol.into(),
191            kind: kind.into(),
192            actor,
193            created_at,
194            nonce,
195            room_id: None,
196            payload,
197            extra: BTreeMap::new(),
198        }
199    }
200
201    pub fn with_room_id(mut self, room_id: impl Into<String>) -> Self {
202        self.room_id = Some(room_id.into());
203        self
204    }
205}
206
207#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
208pub struct Envelope<P = Value> {
209    pub hash: String,
210    pub event: Event<P>,
211    pub signature: String,
212}
213
214pub fn canonical_event_bytes<P>(event: &Event<P>) -> Result<Vec<u8>>
215where
216    P: Serialize,
217{
218    serde_jcs::to_vec(event).map_err(|err| SdkError::CanonicalJson(err.to_string()))
219}
220
221pub fn event_hash<P>(event: &Event<P>) -> Result<String>
222where
223    P: Serialize,
224{
225    let digest = Sha3_256::digest(canonical_event_bytes(event)?);
226    Ok(URL_SAFE_NO_PAD.encode(digest))
227}
228
229pub fn sign_event<P>(signing_key: &SigningKey, event: &Event<P>) -> Result<String>
230where
231    P: Serialize,
232{
233    let bytes = canonical_event_bytes(event)?;
234    let signature = signing_key.sign(&bytes);
235    Ok(URL_SAFE_NO_PAD.encode(signature.to_bytes()))
236}
237
238pub fn verify_event_hash<P>(envelope: &Envelope<P>) -> Result<()>
239where
240    P: Serialize,
241{
242    let expected = event_hash(&envelope.event)?;
243    if expected == envelope.hash {
244        Ok(())
245    } else {
246        Err(SdkError::InvalidEventHash {
247            expected,
248            actual: envelope.hash.clone(),
249        })
250    }
251}
252
253pub fn verify_signature<P>(envelope: &Envelope<P>) -> Result<()>
254where
255    P: Serialize,
256{
257    let bytes = canonical_event_bytes(&envelope.event)?;
258    let signature_bytes = URL_SAFE_NO_PAD.decode(&envelope.signature)?;
259    let signature_bytes: [u8; 64] = signature_bytes
260        .try_into()
261        .map_err(|bytes: Vec<u8>| SdkError::InvalidSignatureLength(bytes.len()))?;
262    let signature = Signature::from_bytes(&signature_bytes);
263    envelope
264        .event
265        .actor
266        .verifying_key()?
267        .verify_strict(&bytes, &signature)?;
268    Ok(())
269}
270
271pub fn verify_envelope<P>(envelope: &Envelope<P>) -> Result<()>
272where
273    P: Serialize,
274{
275    verify_event_hash(envelope)?;
276    verify_signature(envelope)
277}
278
279pub fn verify_timestamp(created_at: i64, now_ms: i64, window_ms: i64) -> Result<()> {
280    if window_ms < 0 || created_at.abs_diff(now_ms) > window_ms as u64 {
281        Err(SdkError::TimestampOutOfWindow)
282    } else {
283        Ok(())
284    }
285}
286
287#[derive(Clone, Copy, Debug, Eq, PartialEq)]
288pub struct NonceRecord {
289    pub max_nonce: u64,
290    pub expires_at: i64,
291}
292
293pub trait NonceStore {
294    fn check_and_update(
295        &mut self,
296        actor: &AgentId,
297        nonce: u64,
298        now_ms: i64,
299        ttl_ms: i64,
300    ) -> Result<u64>;
301
302    fn max_nonce(&self, actor: &AgentId, now_ms: i64) -> Option<u64>;
303}
304
305#[derive(Clone, Debug, Default)]
306pub struct MemoryNonceStore {
307    records: HashMap<AgentId, NonceRecord>,
308}
309
310impl MemoryNonceStore {
311    pub fn new() -> Self {
312        Self::default()
313    }
314}
315
316impl NonceStore for MemoryNonceStore {
317    fn check_and_update(
318        &mut self,
319        actor: &AgentId,
320        nonce: u64,
321        now_ms: i64,
322        ttl_ms: i64,
323    ) -> Result<u64> {
324        validate_nonce(nonce)?;
325        if ttl_ms < 0 {
326            return Err(SdkError::InvalidNonce(
327                "nonce cache ttl must be non-negative".to_owned(),
328            ));
329        }
330
331        if let Some(record) = self.records.get(actor) {
332            if record.expires_at > now_ms && nonce <= record.max_nonce {
333                return Err(SdkError::NonceNotGreater {
334                    max_nonce: record.max_nonce,
335                });
336            }
337        }
338
339        let record = NonceRecord {
340            max_nonce: nonce,
341            expires_at: now_ms.saturating_add(ttl_ms),
342        };
343        self.records.insert(actor.clone(), record);
344        Ok(record.max_nonce)
345    }
346
347    fn max_nonce(&self, actor: &AgentId, now_ms: i64) -> Option<u64> {
348        self.records
349            .get(actor)
350            .filter(|record| record.expires_at > now_ms)
351            .map(|record| record.max_nonce)
352    }
353}
354
355#[derive(Clone, Debug, Eq, PartialEq)]
356pub struct ClientNonceManager {
357    next_nonce: u64,
358}
359
360impl Default for ClientNonceManager {
361    fn default() -> Self {
362        Self { next_nonce: 1 }
363    }
364}
365
366impl ClientNonceManager {
367    pub fn new() -> Self {
368        Self::default()
369    }
370
371    pub fn with_next(next_nonce: u64) -> Result<Self> {
372        validate_nonce(next_nonce)?;
373        Ok(Self { next_nonce })
374    }
375
376    pub fn peek(&self) -> u64 {
377        self.next_nonce
378    }
379
380    pub fn next_nonce(&mut self) -> Result<u64> {
381        let nonce = self.next_nonce;
382        self.next_nonce = self
383            .next_nonce
384            .checked_add(1)
385            .ok_or_else(|| SdkError::InvalidNonce("nonce counter overflow".to_owned()))?;
386        Ok(nonce)
387    }
388
389    pub fn observe_max_nonce(&mut self, max_nonce: u64) {
390        if max_nonce >= self.next_nonce {
391            self.next_nonce = max_nonce.saturating_add(1);
392        }
393    }
394
395    pub fn observe_max_nonce_header(&mut self, value: &str) -> Result<()> {
396        let max_nonce = value
397            .parse::<u64>()
398            .map_err(|_| SdkError::InvalidNonce("invalid max nonce header".to_owned()))?;
399        self.observe_max_nonce(max_nonce);
400        Ok(())
401    }
402}
403
404pub fn validate_nonce(nonce: u64) -> Result<()> {
405    if nonce == 0 || nonce > 0x1FFFFFFFFFFFFF {
406        // Number.MAX_SAFE_INTEGER in JavaScript, to prevent interoperability issues with JS clients using Number for nonces
407        Err(SdkError::InvalidNonce(
408            "nonce must be a positive integer less than or equal to 9007199254740991".to_owned(),
409        ))
410    } else {
411        Ok(())
412    }
413}
414
415#[derive(Clone, Debug)]
416pub struct LiveWriteOptions {
417    pub now_ms: i64,
418    pub window_ms: i64,
419    pub nonce_ttl_ms: i64,
420}
421
422impl Default for LiveWriteOptions {
423    fn default() -> Self {
424        Self {
425            now_ms: unix_ms(),
426            window_ms: DEFAULT_LIVE_WRITE_WINDOW_MS,
427            nonce_ttl_ms: DEFAULT_NONCE_TTL_MS,
428        }
429    }
430}
431
432pub fn verify_live_envelope<P, S>(
433    envelope: &Envelope<P>,
434    options: &LiveWriteOptions,
435    nonce_store: &mut S,
436) -> Result<u64>
437where
438    P: Serialize,
439    S: NonceStore,
440{
441    verify_envelope(envelope)?;
442    verify_timestamp(envelope.event.created_at, options.now_ms, options.window_ms)?;
443    nonce_store.check_and_update(
444        &envelope.event.actor,
445        envelope.event.nonce,
446        options.now_ms,
447        options.nonce_ttl_ms,
448    )
449}
450
451#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
452pub struct RequestJwtHeader {
453    pub alg: String,
454    pub typ: String,
455    pub kid: AgentId,
456}
457
458#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
459pub struct RequestJwtClaims {
460    pub iss: AgentId,
461    pub sub: AgentId,
462    pub aud: String,
463    pub iat: i64,
464    pub exp: i64,
465}
466
467#[derive(Clone, Debug, PartialEq, Eq)]
468pub struct RequestBinding {
469    pub audience: String,
470}
471
472impl RequestBinding {
473    pub fn new(audience: impl Into<String>) -> Self {
474        Self {
475            audience: audience.into(),
476        }
477    }
478}
479
480impl RequestJwtClaims {
481    pub fn new(agent_id: AgentId, binding: RequestBinding, issued_at: i64, ttl_secs: i64) -> Self {
482        Self {
483            iss: agent_id.clone(),
484            sub: agent_id,
485            aud: binding.audience,
486            iat: issued_at,
487            exp: issued_at + ttl_secs,
488        }
489    }
490}
491
492#[derive(Clone, Debug)]
493pub struct RequestAuthContext {
494    pub audience: String,
495    pub now_secs: i64,
496    pub max_ttl_secs: i64,
497}
498
499impl RequestAuthContext {
500    pub fn new(audience: impl Into<String>) -> Self {
501        let binding = RequestBinding::new(audience);
502        Self {
503            audience: binding.audience,
504            now_secs: unix_secs(),
505            max_ttl_secs: DEFAULT_REQUEST_JWT_TTL_SECS,
506        }
507    }
508}
509
510pub fn verify_request_jwt(token: &str, context: &RequestAuthContext) -> Result<RequestJwtClaims> {
511    let parts: Vec<&str> = token.split('.').collect();
512    if parts.len() != 3 {
513        return Err(SdkError::InvalidJwt(
514            "expected three compact JWS parts".into(),
515        ));
516    }
517
518    let header: RequestJwtHeader = serde_json::from_slice(&URL_SAFE_NO_PAD.decode(parts[0])?)?;
519    let claims: RequestJwtClaims = serde_json::from_slice(&URL_SAFE_NO_PAD.decode(parts[1])?)?;
520    let signature_bytes = URL_SAFE_NO_PAD.decode(parts[2])?;
521    let signature_bytes: [u8; 64] = signature_bytes
522        .try_into()
523        .map_err(|bytes: Vec<u8>| SdkError::InvalidSignatureLength(bytes.len()))?;
524    let signature = Signature::from_bytes(&signature_bytes);
525    let signing_input = format!("{}.{}", parts[0], parts[1]);
526
527    if header.alg != "EdDSA" {
528        return Err(SdkError::InvalidJwtClaim("alg"));
529    }
530    if header.typ != "JWT" {
531        return Err(SdkError::InvalidJwtClaim("typ"));
532    }
533    if header.kid != claims.iss || claims.iss != claims.sub {
534        return Err(SdkError::InvalidJwtClaim("kid/iss/sub"));
535    }
536
537    header
538        .kid
539        .verifying_key()?
540        .verify_strict(signing_input.as_bytes(), &signature)?;
541
542    if claims.aud != context.audience {
543        return Err(SdkError::InvalidJwtClaim("aud"));
544    }
545    if claims.iat > context.now_secs || claims.exp < context.now_secs {
546        return Err(SdkError::InvalidJwtClaim("iat/exp"));
547    }
548    if claims.exp - claims.iat > context.max_ttl_secs {
549        return Err(SdkError::InvalidJwtClaim("ttl"));
550    }
551
552    Ok(claims)
553}
554
555pub fn unix_ms() -> i64 {
556    SystemTime::now()
557        .duration_since(UNIX_EPOCH)
558        .map(|duration| duration.as_millis().min(i64::MAX as u128) as i64)
559        .unwrap_or(0)
560}
561
562pub fn unix_secs() -> i64 {
563    unix_ms() / 1000
564}
565
566#[cfg(test)]
567mod tests {
568    use serde_json::json;
569
570    use super::*;
571
572    #[test]
573    fn signs_and_verifies_event_envelope() {
574        let signer = AgentSigner::from_seed([7; 32]);
575        let event = Event::new(
576            "agent-profile/1.0",
577            "profile.update",
578            signer.agent_id(),
579            1_779_753_600_000,
580            1,
581            json!({"agent_id": signer.agent_id(), "name": "ResearchAgent"}),
582        );
583
584        let envelope = signer.sign_event(event).unwrap();
585
586        assert!(!envelope.hash.starts_with("evt_"));
587        assert_eq!(envelope.hash.len(), 43);
588        verify_envelope(&envelope).unwrap();
589    }
590
591    #[test]
592    fn rejects_tampered_event_payload() {
593        let signer = AgentSigner::from_seed([8; 32]);
594        let event = Event::new(
595            "agent-profile/1.0",
596            "profile.update",
597            signer.agent_id(),
598            1_779_753_600_000,
599            1,
600            json!({"name": "before"}),
601        );
602        let mut envelope = signer.sign_event(event).unwrap();
603        envelope.event.payload = json!({"name": "after"});
604
605        assert!(verify_envelope(&envelope).is_err());
606    }
607
608    #[test]
609    fn rejects_nonce_reuse() {
610        let signer = AgentSigner::from_seed([9; 32]);
611        let options = LiveWriteOptions {
612            now_ms: 1000,
613            window_ms: 1000,
614            nonce_ttl_ms: 1000,
615        };
616        let event = Event::new(
617            "agent-profile/1.0",
618            "profile.update",
619            signer.agent_id(),
620            1000,
621            1,
622            json!({"name": "ResearchAgent"}),
623        );
624        let envelope = signer.sign_event(event).unwrap();
625        let mut store = MemoryNonceStore::new();
626
627        let max_nonce = verify_live_envelope(&envelope, &options, &mut store).unwrap();
628        assert_eq!(max_nonce, 1);
629        assert!(matches!(
630            verify_live_envelope(&envelope, &options, &mut store),
631            Err(SdkError::NonceNotGreater { max_nonce: 1 })
632        ));
633    }
634
635    #[test]
636    fn client_nonce_manager_observes_server_max() {
637        let mut manager = ClientNonceManager::new();
638
639        assert_eq!(manager.next_nonce().unwrap(), 1);
640        manager.observe_max_nonce(5);
641
642        assert_eq!(manager.peek(), 6);
643        assert_eq!(manager.next_nonce().unwrap(), 6);
644    }
645
646    #[test]
647    fn signs_and_verifies_request_jwt() {
648        let signer = AgentSigner::from_seed([10; 32]);
649        let claims = RequestJwtClaims::new(
650            signer.agent_id(),
651            RequestBinding::new("https://api.example.com"),
652            100,
653            300,
654        );
655        let token = signer.sign_request_jwt(&claims).unwrap();
656        let context = RequestAuthContext {
657            audience: "https://api.example.com".to_owned(),
658            now_secs: 120,
659            max_ttl_secs: 300,
660        };
661        let verified = verify_request_jwt(&token, &context).unwrap();
662
663        assert_eq!(verified.iss, signer.agent_id());
664    }
665}
agent_protocols/identity.rs

agent_protocols/
identity.rs
