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

agent_protocols/
identity.rs
