Skip to main content

aep_core/
jwt.rs

1use std::{collections::BTreeMap, time::SystemTime};
2
3use chrono::{DateTime, Utc};
4use ed25519_dalek::pkcs8::DecodePrivateKey as _;
5use jwt_compact::{
6    AlgorithmExt, Claims, Header, UntrustedToken,
7    alg::{Ed25519, Es256},
8};
9use serde::{Deserialize, Serialize};
10use serde_json::Value;
11
12use crate::{
13    AdditionalMembers, AssertionOperation, ClientAssertionClaims, ClientAssertionValidationOptions,
14    CoreError, RECOMMENDED_CLOCK_SKEW, SigningAlgorithm,
15    validate_client_assertion_claims_with_options,
16};
17
18#[derive(Clone)]
19pub struct ClientAssertionSigningKey(SigningKeyInner);
20
21#[derive(Clone)]
22enum SigningKeyInner {
23    Ed25519(ed25519_dalek::SigningKey),
24    Es256(p256::ecdsa::SigningKey),
25}
26
27impl ClientAssertionSigningKey {
28    pub fn ed25519_from_seed(seed: [u8; 32]) -> Self {
29        Self(SigningKeyInner::Ed25519(
30            ed25519_dalek::SigningKey::from_bytes(&seed),
31        ))
32    }
33
34    pub fn ed25519_from_pkcs8_pem(pem: &str) -> Result<Self, CoreError> {
35        ed25519_dalek::SigningKey::from_pkcs8_pem(pem)
36            .map(|key| Self(SigningKeyInner::Ed25519(key)))
37            .map_err(jwt_error)
38    }
39
40    pub fn es256_from_bytes(bytes: &[u8]) -> Result<Self, CoreError> {
41        p256::ecdsa::SigningKey::from_slice(bytes)
42            .map(|key| Self(SigningKeyInner::Es256(key)))
43            .map_err(jwt_error)
44    }
45
46    pub fn es256_from_pkcs8_pem(pem: &str) -> Result<Self, CoreError> {
47        p256::ecdsa::SigningKey::from_pkcs8_pem(pem)
48            .map(|key| Self(SigningKeyInner::Es256(key)))
49            .map_err(jwt_error)
50    }
51
52    pub fn algorithm(&self) -> SigningAlgorithm {
53        match self.0 {
54            SigningKeyInner::Ed25519(_) => SigningAlgorithm::EdDsa,
55            SigningKeyInner::Es256(_) => SigningAlgorithm::Es256,
56        }
57    }
58
59    pub fn verifying_key(&self) -> ClientAssertionVerifyingKey {
60        match &self.0 {
61            SigningKeyInner::Ed25519(key) => {
62                ClientAssertionVerifyingKey(VerifyingKeyInner::Ed25519(key.verifying_key()))
63            }
64            SigningKeyInner::Es256(key) => {
65                ClientAssertionVerifyingKey(VerifyingKeyInner::Es256(*key.verifying_key()))
66            }
67        }
68    }
69}
70
71#[derive(Clone, Debug, Eq, PartialEq)]
72pub struct ClientAssertionVerifyingKey(VerifyingKeyInner);
73
74#[derive(Clone, Debug, Eq, PartialEq)]
75enum VerifyingKeyInner {
76    Ed25519(ed25519_dalek::VerifyingKey),
77    Es256(p256::ecdsa::VerifyingKey),
78}
79
80impl ClientAssertionVerifyingKey {
81    pub fn ed25519_from_bytes(bytes: &[u8]) -> Result<Self, CoreError> {
82        let bytes = <&[u8; 32]>::try_from(bytes)
83            .map_err(|_| CoreError::Jwt("Ed25519 public key must contain 32 bytes".to_owned()))?;
84        ed25519_dalek::VerifyingKey::from_bytes(bytes)
85            .map(|key| Self(VerifyingKeyInner::Ed25519(key)))
86            .map_err(jwt_error)
87    }
88
89    pub fn es256_from_sec1_bytes(bytes: &[u8]) -> Result<Self, CoreError> {
90        p256::ecdsa::VerifyingKey::from_sec1_bytes(bytes)
91            .map(|key| Self(VerifyingKeyInner::Es256(key)))
92            .map_err(jwt_error)
93    }
94
95    pub fn algorithm(&self) -> SigningAlgorithm {
96        match self.0 {
97            VerifyingKeyInner::Ed25519(_) => SigningAlgorithm::EdDsa,
98            VerifyingKeyInner::Es256(_) => SigningAlgorithm::Es256,
99        }
100    }
101
102    pub(crate) fn from_jwk(
103        jwk: &jwt_compact::jwk::JsonWebKey<'_>,
104        algorithm: &SigningAlgorithm,
105    ) -> Result<Self, CoreError> {
106        match algorithm {
107            SigningAlgorithm::EdDsa => ed25519_dalek::VerifyingKey::try_from(jwk)
108                .map(|key| Self(VerifyingKeyInner::Ed25519(key)))
109                .map_err(jwt_error),
110            SigningAlgorithm::Es256 => p256::ecdsa::VerifyingKey::try_from(jwk)
111                .map(|key| Self(VerifyingKeyInner::Es256(key)))
112                .map_err(jwt_error),
113            SigningAlgorithm::Other(value) => Err(CoreError::Invalid(format!(
114                "unsupported AEP signing algorithm {value:?}"
115            ))),
116        }
117    }
118}
119
120pub struct SignClientAssertionOptions<'a> {
121    pub allow_insecure_loopback: bool,
122    pub key: &'a ClientAssertionSigningKey,
123    pub key_id: &'a str,
124}
125
126#[derive(Clone, Debug, Default, Eq, PartialEq)]
127pub struct VerifyClientAssertionOptions {
128    pub algorithms: Vec<SigningAlgorithm>,
129    pub allow_insecure_loopback: bool,
130    pub audience: Option<String>,
131    pub clock_tolerance_seconds: Option<u64>,
132    pub current_time: Option<i64>,
133    pub issuer: Option<String>,
134    pub operation: Option<AssertionOperation>,
135    pub resource: Option<String>,
136    pub subject: Option<String>,
137}
138
139#[derive(Clone, Debug, PartialEq)]
140pub struct JwtHeader {
141    pub algorithm: SigningAlgorithm,
142    pub key_id: Option<String>,
143    pub token_type: Option<String>,
144    pub additional: BTreeMap<String, Value>,
145}
146
147#[derive(Clone, Debug, PartialEq)]
148pub struct DecodedJwt {
149    pub header: JwtHeader,
150    pub payload: BTreeMap<String, Value>,
151}
152
153#[derive(Clone, Debug, Deserialize, Serialize)]
154struct AssertionCustomClaims {
155    aud: String,
156    iss: String,
157    jti: String,
158    op: AssertionOperation,
159    #[serde(
160        default,
161        deserialize_with = "crate::validation::deserialize_optional_non_null",
162        skip_serializing_if = "Option::is_none"
163    )]
164    resource: Option<String>,
165    sub: String,
166    #[serde(flatten)]
167    additional: AdditionalMembers,
168}
169
170pub fn sign_client_assertion(
171    claims: &ClientAssertionClaims,
172    options: SignClientAssertionOptions<'_>,
173) -> Result<String, CoreError> {
174    validate_client_assertion_claims_with_options(
175        claims,
176        ClientAssertionValidationOptions {
177            allow_insecure_loopback: options.allow_insecure_loopback,
178        },
179    )?;
180    validate_assertion_key_id(options.key_id, claims)?;
181    let claims = to_compact_claims(claims)?;
182    let header = Header::empty()
183        .with_key_id(options.key_id)
184        .with_token_type("JWT");
185    match &options.key.0 {
186        SigningKeyInner::Ed25519(key) => Ed25519.token(&header, &claims, key).map_err(jwt_error),
187        SigningKeyInner::Es256(key) => Es256.token(&header, &claims, key).map_err(jwt_error),
188    }
189}
190
191pub fn verify_client_assertion(
192    assertion: &str,
193    key: &ClientAssertionVerifyingKey,
194    options: &VerifyClientAssertionOptions,
195) -> Result<ClientAssertionClaims, CoreError> {
196    let token =
197        UntrustedToken::<BTreeMap<String, Value>>::try_from(assertion).map_err(jwt_error)?;
198    if token.header().token_type.as_deref() != Some("JWT") {
199        return Err(CoreError::Invalid(
200            "AEP client assertion typ must be JWT".to_owned(),
201        ));
202    }
203    if token.header().other_fields.contains_key("crit") {
204        return Err(CoreError::Invalid(
205            "AEP client assertion contains unsupported critical JOSE parameters".to_owned(),
206        ));
207    }
208    let key_id = token
209        .header()
210        .key_id
211        .as_deref()
212        .ok_or_else(|| CoreError::Invalid("AEP client assertion kid is required".to_owned()))?;
213    let algorithm = SigningAlgorithm::from(token.algorithm());
214    validate_allowed_algorithm(&algorithm, key, options)?;
215    let token = match &key.0 {
216        VerifyingKeyInner::Ed25519(key) => Ed25519
217            .validator::<AssertionCustomClaims>(key)
218            .validate(&token)
219            .map_err(jwt_error)?,
220        VerifyingKeyInner::Es256(key) => Es256
221            .validator::<AssertionCustomClaims>(key)
222            .validate(&token)
223            .map_err(jwt_error)?,
224    };
225    let claims = from_compact_claims(token.claims())?;
226    validate_client_assertion_claims_with_options(
227        &claims,
228        ClientAssertionValidationOptions {
229            allow_insecure_loopback: options.allow_insecure_loopback,
230        },
231    )?;
232    validate_assertion_key_id(key_id, &claims)?;
233    validate_expected_claims(
234        &claims,
235        token.claims().not_before.map(|value| value.timestamp()),
236        options,
237    )?;
238    Ok(claims)
239}
240
241pub fn decode_jwt_unverified(assertion: &str) -> Result<DecodedJwt, CoreError> {
242    let token =
243        UntrustedToken::<BTreeMap<String, Value>>::try_from(assertion).map_err(jwt_error)?;
244    let claims = token
245        .deserialize_claims_unchecked::<BTreeMap<String, Value>>()
246        .map_err(jwt_error)?;
247    let payload = serde_json::to_value(claims)?
248        .as_object()
249        .cloned()
250        .ok_or_else(|| CoreError::Jwt("JWT claims must be a JSON object".to_owned()))?
251        .into_iter()
252        .collect();
253    Ok(DecodedJwt {
254        header: JwtHeader {
255            algorithm: SigningAlgorithm::from(token.algorithm()),
256            key_id: token.header().key_id.clone(),
257            token_type: token.header().token_type.clone(),
258            additional: token.header().other_fields.clone(),
259        },
260        payload,
261    })
262}
263
264fn to_compact_claims(
265    claims: &ClientAssertionClaims,
266) -> Result<Claims<AssertionCustomClaims>, CoreError> {
267    let mut compact = Claims::new(AssertionCustomClaims {
268        aud: claims.aud.clone(),
269        iss: claims.iss.clone(),
270        jti: claims.jti.clone(),
271        op: claims.op,
272        resource: claims.resource.clone(),
273        sub: claims.sub.clone(),
274        additional: claims.additional.clone(),
275    });
276    compact.expiration = Some(timestamp(claims.exp)?);
277    compact.issued_at = Some(timestamp(claims.iat)?);
278    Ok(compact)
279}
280
281fn from_compact_claims(
282    claims: &Claims<AssertionCustomClaims>,
283) -> Result<ClientAssertionClaims, CoreError> {
284    Ok(ClientAssertionClaims {
285        aud: claims.custom.aud.clone(),
286        exp: claims
287            .expiration
288            .ok_or_else(|| CoreError::Invalid("AEP client assertion exp is required".to_owned()))?
289            .timestamp(),
290        iat: claims
291            .issued_at
292            .ok_or_else(|| CoreError::Invalid("AEP client assertion iat is required".to_owned()))?
293            .timestamp(),
294        iss: claims.custom.iss.clone(),
295        jti: claims.custom.jti.clone(),
296        op: claims.custom.op,
297        resource: claims.custom.resource.clone(),
298        sub: claims.custom.sub.clone(),
299        additional: claims.custom.additional.clone(),
300    })
301}
302
303fn timestamp(value: i64) -> Result<DateTime<Utc>, CoreError> {
304    DateTime::from_timestamp(value, 0)
305        .ok_or_else(|| CoreError::Invalid("AEP client assertion timestamp is invalid".to_owned()))
306}
307
308fn validate_allowed_algorithm(
309    algorithm: &SigningAlgorithm,
310    key: &ClientAssertionVerifyingKey,
311    options: &VerifyClientAssertionOptions,
312) -> Result<(), CoreError> {
313    if algorithm != &key.algorithm() {
314        return Err(CoreError::Invalid(
315            "AEP client assertion signing key does not match alg".to_owned(),
316        ));
317    }
318    let allowed = if options.algorithms.is_empty() {
319        matches!(algorithm, SigningAlgorithm::EdDsa | SigningAlgorithm::Es256)
320    } else {
321        options.algorithms.contains(algorithm)
322    };
323    if !allowed {
324        return Err(CoreError::Invalid(
325            "AEP client assertion signing algorithm is not allowed".to_owned(),
326        ));
327    }
328    Ok(())
329}
330
331fn validate_expected_claims(
332    claims: &ClientAssertionClaims,
333    not_before: Option<i64>,
334    options: &VerifyClientAssertionOptions,
335) -> Result<(), CoreError> {
336    for (actual, expected, message) in [
337        (&claims.aud, options.audience.as_ref(), "audience"),
338        (&claims.iss, options.issuer.as_ref(), "issuer"),
339        (&claims.sub, options.subject.as_ref(), "subject"),
340    ] {
341        if expected.is_some_and(|expected| actual != expected) {
342            return Err(CoreError::Invalid(format!(
343                "AEP client assertion {message} does not match"
344            )));
345        }
346    }
347    if options
348        .operation
349        .is_some_and(|operation| claims.op != operation)
350    {
351        return Err(CoreError::Invalid(
352            "AEP client assertion operation does not match".to_owned(),
353        ));
354    }
355    if options
356        .resource
357        .as_ref()
358        .is_some_and(|resource| claims.resource.as_ref() != Some(resource))
359    {
360        return Err(CoreError::Invalid(
361            "AEP client assertion resource does not match".to_owned(),
362        ));
363    }
364    let now = options.current_time.unwrap_or_else(current_timestamp);
365    let tolerance = options
366        .clock_tolerance_seconds
367        .unwrap_or(RECOMMENDED_CLOCK_SKEW.as_secs());
368    let tolerance = i64::try_from(tolerance)
369        .map_err(|_| CoreError::Invalid("clock tolerance is too large".to_owned()))?;
370    if claims.iat > now.saturating_add(tolerance)
371        || claims.exp <= now.saturating_sub(tolerance)
372        || not_before.is_some_and(|value| value > now.saturating_add(tolerance))
373    {
374        return Err(CoreError::Invalid(
375            "AEP client assertion is outside its validity window".to_owned(),
376        ));
377    }
378    Ok(())
379}
380
381fn current_timestamp() -> i64 {
382    SystemTime::UNIX_EPOCH
383        .elapsed()
384        .ok()
385        .and_then(|duration| i64::try_from(duration.as_secs()).ok())
386        .unwrap_or(i64::MAX)
387}
388
389fn validate_assertion_key_id(
390    key_id: &str,
391    claims: &ClientAssertionClaims,
392) -> Result<(), CoreError> {
393    if key_id.is_empty() {
394        return Err(CoreError::Invalid(
395            "AEP client assertion kid is required".to_owned(),
396        ));
397    }
398    let key_did = key_id.split_once('#').map_or(key_id, |part| part.0);
399    if key_did != claims.iss || key_did != claims.sub {
400        return Err(CoreError::Invalid(
401            "AEP client assertion kid must identify the Agent DID".to_owned(),
402        ));
403    }
404    Ok(())
405}
406
407fn jwt_error(error: impl std::fmt::Display) -> CoreError {
408    CoreError::Jwt(error.to_string())
409}
410
411#[cfg(test)]
412mod tests {
413    use super::*;
414
415    fn claims(operation: AssertionOperation, jti: &str) -> ClientAssertionClaims {
416        ClientAssertionClaims {
417            aud: "did:web:service.example".to_owned(),
418            exp: 1_748_428_860,
419            iat: 1_748_428_800,
420            iss: "did:web:agent.example".to_owned(),
421            jti: jti.to_owned(),
422            op: operation,
423            resource: None,
424            sub: "did:web:agent.example".to_owned(),
425            additional: Default::default(),
426        }
427    }
428
429    #[test]
430    fn rejects_an_invalid_unverified_token() {
431        assert!(decode_jwt_unverified("not-a-jwt").is_err());
432    }
433
434    #[test]
435    fn imports_supported_key_encodings() {
436        let ed25519_pem = r#"-----BEGIN PRIVATE KEY-----
437MC4CAQAwBQYDK2VwBCIEIAABAgMEBQYHCAkKCwwNDg8QERITFBUWFxgZGhscHR4f
438-----END PRIVATE KEY-----"#;
439        let ed25519 = ClientAssertionSigningKey::ed25519_from_pkcs8_pem(ed25519_pem)
440            .expect("Ed25519 signing key");
441        assert_eq!(ed25519.algorithm(), SigningAlgorithm::EdDsa);
442        let ed25519_verifying = ed25519.verifying_key();
443        let VerifyingKeyInner::Ed25519(ed25519_bytes) = &ed25519_verifying.0 else {
444            panic!("expected Ed25519 key");
445        };
446        ClientAssertionVerifyingKey::ed25519_from_bytes(ed25519_bytes.as_bytes())
447            .expect("Ed25519 verifying key");
448
449        let es256_pem = r#"-----BEGIN PRIVATE KEY-----
450MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQg8AF0ffHVA9WPTRaT
4513ITlDvQ3VHQLX/xwgR6esnFXphKhRANCAAT6SMwbSLTZ9wPLUy5ilCP1HgjJ/xbS
452HMacx3g7kP+jGzNEdPXNVpmcqyQxe3Ffb2VWVcxnWIu7fR1d/Il52p9E
453-----END PRIVATE KEY-----"#;
454        let es256 =
455            ClientAssertionSigningKey::es256_from_pkcs8_pem(es256_pem).expect("ES256 signing key");
456        assert_eq!(es256.algorithm(), SigningAlgorithm::Es256);
457        let es256_verifying = es256.verifying_key();
458        let VerifyingKeyInner::Es256(es256_key) = &es256_verifying.0 else {
459            panic!("expected ES256 key");
460        };
461        let point = es256_key.to_encoded_point(false);
462        ClientAssertionVerifyingKey::es256_from_sec1_bytes(point.as_bytes())
463            .expect("ES256 verifying key");
464    }
465
466    #[test]
467    fn signs_and_verifies_an_eddsa_assertion() {
468        let signing_key = ClientAssertionSigningKey::ed25519_from_seed([7; 32]);
469        let verifying_key = signing_key.verifying_key();
470        let claims = claims(AssertionOperation::Status, "jti-1");
471        let assertion = sign_client_assertion(
472            &claims,
473            SignClientAssertionOptions {
474                allow_insecure_loopback: false,
475                key: &signing_key,
476                key_id: "did:web:agent.example#key-1",
477            },
478        )
479        .expect("signed assertion");
480        let verified = verify_client_assertion(
481            &assertion,
482            &verifying_key,
483            &VerifyClientAssertionOptions {
484                audience: Some(claims.aud.clone()),
485                current_time: Some(1_748_428_830),
486                operation: Some(AssertionOperation::Status),
487                ..VerifyClientAssertionOptions::default()
488            },
489        )
490        .expect("verified assertion");
491        assert_eq!(verified.jti, "jti-1");
492        assert!(
493            verify_client_assertion(
494                &assertion,
495                &verifying_key,
496                &VerifyClientAssertionOptions {
497                    audience: Some("did:web:other.example".to_owned()),
498                    current_time: Some(1_748_428_830),
499                    ..VerifyClientAssertionOptions::default()
500                },
501            )
502            .is_err()
503        );
504        let decoded = decode_jwt_unverified(&assertion).expect("decoded assertion");
505        assert_eq!(
506            decoded.header.key_id.as_deref(),
507            Some("did:web:agent.example#key-1")
508        );
509        assert_eq!(decoded.payload.get("jti"), Some(&Value::from("jti-1")));
510    }
511
512    #[test]
513    fn rejects_the_expiration_boundary() {
514        let signing_key = ClientAssertionSigningKey::ed25519_from_seed([7; 32]);
515        let assertion = sign_client_assertion(
516            &claims(AssertionOperation::Status, "expiration-boundary"),
517            SignClientAssertionOptions {
518                allow_insecure_loopback: false,
519                key: &signing_key,
520                key_id: "did:web:agent.example#key-1",
521            },
522        )
523        .expect("signed assertion");
524
525        assert!(
526            verify_client_assertion(
527                &assertion,
528                &signing_key.verifying_key(),
529                &VerifyClientAssertionOptions {
530                    clock_tolerance_seconds: Some(30),
531                    current_time: Some(1_748_428_890),
532                    ..VerifyClientAssertionOptions::default()
533                },
534            )
535            .is_err()
536        );
537    }
538
539    #[test]
540    fn signs_and_verifies_an_es256_assertion() {
541        let signing_key =
542            ClientAssertionSigningKey::es256_from_bytes(&[3; 32]).expect("signing key");
543        let verifying_key = signing_key.verifying_key();
544        let claims = claims(AssertionOperation::Grant, "jti-2");
545        let assertion = sign_client_assertion(
546            &claims,
547            SignClientAssertionOptions {
548                allow_insecure_loopback: false,
549                key: &signing_key,
550                key_id: "did:web:agent.example#key-2",
551            },
552        )
553        .expect("signed assertion");
554        let verified = verify_client_assertion(
555            &assertion,
556            &verifying_key,
557            &VerifyClientAssertionOptions {
558                algorithms: vec![SigningAlgorithm::Es256],
559                current_time: Some(1_748_428_830),
560                ..VerifyClientAssertionOptions::default()
561            },
562        )
563        .expect("verified assertion");
564        assert_eq!(verified.jti, "jti-2");
565    }
566
567    #[test]
568    fn rejects_a_mismatched_algorithm_and_key() {
569        let ed25519 = ClientAssertionSigningKey::ed25519_from_seed([7; 32]);
570        let es256 = ClientAssertionSigningKey::es256_from_bytes(&[3; 32])
571            .expect("ES256 signing key")
572            .verifying_key();
573        let assertion = sign_client_assertion(
574            &claims(AssertionOperation::Status, "jti-3"),
575            SignClientAssertionOptions {
576                allow_insecure_loopback: false,
577                key: &ed25519,
578                key_id: "did:web:agent.example#key-1",
579            },
580        )
581        .expect("signed assertion");
582        assert!(
583            verify_client_assertion(
584                &assertion,
585                &es256,
586                &VerifyClientAssertionOptions {
587                    current_time: Some(1_748_428_830),
588                    ..VerifyClientAssertionOptions::default()
589                },
590            )
591            .is_err()
592        );
593    }
594
595    #[test]
596    fn rejects_invalid_signing_inputs() {
597        let key = ClientAssertionSigningKey::ed25519_from_seed([7; 32]);
598        let claims = claims(AssertionOperation::Status, "jti");
599        assert!(
600            sign_client_assertion(
601                &claims,
602                SignClientAssertionOptions {
603                    allow_insecure_loopback: false,
604                    key: &key,
605                    key_id: "did:web:other.example#key-1",
606                },
607            )
608            .is_err()
609        );
610    }
611}