pdk-jwt-lib 1.9.0

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

use crate::model::claims::JWTClaims;
use jwt_simple::reexports::anyhow;

#[non_exhaustive]
#[warn(unused_imports)]
#[derive(Debug, thiserror::Error)]
/// Error type for JWT generator operations
pub enum GeneratorError {
    #[error("Invalid algorithm {0:?}")]
    /// The provided algorithm and length provided are not supported by the generator.
    InvalidAlgorithm(String),

    #[error("Error generating the jwt: {0}.")]
    /// There was an unexpected error while creating the JWT
    SigningError(anyhow::Error),

    #[error("Error parsing key: {0}. Key must be provided in pkcs#8 format.")]
    /// There was an error processing the key.
    KeyParsing(anyhow::Error),

    #[error("Error serializing values: {0}.")]
    /// Failed to serialize the claims and headers.
    Serialization(#[from] serde_json::Error),
}

pub(crate) fn base64<D: AsRef<[u8]>>(data: D) -> String {
    base64::encode_config(data.as_ref(), base64::URL_SAFE_NO_PAD)
}

trait JwtGeneratorImpl {
    fn jwt(&self, claims: JWTClaims) -> Result<String, GeneratorError>;
}

cfg_if::cfg_if! {
    if #[cfg(not(fips))] {
        mod native_generator;
    } else {
        mod proxy_wasm_generator;
    }
}

/// Helper to generate JWT tokens
pub struct JwtGenerator {
    inner: Box<dyn JwtGeneratorImpl>,
}

impl JwtGenerator {
    /// Create a new JWT generator. Key must be provided as a PEM with pkcs8 format for RSA and ES algorithms, and plain text for HMAC.
    pub fn new(
        algorithm: crate::api::model::SigningAlgorithm,
        length: crate::api::model::SigningKeyLength,
        key: &str,
    ) -> Result<Self, GeneratorError> {
        #[cfg(not(fips))]
        {
            let inner = native_generator::create_native_generator(algorithm, length, key)?;
            Ok(Self { inner })
        }

        #[cfg(fips)]
        {
            let inner = Box::new(proxy_wasm_generator::ProxyWasmJwtGenerator::new(
                algorithm, length, key,
            )?);
            Ok(Self { inner })
        }
    }

    /// Create a JWT token from the provided claims and headers.
    pub fn jwt(&self, claims: JWTClaims) -> Result<String, GeneratorError> {
        self.inner.jwt(claims)
    }
}

#[cfg(test)]
#[warn(unused_imports)]
mod test {
    type DynamicSimpleClaims = jwt_simple::claims::JWTClaims<HashMap<String, Value>>;
    use crate::api::model::{JWTClaims, SigningAlgorithm, SigningKeyLength};
    use crate::generator::{GeneratorError, JwtGenerator};
    use chrono::{TimeZone, Utc};
    use jwt_simple::algorithms::{
        ECDSAP256PublicKeyLike, ECDSAP384PublicKeyLike, ES256PublicKey, ES384PublicKey, HS256Key,
        HS384Key, HS512Key, MACLike, RS256PublicKey, RS384PublicKey, RS512PublicKey,
        RSAPublicKeyLike,
    };
    use serde_json::Value;
    use std::collections::HashMap;

    // openssl ecparam -name secp256r1 -genkey -noout -out jwtES256key.pem
    // openssl pkcs8 -topk8 -inform PEM -outform PEM -nocrypt -in jwtES256key.pem -out jwtES256key-pkcs8.pem
    // @@Credential-Scanning-Service-Mock@@
    const TEST_EC_PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY-----
MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgyzC5ONgTRTCKUko8
elZbv5YjMntLt/exJgdjmm1l7IWhRANCAARLDDmGlhxW5yzgu7SrO86MtQcKtm1j
2ydR/7bsAOVDGeQB6kVBqvOGZXKmVaT8X2mBc/VZrZkYh5PVapnMngmC
-----END PRIVATE KEY-----";

    // openssl ec -in ./jwtES256key-pkcs8.pem -pubout -out jwtES256key-pkcs8-public.pem
    // @@Credential-Scanning-Service-Mock@@
    const TEST_EC_PUBLIC_KEY: &str = "-----BEGIN PUBLIC KEY-----
MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAESww5hpYcVucs4Lu0qzvOjLUHCrZt
Y9snUf+27ADlQxnkAepFQarzhmVyplWk/F9pgXP1Wa2ZGIeT1WqZzJ4Jgg==
-----END PUBLIC KEY-----";

