Skip to main content

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