pdk-jwt-lib 1.10.0

// Copyright (c) 2026, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

//! # Untrusted Token Handling
//!
//! This module provides utilities for handling untrusted JWT tokens before validation,
//! including extraction of key IDs and token metadata.

/// Trait to get the key ID from an untrusted token.
pub trait UntrustedToken {
    fn kid(&self) -> Option<&String>;
}

/// FIPS-safe untrusted-token handling. Validates JWT structure and extracts the
/// header (and its `kid`) with plain base64 + serde_json, carrying no RustCrypto
/// dependency. Used by the shared claims parser and the FIPS signature
/// validator so neither pulls in jwt-compact.
pub mod untrusted_jwt_raw {
    use super::UntrustedToken;
    use crate::error::jwt_error::JWTError;
    use base64::{decode_config, URL_SAFE_NO_PAD};
    use serde_json::Value;
    use std::collections::HashMap;

    /// Structurally-validated (but unverified) JWT.
    pub struct UntrustedJWT {
        kid: Option<String>,
        header: HashMap<String, Value>,
    }

    impl UntrustedJWT {
        /// Create a new [`UntrustedJWT`] from a token string.
        ///
        /// Mirrors jwt-compact's structural check: the token must have exactly
        /// three `.`-separated segments (a trailing empty signature is allowed),
        /// the header segment must base64url-decode to a JSON object carrying a
        /// string `alg`, and the signature segment must be valid base64url.
        pub fn new<S: AsRef<str> + ?Sized>(token_string: &S) -> Result<Self, JWTError> {
            let structure_error = || {
                JWTError::TokenParseFailed(String::from(
                    "A JWT should have xxxxx.yyyyy.zzzzz or xxxxx.yyyyy. format",
                ))
            };

            let parts: Vec<_> = token_string.as_ref().splitn(4, '.').collect();
            let [header_b64, _claims_b64, signature_b64] = parts[..] else {
                return Err(structure_error());
            };

            let header_bytes =
                decode_config(header_b64, URL_SAFE_NO_PAD).map_err(|_| structure_error())?;
            let header: HashMap<String, Value> =
                serde_json::from_slice(&header_bytes).map_err(|_| structure_error())?;

            if !header.get("alg").map(Value::is_string).unwrap_or(false) {
                return Err(structure_error());
            }

            // A non-empty signature must still be valid base64url.
            if !signature_b64.is_empty() {
                decode_config(signature_b64, URL_SAFE_NO_PAD).map_err(|_| structure_error())?;
            }

            let kid = header.get("kid").and_then(Value::as_str).map(String::from);

            Ok(UntrustedJWT { kid, header })
        }

        /// Get the decoded JWT header.
        pub fn header(&self) -> &HashMap<String, Value> {
            &self.header
        }
    }

    impl UntrustedToken for UntrustedJWT {
        fn kid(&self) -> Option<&String> {
            self.kid.as_ref()
        }
    }

    #[cfg(test)]
    mod tests {
        use super::{UntrustedJWT, UntrustedToken};

        // Header {"alg":"RS256","kid":"key-1"}, payload {"sub":"s"}, non-empty sig.
        const VALID_WITH_KID: &str =
            "eyJhbGciOiJSUzI1NiIsImtpZCI6ImtleS0xIn0.eyJzdWIiOiJzIn0.c2lnbmF0dXJl";
        // Header {"alg":"HS256"}, payload {"sub":"s"}, non-empty sig, no kid.
        const VALID_NO_KID: &str = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJzIn0.c2lnbmF0dXJl";
        // Same header/payload, empty (unsigned) signature segment.
        const EMPTY_SIG: &str = "eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJzIn0.";

        #[test]
        fn parses_valid_token_and_extracts_kid() {
            let token = UntrustedJWT::new(VALID_WITH_KID).unwrap();
            assert_eq!(token.kid(), Some(&"key-1".to_string()));
            assert_eq!(
                token.header().get("alg").and_then(|v| v.as_str()),
                Some("RS256")
            );
        }

        #[test]
        fn parses_valid_token_without_kid() {
            let token = UntrustedJWT::new(VALID_NO_KID).unwrap();
            assert_eq!(token.kid(), None);
        }

