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(skip_serializing_if = "Option::is_none")]
160    resource: Option<String>,
161    sub: String,
162    #[serde(flatten)]
163    additional: AdditionalMembers,
164}
165
166pub fn sign_client_assertion(
167    claims: &ClientAssertionClaims,
168    options: SignClientAssertionOptions<'_>,
169) -> Result<String, CoreError> {
170    validate_client_assertion_claims_with_options(
171        claims,
172        ClientAssertionValidationOptions {
173            allow_insecure_loopback: options.allow_insecure_loopback,
174        },
175    )?;
176    validate_assertion_key_id(options.key_id, claims)?;
177    let claims = to_compact_claims(claims)?;
178    let header = Header::empty()
179        .with_key_id(options.key_id)
180        .with_token_type("JWT");
181    match &options.key.0 {
182        SigningKeyInner::Ed25519(key) => Ed25519.token(&header, &claims, key).map_err(jwt_error),
183        SigningKeyInner::Es256(key) => Es256.token(&header, &claims, key).map_err(jwt_error),
184    }
185}
186
187pub fn verify_client_assertion(
188    assertion: &str,
189    key: &ClientAssertionVerifyingKey,
190    options: &VerifyClientAssertionOptions,
191) -> Result<ClientAssertionClaims, CoreError> {
192    let token =
193        UntrustedToken::<BTreeMap<String, Value>>::try_from(assertion).map_err(jwt_error)?;
194    if token.header().token_type.as_deref() != Some("JWT") {
195        return Err(CoreError::Invalid(
196            "AEP client assertion typ must be JWT".to_owned(),
197        ));
198    }
199    if token.header().other_fields.contains_key("crit") {
200        return Err(CoreError::Invalid(
201            "AEP client assertion contains unsupported critical JOSE parameters".to_owned(),
202        ));
203    }
204    let key_id = token
205        .header()
206        .key_id
207        .as_deref()
208        .ok_or_else(|| CoreError::Invalid("AEP client assertion kid is required".to_owned()))?;
209    let algorithm = SigningAlgorithm::from(token.algorithm());
210    validate_allowed_algorithm(&algorithm, key, options)?;
211    let token = match &key.0 {
212        VerifyingKeyInner::Ed25519(key) => Ed25519
213            .validator::<AssertionCustomClaims>(key)
214            .validate(&token)
215            .map_err(jwt_error)?,
216        VerifyingKeyInner::Es256(key) => Es256
217            .validator::<AssertionCustomClaims>(key)
218            .validate(&token)
219            .map_err(jwt_error)?,
220    };
221    let claims = from_compact_claims(token.claims())?;
222    validate_client_assertion_claims_with_options(
223        &claims,
224        ClientAssertionValidationOptions {
225            allow_insecure_loopback: options.allow_insecure_loopback,
226        },
227    )?;
228    validate_assertion_key_id(key_id, &claims)?;
229    validate_expected_claims(&claims, options)?;
230    Ok(claims)
231}
232
233pub fn decode_jwt_unverified(assertion: &str) -> Result<DecodedJwt, CoreError> {
234    let token =
235        UntrustedToken::<BTreeMap<String, Value>>::try_from(assertion).map_err(jwt_error)?;
236    let claims = token
237        .deserialize_claims_unchecked::<BTreeMap<String, Value>>()
238        .map_err(jwt_error)?;
239    let payload = serde_json::to_value(claims)?
240        .as_object()
241        .cloned()
242        .ok_or_else(|| CoreError::Jwt("JWT claims must be a JSON object".to_owned()))?
243        .into_iter()
244        .collect();
245    Ok(DecodedJwt {
246        header: JwtHeader {
247            algorithm: SigningAlgorithm::from(token.algorithm()),
248            key_id: token.header().key_id.clone(),
249            token_type: token.header().token_type.clone(),
250            additional: token.header().other_fields.clone(),
251        },
252        payload,
253    })
254}
255
256fn to_compact_claims(
257    claims: &ClientAssertionClaims,
258) -> Result<Claims<AssertionCustomClaims>, CoreError> {
259    let mut compact = Claims::new(AssertionCustomClaims {
260        aud: claims.aud.clone(),
261        iss: claims.iss.clone(),
262        jti: claims.jti.clone(),
263        op: claims.op,
264        resource: claims.resource.clone(),
265        sub: claims.sub.clone(),
266        additional: claims.additional.clone(),
267    });
268    compact.expiration = Some(timestamp(claims.exp)?);
269    compact.issued_at = Some(timestamp(claims.iat)?);
270    Ok(compact)
271}
272
273fn from_compact_claims(
274    claims: &Claims<AssertionCustomClaims>,
275) -> Result<ClientAssertionClaims, CoreError> {
276    Ok(ClientAssertionClaims {
277        aud: claims.custom.aud.clone(),
278        exp: claims
279            .expiration
280            .ok_or_else(|| CoreError::Invalid("AEP client assertion exp is required".to_owned()))?
281            .timestamp(),
282        iat: claims
283            .issued_at
284            .ok_or_else(|| CoreError::Invalid("AEP client assertion iat is required".to_owned()))?
285            .timestamp(),
286        iss: claims.custom.iss.clone(),
287        jti: claims.custom.jti.clone(),
288        op: claims.custom.op,
289        resource: claims.custom.resource.clone(),
290        sub: claims.custom.sub.clone(),
291        additional: claims.custom.additional.clone(),
292    })
293}
294
295fn timestamp(value: i64) -> Result<DateTime<Utc>, CoreError> {
296    DateTime::from_timestamp(value, 0)
297        .ok_or_else(|| CoreError::Invalid("AEP client assertion timestamp is invalid".to_owned()))
298}
299
300fn validate_allowed_algorithm(
301    algorithm: &SigningAlgorithm,
302    key: &ClientAssertionVerifyingKey,
303    options: &VerifyClientAssertionOptions,
304) -> Result<(), CoreError> {
305    if algorithm != &key.algorithm() {
306        return Err(CoreError::Invalid(
307            "AEP client assertion signing key does not match alg".to_owned(),
308        ));
309    }
310    let allowed = if options.algorithms.is_empty() {
311        matches!(algorithm, SigningAlgorithm::EdDsa | SigningAlgorithm::Es256)
312    } else {
313        options.algorithms.contains(algorithm)
314    };
315    if !allowed {
316        return Err(CoreError::Invalid(
317            "AEP client assertion signing algorithm is not allowed".to_owned(),
318        ));
319    }
320    Ok(())
321}
322
323fn validate_expected_claims(
324    claims: &ClientAssertionClaims,
325    options: &VerifyClientAssertionOptions,
326) -> Result<(), CoreError> {
327    for (actual, expected, message) in [
328        (&claims.aud, options.audience.as_ref(), "audience"),
329        (&claims.iss, options.issuer.as_ref(), "issuer"),
330        (&claims.sub, options.subject.as_ref(), "subject"),
331    ] {
332        if expected.is_some_and(|expected| actual != expected) {
333            return Err(CoreError::Invalid(format!(
334                "AEP client assertion {message} does not match"
335            )));
336        }
337    }
338    if options
339        .operation
340        .is_some_and(|operation| claims.op != operation)
341    {
342        return Err(CoreError::Invalid(
343            "AEP client assertion operation does not match".to_owned(),
344        ));
345    }
346    if options
347        .resource
348        .as_ref()
349        .is_some_and(|resource| claims.resource.as_ref() != Some(resource))
350    {
351        return Err(CoreError::Invalid(
352            "AEP client assertion resource does not match".to_owned(),
353        ));
354    }
355    let now = options.current_time.unwrap_or_else(current_timestamp);
356    let tolerance = options
357        .clock_tolerance_seconds
358        .unwrap_or(RECOMMENDED_CLOCK_SKEW.as_secs());
359    let tolerance = i64::try_from(tolerance)
360        .map_err(|_| CoreError::Invalid("clock tolerance is too large".to_owned()))?;
361    if claims.iat > now.saturating_add(tolerance) || claims.exp < now.saturating_sub(tolerance) {
362        return Err(CoreError::Invalid(
363            "AEP client assertion is outside its validity window".to_owned(),
364        ));
365    }
366    Ok(())
367}
368
369fn current_timestamp() -> i64 {
370    SystemTime::UNIX_EPOCH
371        .elapsed()
372        .ok()
373        .and_then(|duration| i64::try_from(duration.as_secs()).ok())
374        .unwrap_or(i64::MAX)
375}
376
377fn validate_assertion_key_id(
378    key_id: &str,
379    claims: &ClientAssertionClaims,
380) -> Result<(), CoreError> {
381    if key_id.is_empty() {
382        return Err(CoreError::Invalid(
383            "AEP client assertion kid is required".to_owned(),
384        ));
385    }
386    let key_did = key_id.split_once('#').map_or(key_id, |part| part.0);
387    if key_did != claims.iss || key_did != claims.sub {
388        return Err(CoreError::Invalid(
389            "AEP client assertion kid must identify the Agent DID".to_owned(),
390        ));
391    }
392    Ok(())
393}
394
395fn jwt_error(error: impl std::fmt::Display) -> CoreError {
396    CoreError::Jwt(error.to_string())
397}
398
399#[cfg(test)]
400mod tests {
401    use super::*;
402
403    fn claims(operation: AssertionOperation, jti: &str) -> ClientAssertionClaims {
404        ClientAssertionClaims {
405            aud: "did:web:service.example".to_owned(),
406            exp: 1_748_428_860,
407            iat: 1_748_428_800,
408            iss: "did:web:agent.example".to_owned(),
409            jti: jti.to_owned(),
410            op: operation,
411            resource: None,
412            sub: "did:web:agent.example".to_owned(),
413            additional: Default::default(),
414        }
415    }
416
417    #[test]
418    fn rejects_an_invalid_unverified_token() {
419        assert!(decode_jwt_unverified("not-a-jwt").is_err());
420    }
421
422    #[test]
423    fn imports_supported_key_encodings() {
424        let ed25519_pem = r#"-----BEGIN PRIVATE KEY-----
425MC4CAQAwBQYDK2VwBCIEIAABAgMEBQYHCAkKCwwNDg8QERITFBUWFxgZGhscHR4f
426-----END PRIVATE KEY-----"#;
427        let ed25519 = ClientAssertionSigningKey::ed25519_from_pkcs8_pem(ed25519_pem)
428            .expect("Ed25519 signing key");
429        assert_eq!(ed25519.algorithm(), SigningAlgorithm::EdDsa);
430        let ed25519_verifying = ed25519.verifying_key();
431        let VerifyingKeyInner::Ed25519(ed25519_bytes) = &ed25519_verifying.0 else {
432            panic!("expected Ed25519 key");
433        };
434        ClientAssertionVerifyingKey::ed25519_from_bytes(ed25519_bytes.as_bytes())
435            .expect("Ed25519 verifying key");
436
437        let es256_pem = r#"-----BEGIN PRIVATE KEY-----
438MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQg8AF0ffHVA9WPTRaT
4393ITlDvQ3VHQLX/xwgR6esnFXphKhRANCAAT6SMwbSLTZ9wPLUy5ilCP1HgjJ/xbS
440HMacx3g7kP+jGzNEdPXNVpmcqyQxe3Ffb2VWVcxnWIu7fR1d/Il52p9E
441-----END PRIVATE KEY-----"#;
442        let es256 =
443            ClientAssertionSigningKey::es256_from_pkcs8_pem(es256_pem).expect("ES256 signing key");
444        assert_eq!(es256.algorithm(), SigningAlgorithm::Es256);
445        let es256_verifying = es256.verifying_key();
446        let VerifyingKeyInner::Es256(es256_key) = &es256_verifying.0 else {
447            panic!("expected ES256 key");
448        };
449        let point = es256_key.to_encoded_point(false);
450        ClientAssertionVerifyingKey::es256_from_sec1_bytes(point.as_bytes())
451            .expect("ES256 verifying key");
452    }
453
454    #[test]
455    fn signs_and_verifies_an_eddsa_assertion() {
456        let signing_key = ClientAssertionSigningKey::ed25519_from_seed([7; 32]);
457        let verifying_key = signing_key.verifying_key();
458        let claims = claims(AssertionOperation::Status, "jti-1");
459        let assertion = sign_client_assertion(
460            &claims,
461            SignClientAssertionOptions {
462                allow_insecure_loopback: false,
463                key: &signing_key,
464                key_id: "did:web:agent.example#key-1",
465            },
466        )
467        .expect("signed assertion");
468        let verified = verify_client_assertion(
469            &assertion,
470            &verifying_key,
471            &VerifyClientAssertionOptions {
472                audience: Some(claims.aud.clone()),
473                current_time: Some(1_748_428_830),
474                operation: Some(AssertionOperation::Status),
475                ..VerifyClientAssertionOptions::default()
476            },
477        )
478        .expect("verified assertion");
479        assert_eq!(verified.jti, "jti-1");
480        assert!(
481            verify_client_assertion(
482                &assertion,
483                &verifying_key,
484                &VerifyClientAssertionOptions {
485                    audience: Some("did:web:other.example".to_owned()),
486                    current_time: Some(1_748_428_830),
487                    ..VerifyClientAssertionOptions::default()
488                },
489            )
490            .is_err()
491        );
492        let decoded = decode_jwt_unverified(&assertion).expect("decoded assertion");
493        assert_eq!(
494            decoded.header.key_id.as_deref(),
495            Some("did:web:agent.example#key-1")
496        );
497        assert_eq!(decoded.payload.get("jti"), Some(&Value::from("jti-1")));
498    }
499
500    #[test]
501    fn signs_and_verifies_an_es256_assertion() {
502        let signing_key =
503            ClientAssertionSigningKey::es256_from_bytes(&[3; 32]).expect("signing key");
504        let verifying_key = signing_key.verifying_key();
505        let claims = claims(AssertionOperation::Grant, "jti-2");
506        let assertion = sign_client_assertion(
507            &claims,
508            SignClientAssertionOptions {
509                allow_insecure_loopback: false,
510                key: &signing_key,
511                key_id: "did:web:agent.example#key-2",
512            },
513        )
514        .expect("signed assertion");
515        let verified = verify_client_assertion(
516            &assertion,
517            &verifying_key,
518            &VerifyClientAssertionOptions {
519                algorithms: vec![SigningAlgorithm::Es256],
520                current_time: Some(1_748_428_830),
521                ..VerifyClientAssertionOptions::default()
522            },
523        )
524        .expect("verified assertion");
525        assert_eq!(verified.jti, "jti-2");
526    }
527
528    #[test]
529    fn rejects_a_mismatched_algorithm_and_key() {
530        let ed25519 = ClientAssertionSigningKey::ed25519_from_seed([7; 32]);
531        let es256 = ClientAssertionSigningKey::es256_from_bytes(&[3; 32])
532            .expect("ES256 signing key")
533            .verifying_key();
534        let assertion = sign_client_assertion(
535            &claims(AssertionOperation::Status, "jti-3"),
536            SignClientAssertionOptions {
537                allow_insecure_loopback: false,
538                key: &ed25519,
539                key_id: "did:web:agent.example#key-1",
540            },
541        )
542        .expect("signed assertion");
543        assert!(
544            verify_client_assertion(
545                &assertion,
546                &es256,
547                &VerifyClientAssertionOptions {
548                    current_time: Some(1_748_428_830),
549                    ..VerifyClientAssertionOptions::default()
550                },
551            )
552            .is_err()
553        );
554    }
555
556    #[test]
557    fn rejects_invalid_signing_inputs() {
558        let key = ClientAssertionSigningKey::ed25519_from_seed([7; 32]);
559        let claims = claims(AssertionOperation::Status, "jti");
560        assert!(
561            sign_client_assertion(
562                &claims,
563                SignClientAssertionOptions {
564                    allow_insecure_loopback: false,
565                    key: &key,
566                    key_id: "did:web:other.example#key-1",
567                },
568            )
569            .is_err()
570        );
571    }
572}