    // openssl ecparam -name secp384r1 -genkey -noout -out jwtES384key.pem
    // openssl pkcs8 -topk8 -inform PEM -outform PEM -nocrypt -in jwtES384key.pem -out jwtES384key-pkcs8.pem
    // @@Credential-Scanning-Service-Mock@@
    const TEST_EC_PRIVATE_KEY_384: &str = "-----BEGIN PRIVATE KEY-----
MIG2AgEAMBAGByqGSM49AgEGBSuBBAAiBIGeMIGbAgEBBDA8W5KSitwokEoKlCP7
h6DdD7fYD567VjEEEl0L9LcpCfrUWp6CMhTfLoyyj7CoW8qhZANiAARTfi70JZsD
Vf0iS0Gw3+8ylH/7IQwRA3iZpHC5KQDEOrgXZmcEocgxkZD/Ud0X/tZVwJNAqVt9
XhvK7gvogAUePeAklVfRJTuZuJhb7x3NbAczz43nHvLtgSHIWH1nyrk=
-----END PRIVATE KEY-----";

    // openssl ec -in ./jwtES384key-pkcs8.pem -pubout -out jwtES384key-pkcs8-public.pem
    // @@Credential-Scanning-Service-Mock@@
    const TEST_EC_PUBLIC_KEY_384: &str = "-----BEGIN PUBLIC KEY-----
MHYwEAYHKoZIzj0CAQYFK4EEACIDYgAEU34u9CWbA1X9IktBsN/vMpR/+yEMEQN4
maRwuSkAxDq4F2ZnBKHIMZGQ/1HdF/7WVcCTQKlbfV4byu4L6IAFHj3gJJVX0SU7
mbiYW+8dzWwHM8+N5x7y7YEhyFh9Z8q5
-----END PUBLIC KEY-----";

