From 7fe7d62c520782b44b818640b73e6e54d33f2246 Mon Sep 17 00:00:00 2001
From: "Sebastian H. Gabrielli" <sebgab@sebaweb.no>
Date: Sat, 30 Dec 2023 22:31:28 +0100
Subject: [PATCH] Handle all errors in validate_jwt

---
 src/main.rs | 43 ++++++++++++++++++++++++++++++++++++-------
 1 file changed, 36 insertions(+), 7 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 1ba1801..be3027f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -44,7 +44,7 @@ struct JwtInfo {
     public_keys: HashMap<String, String>,
 }
 
-fn validate_jwt(token: &str, jwt_info: &JwtInfo) -> Result<MyClaims> {
+fn validate_jwt(token: &str, jwt_info: &mut JwtInfo) -> Result<MyClaims> {
     // Decode the header to give info about the crypto
     let jwt_header = decode_header(token)?;
 
@@ -54,10 +54,39 @@ fn validate_jwt(token: &str, jwt_info: &JwtInfo) -> Result<MyClaims> {
     validation.set_audience(&jwt_info.audience);
     validation.set_issuer(&jwt_info.issuer);
 
-    // Fetch the JWT kid
-    let kid = jwt_header.kid.unwrap();
+    // Extract the JWT kid
+    let kid: String;
+    match jwt_header.kid {
+        Some(fetched_kid) => kid = fetched_kid,
+        None => {
+            eprintln!("Unable to extract KID from jwt header");
+            return Err(jsonwebtoken::errors::ErrorKind::InvalidToken.into());
+        }
+    }
+
     // Fetch the corresponding public key
-    let public_key_pem = jwt_info.public_keys.get(&kid).unwrap();
+    let public_key_pem: &String;
+    match jwt_info.public_keys.get(&kid) {
+        Some(key) => public_key_pem = key,
+        None => {
+            // If the key doesn't exist look up the keys again
+            match fetch_jwt_certificates(jwt_info) {
+                Some(key_map) => jwt_info.public_keys = key_map,
+                None => {
+                    eprintln!("Failed to fetch jwt pem certificates");
+                }
+            }
+
+            // Try to get the keys once more
+            match jwt_info.public_keys.get(&kid) {
+                Some(key) => public_key_pem = key,
+                None => {
+                    eprintln!("Failed to fetch find matching certificates for given KID. {}", kid);
+                    return Err(jsonwebtoken::errors::ErrorKind::InvalidToken.into());
+                }
+            }
+        }
+    }
 
     // Decode the JWT token
     let token_data: TokenData<MyClaims>;
@@ -65,14 +94,14 @@ fn validate_jwt(token: &str, jwt_info: &JwtInfo) -> Result<MyClaims> {
         Algorithm::RS256 => {
             token_data = decode::<MyClaims>(
                 token,
-                &DecodingKey::from_rsa_pem(public_key_pem.as_bytes()).unwrap(),
+                &DecodingKey::from_rsa_pem(public_key_pem.as_bytes())?,
                 &validation,
                 )?;
         },
         Algorithm::ES256 => {
             token_data = decode::<MyClaims>(
                 token,
-                &DecodingKey::from_ec_pem(public_key_pem.as_bytes()).unwrap(),
+                &DecodingKey::from_ec_pem(public_key_pem.as_bytes())?,
                 &validation,
                 )?;
         },
@@ -102,7 +131,7 @@ fn fetch_jwt_certificates(jwt_info: &JwtInfo) -> Option<HashMap<String, String>>
         // Extract the x5c key data
         let x5c = key.x5c.get(0).unwrap();
 
-        // Append the PEM info in to the x5c
+        // Add the PEM info in to the x5c
         let pem_data = format!(
             "-----BEGIN CERTIFICATE-----\n{}\n-----END CERTIFICATE-----",
             x5c