pdk-jwt-lib 1.8.0

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

use crate::api::model::{SigningAlgorithm, SigningKeyLength};
use crate::model::claims::JWTClaims;
use jwt_simple::algorithms::{
    ECDSAP256KeyPairLike, ECDSAP384KeyPairLike, ES256KeyPair, ES384KeyPair, HS256Key, HS384Key,
    HS512Key, MACLike, RS256KeyPair, RS384KeyPair, RS512KeyPair, RSAKeyPairLike,
};
use jwt_simple::reexports::rand;
use p256::ecdsa::signature::RandomizedDigestSigner;

use super::{GeneratorError, JwtGeneratorImpl};

trait JwtAlgorithm {
    fn alg(&self) -> &'static str;
    fn sign(&self, message: &str) -> Result<Vec<u8>, GeneratorError>;

    fn jwt(&self, claims: JWTClaims) -> Result<String, GeneratorError> {
        let mut claims = claims;
        claims.headers.insert("alg".to_string(), self.alg().into());
        claims.headers.insert("typ".to_string(), "JWT".into());
        let headers = super::base64(serde_json::to_vec(&claims.headers)?);
        let body = super::base64(serde_json::to_vec(&claims.claims)?);

        let message = format!("{headers}.{body}");
        let signature = super::base64(self.sign(message.as_str())?);

        Ok(format!("{message}.{signature}"))
    }
}

macro_rules! jwt_rsa {
    ($type_name:ty) => {
        impl JwtAlgorithm for $type_name {
            fn alg(&self) -> &'static str {
                Self::jwt_alg_name()
            }

            fn sign(&self, message: &str) -> Result<Vec<u8>, GeneratorError> {
                let digest = Self::hash(message.as_bytes());
                let mut rng = rand::thread_rng();
                self.key_pair()
                    .as_ref()
                    .sign_blinded(&mut rng, self.padding_scheme(), &digest)
                    .map_err(|e| GeneratorError::SigningError(e.into()))
            }
        }
    };
}

jwt_rsa!(RS256KeyPair);
jwt_rsa!(RS384KeyPair);
jwt_rsa!(RS512KeyPair);

macro_rules! jwt_hmac {
    ($type_name:ty) => {
        impl JwtAlgorithm for $type_name {
            fn alg(&self) -> &'static str {
                Self::jwt_alg_name()
            }

            fn sign(&self, message: &str) -> Result<Vec<u8>, GeneratorError> {
                Ok(self.authentication_tag(message))
            }
        }
    };
}

jwt_hmac!(HS256Key);
jwt_hmac!(HS384Key);
jwt_hmac!(HS512Key);

impl JwtAlgorithm for ES256KeyPair {
    fn alg(&self) -> &'static str {
        Self::jwt_alg_name()
    }

    fn sign(&self, message: &str) -> Result<Vec<u8>, GeneratorError> {
        let mut digest = hmac_sha256::Hash::new();
        digest.update(message.as_bytes());
        let mut rng = rand::thread_rng();
        let signature: p256::ecdsa::Signature = self
            .key_pair()
            .as_ref()
            .sign_digest_with_rng(&mut rng, digest);

        Ok(signature.to_vec())
    }
}

impl JwtAlgorithm for ES384KeyPair {
    fn alg(&self) -> &'static str {
        Self::jwt_alg_name()
    }

    fn sign(&self, message: &str) -> Result<Vec<u8>, GeneratorError> {
        let mut digest = hmac_sha512::sha384::Hash::new();
        digest.update(message.as_bytes());
        let mut rng = rand::thread_rng();
        let signature: p384::ecdsa::Signature = self
            .key_pair()
            .as_ref()
            .sign_digest_with_rng(&mut rng, digest);

        Ok(signature.to_vec())
    }
}

macro_rules! from_pem {
    ($type_name:ty, $key:expr) => {{
        Box::new(<$type_name>::from_pem($key).map_err(GeneratorError::KeyParsing)?)
    }};
}

macro_rules! from_bytes {
    ($type_name:ty, $key:expr) => {{
        Box::new(<$type_name>::from_bytes($key.as_bytes()))
    }};
}

pub(super) fn create_native_generator(
    algorithm: SigningAlgorithm,
    length: SigningKeyLength,
    key: &str,
) -> Result<Box<dyn JwtGeneratorImpl>, GeneratorError> {
    let alg: Box<dyn JwtAlgorithm> = match (&algorithm, &length) {
        (SigningAlgorithm::Hmac, SigningKeyLength::Len256) => from_bytes!(HS256Key, key),
        (SigningAlgorithm::Hmac, SigningKeyLength::Len384) => from_bytes!(HS384Key, key),
        (SigningAlgorithm::Hmac, SigningKeyLength::Len512) => from_bytes!(HS512Key, key),
        (SigningAlgorithm::Rsa, SigningKeyLength::Len256) => from_pem!(RS256KeyPair, key),
        (SigningAlgorithm::Rsa, SigningKeyLength::Len384) => from_pem!(RS384KeyPair, key),
        (SigningAlgorithm::Rsa, SigningKeyLength::Len512) => from_pem!(RS512KeyPair, key),
        (SigningAlgorithm::Es, SigningKeyLength::Len256) => from_pem!(ES256KeyPair, key),
        (SigningAlgorithm::Es, SigningKeyLength::Len384) => from_pem!(ES384KeyPair, key),
        (SigningAlgorithm::Es, SigningKeyLength::Len512) => {
            return Err(GeneratorError::InvalidAlgorithm("ES512".to_string()));
        }
    };

    Ok(Box::new(NativeJwtGenerator { alg }))
}

struct NativeJwtGenerator {
    alg: Box<dyn JwtAlgorithm>,
}

impl JwtGeneratorImpl for NativeJwtGenerator {
    fn jwt(&self, claims: JWTClaims) -> Result<String, GeneratorError> {
        self.alg.jwt(claims)
    }
}