    // @@Credential-Scanning-Service-Mock@@
    const TEST_RSA_PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY-----
MIIJQQIBADANBgkqhkiG9w0BAQEFAASCCSswggknAgEAAoICAQC/OrFHGMNbsquO
eilUI39KUhCknuBlxf+3mXODrDp8iPm9/tQ/E1S2Jatp5ip9cGFv/FA7ndImgUMI
HZnRVkNaCk6Yl/+SA+0C1IIcBvR4Z7rroSXH4vQtBBTnPVyxEDOZrupp/N5Aatht
qgqCldERiPu6y/GSil0UhO/nFrlOZ74NA9BeAtQhRlENkTML9ObaxO/Dv5jzJ4Ik
IZJ8XKforeXq3CZBEm2g1EZEU2WD2vDZxv7WdP1z/Tx6+PLbnLBTy5E7ZUOx9bHo
aM+oXPCzrA2e5c9xhY6M1JP8ccl9qoOUhVUBMiXIRSR0OeasAz/dEC+qlx8YT8Fp
JCQMadR6NuogAWXk2SL1fxtmc0zKYWXJj7Diit3pBmUsXsmOxQGGV9uKeWJuz4aV
3bc0uMGSWcBGoVWOVQtw8dLNFWUdZYvr06hARcSLCRhLbP4CCQBgXe5MqKXDxa4X
wYymhFcxva/4XpenrEdxQeuTacDzPoJAnIqsEOHAl7POcWzJbHXm0PWnQ4GVY+4U
RTide+bW5YfhAv/FojAOG8+nM49laVZuuQ4PkKD5tzoLMVh5mIaUveiJG//cFXvM
GtGALgUNzlcotbG8GYXT2U1EL5znwAOrxHxWrEjBIrT13B+OwmaJ21fpcd8KIzla
s9g4SzAMYxOLbgxF0f2p+mFchc6tbwIDAQABAoICAB2z3KBZ8NI37NzLDctTXiyp
lYs0YE9+kyst6xrbMBRy5DPGNqp7cq9+J2NiDFyCjafqzX2NFHzFnCdRDbjNyNVd
/3pFNb203WYQowr+a4+eMRLza15iWqH5XdPTHKgmB5XJ7QA8djsUPXy/KjXBVoF+
QPdxQRsNYcrToT3IMk1C4Oq9mmpXzyJB/Un5sS+cwRTe/QzvIC84hkbdbhbh/3St
OiaiPlDiL2QJRMbNG1oBMmLpPWELN+kBvxisvXAuJNdHKc5LetnT+2fJi+OvV/XY
dh8lu/R6lbs7M6dE91KNHzX9BciTRPoX/0MMUU+Li6pnHrhFE9/fV3/gzLae45E5
CX2JfciE89W9cO/moPoh9jtfCeu1ojRY3/HcgeRBxppPmZoybUamoQYnC9LpisA2
3Zn+IsPSXHHznhc7gOhXvqLBgSdRs450hqhOyJ/hwXSeeoKbqtx81YIhZeauUuzZ
8oiYibHfFLJhWoH/sqSSF4D1tJQysybKD8Zw2dQz34JW3hsUXMu/mmcYd9almWJm
9toCtmV8pVL5MlbC635A1rfpnbtPaLLEtPoP1XAnHsquvy/yu6Njhj1zqbe2D3nH
5Q9EYikML5A6cBPzYcxtvmST3I58MyD3gPMxxEqYrF5GFNF+ty96MX9y+B4mYcGQ
Yv2sVyXOgQz4Wz16zMWNAoIBAQD1DlBYEcYsklVT/gIOjr2QlmonbCgtoE7rarr4
XXYFUL96IYlQD9BzKocka3Hu7YJtLf0SOzMZzmidc424DwWEMhIeuyPUrCFIgIYl
KzTxKTijwDpmMQQN2JGFP3R6zBSXyEr4XLUnOAu2FfxWSU5jVqaTKD9iB+bCm7+/
prg54EXPIr7Mh9+qkx4Vq1VZs7/5RnmZtDJ2+UbV0G8d9HbuRBz0IHvY1e7JZ8Nz
4G3VM5zETvxvFjgxg4EfcCDPt26yxRLJMazihsMWrtvMjvMUAhiil+Df+XOJMCXx
zrNDjn5M8A7pgilXlGzXqF7JLzB04cDDX/+CwqpW0QDSwHL1AoIBAQDHxPyqUFwe
ZB7Sgwv09cKaji3ez8F0as6Z3uY6odhaBLIJuei3XNO0DF39iYbgEXS96bHeja0f
5SsIdcJfryX1q+w820kCOyB0ZwL/JU86PDfJSn/nJ6bqrArr/R7RbytcN1wKuEcP
93TI9fXHyq2GyEnlO385G+3YZj6JU/d6ZBS0zsW/eCd2B1ze47O1Ti/mM2wGQy7M
vEoLqC5OwvOAisZZg8qWr0/Z31aR6V2lMo6dMdG8RrmAan5DKai1DG2EmFssAJIH
cAXB+EQ2MhyRrhk33nDgBf2mJT4ZLjszBZcutodw3pbzIM62489V0dUmOWKRgPPX
7d5YM2z88shTAoIBABp0cCIB0TYQmhuWKVyu9jH8uvsEhxXd34c0n3iehlYukG07
35oACw3TwoEhBEy54UGuHEryjyKzEMImrl73aC4MRb6Bj22vI2yzS0gJ8Q4z2AR9
hRBxLDHedl8/KXD0RSjZm5ZSU9AnEcSXfQVHpqm8ugDa8HTBy5youbuT4QGGf6LL
6nMkG/ZLKY1HUNB9QjVD8W6xcF09rfL5LHW8ZXZ1bfbA5v3SopOlmwkQamsAxmS+
7iuD548Y1kCxlyk1cULlWZDUxwgxajAxslLT/9PiIgyzfrhPMrTVuNLw8JNTd7kQ
lVuKDLKCuHlTmN/5My77DBdLbscMAt2adI9L7V0CggEAWQmpe9eZV0pUmosiFyo6
dFyOgVKj7Nl2AArjHproLScOm1srKB7NlOA2PDzByri9CbBRQNpwoVipF3o1CiSs
jJT2FCHApqfnzTnkkgf1CgWw75yu6T45HTtVGt2UkNA1yUI7WePMeIdYnAFUbJof
QYWfufYMvE2AcwUPNnIgSYK13+iRJsfM/sRFVmqyvEp++uFMcnYbM9FwR0XMbfpi
QZaY1WjyMLsuofLzSNF0lZ61BccgrgPvxhaw9AprUVaasZCegjw22e3KAyw+atFm
/l9UihwwvwishxLuXJbId/Mz8PQV5e6v5OloeQeMb7m4gPLuxd9tz34LrdAt8Yfc
VQKCAQBIi8Hq+efP5VeUkJ9aNplDk4mZsUx7bDZXu/x+oNncggJkO1rWSJf0MX+H
noC28uo0yL+OKo7H33XSfCmiTvXKw8EyAwSk5QoPIG3FnDI22IvzUyhRqrNKjA1a
LyitXE6ZcACrGQTvnT2J4pQqQICTVl7gczY5qZKnfhnE9q20jfgwnQyElU+qQtuw
GgHRzWyYVMDY+soTKyAPY161jvoWRgJOrfVC8wvb8xiMRVn1HJXTDuoK1MX35pGK
EjInQztINGCr1RfQeA/LKdDpqJREHaAfsKLJWIjkXFYnSnaeT4uG+GeiU3yQpNeS
AhW8CZzWMXVPOgExNZJBW65D7p/V
-----END PRIVATE KEY-----";

