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}