feat: Implement extraction of KIDs from unvalidated tokens

2 files changed, 32 insertions, 2 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 1633dc255fb3..ba856c176615 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -154,8 +154,29 @@ impl From<ErrorStack> for ValidationError {
 ///
 /// This is only safe if the key set containing the currently allowed
 /// key IDs is fetched from a trusted source.
-pub fn token_kid(jwt: JWT) -> Option<String> {
-    unimplemented!()
+pub fn token_kid(jwt: &JWT) -> JWTResult<Option<String>> {
+    // Fetch the header component of the JWT by splitting it out and
+    // dismissing the rest.
+    let parts: Vec<&str> = jwt.0.splitn(2, '.').collect();
+    if parts.len() != 2 {
+        return Err(ValidationError::MalformedJWT);
+    }
+
+    // The token components are individually base64 decoded, decode
+    // just the first part and deserialise it into the expected
+    // representation.
+    let headers_json = base64::decode_config(parts[0], URL_SAFE)
+        .map_err(|_| ValidationError::MalformedJWT)?;
+
+    #[derive(Deserialize)]
+    struct KidOnly {
+        kid: Option<String>,
+    }
+
+    let kid_only: KidOnly = serde_json::from_slice(&headers_json)
+        .map_err(|_| ValidationError::MalformedJWT)?;
+
+    Ok(kid_only.kid)
 }
 
 /// Validate the signature of a JSON Web Token and optionally apply
diff --git a/src/tests.rs b/src/tests.rs
index 512978f07445..9edd7c12adaa 100644
--- a/src/tests.rs
+++ b/src/tests.rs
@@ -20,6 +20,15 @@ fn test_decode_find_jwks() {
 }
 
 #[test]
+fn test_token_kid() {
+    let jwt = JWT("eyJraWQiOiI4ckRxOFB3MEZaY2FvWFdURVZRbzcrVGYyWXpTTDFmQnhOS1BDZWJhYWk0PSIsImFsZyI6IlJTMjU2IiwidHlwIjoiSldUIn0.eyJpc3MiOiJhdXRoLnRlc3QuYXByaWxhLm5vIiwiaWF0IjoxNTM2MDUwNjkzLCJleHAiOjE1MzYwNTQyOTMsInN1YiI6IjQyIiwiZXh0Ijoic21va2V0ZXN0IiwicHJ2IjoiYXJpc3RpIiwic2NwIjoicHJvY2VzcyJ9.gOLsv98109qLkmRK6Dn7WWRHLW7o8W78WZcWvFZoxPLzVO0qvRXXRLYc9h5chpfvcWreLZ4f1cOdvxv31_qnCRSQQPOeQ7r7hj_sPEDzhKjk-q2aoNHaGGJg1vabI--9EFkFsGQfoS7UbMMssS44dgR68XEnKtjn0Vys-Vzbvz_CBSCH6yQhRLik2SU2jR2L7BoFvh4LGZ6EKoQWzm8Z-CHXLGLUs4Hp5aPhF46dGzgAzwlPFW4t9G4DciX1uB4vv1XnfTc5wqJch6ltjKMde1GZwLR757a8dJSBcmGWze3UNE2YH_VLD7NCwH2kkqr3gh8rn7lWKG4AUIYPxsw9CB".into());
+
+    let kid = token_kid(&jwt).expect("Failed to extract token KID");
+    assert_eq!(Some("8rDq8Pw0FZcaoXWTEVQo7+Tf2YzSL1fBxNKPCebaai4=".into()),
+               kid, "Extracted KID did not match expected KID");
+}
+
+#[test]
 fn test_validate_jwt() {
     let jwks_json = "{\"keys\":[{\"kty\":\"RSA\",\"alg\":\"RS256\",\"use\":\"sig\",\"kid\":\"8rDq8Pw0FZcaoXWTEVQo7+Tf2YzSL1fBxNKPCebaai4=\",\"n\":\"l4UTgk1zr-8C8utt0E57DtBV6qqAPWzVRrIuQS2j0_hp2CviaNl5XzGRDnB8gwk0Hx95YOhJupAe6RNq5ok3fDdxL7DLvppJNRLz3Ag9CsmDLcbXgNEQys33fBJaPw1v3GcaFC4tisU5p-o1f5RfWwvwdBtdBfGiwT1GRvbc5sFx6M4iYjg9uv1lNKW60PqSJW4iDYrfqzZmB0zF1SJ0BL_rnQZ1Wi_UkFmNe9arM8W9tI9T3Ie59HITFuyVSTCt6qQEtSfa1e5PiBaVuV3qoFI2jPBiVZQ6LPGBWEDyz4QtrHLdECPPoTF30NN6TSVwwlRbCuUUrdNdXdjYe2dMFQ\",\"e\":\"DhaD5zC7mzaDvHO192wKT_9sfsVmdy8w8T8C9VG17_b1jG2srd3cmc6Ycw-0blDf53Wrpi9-KGZXKHX6_uIuJK249WhkP7N1SHrTJxO0sUJ8AhK482PLF09Qtu6cUfJqY1X1y1S2vACJZItU4Vjr3YAfiVGQXeA8frAf7Sm4O1CBStCyg6yCcIbGojII0jfh2vSB-GD9ok1F69Nmk-R-bClyqMCV_Oq-5a0gqClVS8pDyGYMgKTww2RHgZaFSUcG13KeLMQsG2UOB2OjSC8FkOXK00NBlAjU3d0Vv-IamaLIszO7FQBY3Oh0uxNOvIE9ofQyCOpB-xIK6V9CTTphxw\"}]}";