    // @@Credential-Scanning-Service-Mock@@
    const TEST_RSA_PUBLIC_KEY: &str = "-----BEGIN PUBLIC KEY-----
MIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEAvzqxRxjDW7KrjnopVCN/
SlIQpJ7gZcX/t5lzg6w6fIj5vf7UPxNUtiWraeYqfXBhb/xQO53SJoFDCB2Z0VZD
WgpOmJf/kgPtAtSCHAb0eGe666Elx+L0LQQU5z1csRAzma7qafzeQGrYbaoKgpXR
EYj7usvxkopdFITv5xa5Tme+DQPQXgLUIUZRDZEzC/Tm2sTvw7+Y8yeCJCGSfFyn
6K3l6twmQRJtoNRGRFNlg9rw2cb+1nT9c/08evjy25ywU8uRO2VDsfWx6GjPqFzw
s6wNnuXPcYWOjNST/HHJfaqDlIVVATIlyEUkdDnmrAM/3RAvqpcfGE/BaSQkDGnU
ejbqIAFl5Nki9X8bZnNMymFlyY+w4ord6QZlLF7JjsUBhlfbinlibs+Gld23NLjB
klnARqFVjlULcPHSzRVlHWWL69OoQEXEiwkYS2z+AgkAYF3uTKilw8WuF8GMpoRX
Mb2v+F6Xp6xHcUHrk2nA8z6CQJyKrBDhwJezznFsyWx15tD1p0OBlWPuFEU4nXvm
1uWH4QL/xaIwDhvPpzOPZWlWbrkOD5Cg+bc6CzFYeZiGlL3oiRv/3BV7zBrRgC4F
Dc5XKLWxvBmF09lNRC+c58ADq8R8VqxIwSK09dwfjsJmidtX6XHfCiM5WrPYOEsw
DGMTi24MRdH9qfphXIXOrW8CAwEAAQ==
-----END PUBLIC KEY-----";

    fn base64decode<D: AsRef<[u8]>>(data: D) -> Vec<u8> {
        base64::decode_config(data, base64::URL_SAFE_NO_PAD).unwrap()
    }

    type Claims = HashMap<String, Value>;

    fn number(map: &Claims, key: &str) -> u64 {
        map.get(key).unwrap().as_u64().unwrap()
    }