        #[test]
        fn accepts_empty_signature_segment() {
            let token = UntrustedJWT::new(EMPTY_SIG).unwrap();
            assert_eq!(token.kid(), None);
        }

        #[test]
        fn rejects_token_with_too_few_segments() {
            assert!(UntrustedJWT::new("eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJzIn0").is_err());
        }

        #[test]
        fn rejects_non_base64url_header() {
            assert!(UntrustedJWT::new("not+valid/base64.eyJzdWIiOiJzIn0.c2ln").is_err());
        }

        #[test]
        fn rejects_header_that_is_not_a_json_object() {
            // Header base64url-decodes to the JSON string "astring", not an object.
            assert!(UntrustedJWT::new("ImFzdHJpbmci.eyJzdWIiOiJzIn0.c2lnbmF0dXJl").is_err());
        }

        #[test]
        fn rejects_missing_alg() {
            // Header {"typ":"JWT"} — no alg.
            assert!(UntrustedJWT::new("eyJ0eXAiOiJKV1QifQ.eyJzdWIiOiJzIn0.c2lnbmF0dXJl").is_err());
        }

        #[test]
        fn rejects_non_string_alg() {
            // Header {"alg":123} — alg present but not a string.
            assert!(UntrustedJWT::new("eyJhbGciOjEyM30.eyJzdWIiOiJzIn0.c2lnbmF0dXJl").is_err());
        }

        #[test]
        fn rejects_non_base64url_signature() {
            let token = format!("eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJzIn0.{}", "not+base64/");
            assert!(UntrustedJWT::new(&token).is_err());
        }
    }
}

/// Module to handle untrusted JWT compact tokens.
#[cfg(not(any(feature = "fips", fips)))]
pub mod untrusted_jwt_compact {
    use super::UntrustedToken;
    use crate::error::jwt_error::JWTError;

    /// Struct to handle untrusted JWT compact tokens.
    pub struct UntrustedJWT<'a> {
        untrusted_token: jwt_compact::UntrustedToken<'a>,
    }

    impl<'a> UntrustedJWT<'a> {
        /// Create a new [`UntrustedJWT`] from a token string.
        pub fn new<S: AsRef<str> + ?Sized>(token_string: &'a S) -> Result<Self, JWTError> {
            let untrusted = jwt_compact::UntrustedToken::new(token_string)?;
            Ok(UntrustedJWT {
                untrusted_token: untrusted,
            })
        }
    }

    impl UntrustedToken for UntrustedJWT<'_> {
        fn kid(&self) -> Option<&String> {
            self.untrusted_token.header().key_id.as_ref()
        }
    }

    impl<'a> From<UntrustedJWT<'a>> for jwt_compact::UntrustedToken<'a> {
        fn from(token: UntrustedJWT<'a>) -> Self {
            token.untrusted_token
        }
    }
}

/// Module to handle untrusted JWT tokens which provides a wrapper around
/// untrusted JWT tokens, extracting Key IDs and storing the original token.
#[cfg(not(any(feature = "fips", fips)))]
pub mod untrusted_jwt {
    use super::UntrustedToken;
    use crate::error::jwt_error::JWTError;
    use jwt_simple::prelude::Token;
    use std::option::Option;

    /// Struct to handle untrusted JWT tokens.
    pub struct UntrustedJWT {
        kid: Option<String>,
        value: Option<String>,
    }

    impl UntrustedJWT {
        /// Create a new [`UntrustedJWT`] from a token string.
        pub fn new(token_string: &str) -> Result<Self, JWTError> {
            let untrusted = Token::decode_metadata(token_string)?;

            Ok(UntrustedJWT {
                kid: untrusted.key_id().map(String::from),
                value: Some(String::from(token_string)),
            })
        }

        /// Get the token value.
        pub fn value(&self) -> &String {
            self.value.as_ref().unwrap()
        }
    }

    impl UntrustedToken for UntrustedJWT {
        fn kid(&self) -> Option<&String> {
            self.kid.as_ref()
        }
    }
}