    fn get_str<'a>(map: &'a Claims, key: &str) -> &'a str {
        map.get(key).unwrap().as_str().unwrap()
    }

    fn decode_jwt(jwt: &str) -> (Claims, Claims) {
        let header: Claims =
            serde_json::from_slice(base64decode(jwt.split('.').next().unwrap()).as_slice())
                .unwrap();
        let claims: Claims =
            serde_json::from_slice(base64decode(jwt.split('.').nth(1).unwrap()).as_slice())
                .unwrap();
        (header, claims)
    }

    pub fn smoke(alg: SigningAlgorithm, length: SigningKeyLength, key: &str) -> String {
        let generator = JwtGenerator::new(alg, length, key).unwrap();

        let mut custom_headers = HashMap::new();
        custom_headers.insert("custom_header".to_string(), Value::from("custom_header"));

        let mut custom_claims = HashMap::new();
        custom_claims.insert("custom_claim".to_string(), Value::from("custom_claim"));

        let claims = JWTClaims::new(
            Some(Utc.timestamp_opt(i32::MAX as i64, 0).single().unwrap()),
            None,
            None,
            custom_claims,
            custom_headers,
        )
        .unwrap();

        let jwt = generator.jwt(claims).unwrap();

        println!("{jwt}");

        jwt
    }

    fn validate(jwt: String, alg: &str) {
        let (header, claims) = decode_jwt(&jwt);

        assert_eq!(header.len(), 3);
        assert_eq!(get_str(&header, "alg"), alg);
        assert_eq!(get_str(&header, "typ"), "JWT");
        assert_eq!(get_str(&header, "custom_header"), "custom_header");

        assert_eq!(claims.len(), 2);
        assert_eq!(number(&claims, "exp"), i32::MAX as u64);
        assert_eq!(get_str(&claims, "custom_claim"), "custom_claim");
    }

    #[test]
    #[cfg(not(fips))]
    pub fn hmac256() {
        let key = "some";
        let jwt = smoke(SigningAlgorithm::Hmac, SigningKeyLength::Len256, key);
        let _: DynamicSimpleClaims = HS256Key::from_bytes(key.as_bytes())
            .verify_token(jwt.as_str(), None)
            .unwrap();

        validate(jwt, "HS256");
    }

    #[test]
    #[cfg(not(fips))]
    pub fn hmac384() {
        let key = "some";
        let jwt = smoke(SigningAlgorithm::Hmac, SigningKeyLength::Len384, key);
        let _: DynamicSimpleClaims = HS384Key::from_bytes(key.as_bytes())
            .verify_token(jwt.as_str(), None)
            .unwrap();

        validate(jwt, "HS384");
    }

    #[test]
    #[cfg(not(fips))]
    pub fn hmac512() {
        let key = "some";
        let jwt = smoke(SigningAlgorithm::Hmac, SigningKeyLength::Len512, key);
        let _: DynamicSimpleClaims = HS512Key::from_bytes(key.as_bytes())
            .verify_token(jwt.as_str(), None)
            .unwrap();

        validate(jwt, "HS512");
    }

    #[test]
    pub fn rsa256() {
        let jwt = smoke(
            SigningAlgorithm::Rsa,
            SigningKeyLength::Len256,
            TEST_RSA_PRIVATE_KEY,
        );

        let _: DynamicSimpleClaims = RS256PublicKey::from_pem(TEST_RSA_PUBLIC_KEY)
            .unwrap()
            .verify_token(jwt.as_str(), None)
            .unwrap();

        validate(jwt, "RS256");
    }

    #[test]
    pub fn rsa384() {
        let jwt = smoke(
            SigningAlgorithm::Rsa,
            SigningKeyLength::Len384,
            TEST_RSA_PRIVATE_KEY,
        );

        let _: DynamicSimpleClaims = RS384PublicKey::from_pem(TEST_RSA_PUBLIC_KEY)
            .unwrap()
            .verify_token(jwt.as_str(), None)
            .unwrap();

        validate(jwt, "RS384");
    }

    #[test]
    pub fn rsa512() {
        let jwt = smoke(
            SigningAlgorithm::Rsa,
            SigningKeyLength::Len512,
            TEST_RSA_PRIVATE_KEY,
        );

        let _: DynamicSimpleClaims = RS512PublicKey::from_pem(TEST_RSA_PUBLIC_KEY)
            .unwrap()
            .verify_token(jwt.as_str(), None)
            .unwrap();

        validate(jwt, "RS512");
    }

    #[test]
    pub fn es256() {
        let jwt = smoke(
            SigningAlgorithm::Es,
            SigningKeyLength::Len256,
            TEST_EC_PRIVATE_KEY,
        );

        let _: DynamicSimpleClaims = ES256PublicKey::from_pem(TEST_EC_PUBLIC_KEY)
            .unwrap()
            .verify_token(jwt.as_str(), None)
            .unwrap();

        validate(jwt, "ES256");
    }

    #[test]
    pub fn es384() {
        let jwt = smoke(
            SigningAlgorithm::Es,
            SigningKeyLength::Len384,
            TEST_EC_PRIVATE_KEY_384,
        );

        let _: DynamicSimpleClaims = ES384PublicKey::from_pem(TEST_EC_PUBLIC_KEY_384)
            .unwrap()
            .verify_token(jwt.as_str(), None)
            .unwrap();

        validate(jwt, "ES384");
    }

    #[test]
    pub fn es512() {
        let result = JwtGenerator::new(
            SigningAlgorithm::Es,
            SigningKeyLength::Len512,
            TEST_EC_PRIVATE_KEY_384,
        );

        let Err(GeneratorError::InvalidAlgorithm(_)) = result else {
            panic!("should not create a validator")
        };
    }

    #[test]
    pub fn invalid_key() {
        let result = JwtGenerator::new(
            SigningAlgorithm::Es,
            SigningKeyLength::Len256,
            TEST_EC_PRIVATE_KEY_384,
        );

        let Err(GeneratorError::KeyParsing(_)) = result else {
            panic!("should not create a validator")
        };
    }
}