pdk-jwt-lib 1.10.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 superboring::pkey::PKey;
use superboring::sign::Signer;

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> {
                // RSA blinding is preserved: superboring's Signer::sign_to_vec internally
                // creates a fresh rand::thread_rng() and passes it to rsa::SigningKey::
                // sign_digest_with_rng, which blinds via rsa_decrypt(Some(rng), ...) to
                // resist timing side-channel attacks (Marvin-style). See superboring
                // implementation.rs:1329-1338 and rsa pkcs1v15/signing_key.rs:130-138.
                let pkey = PKey::from_rsa(self.key_pair().as_ref().clone())
                    .map_err(|e| GeneratorError::SigningError(e.into()))?;
                let mut signer = Signer::new(Self::hash(), &pkey)
                    .map_err(|e| GeneratorError::SigningError(e.into()))?;
                signer
                    .set_rsa_padding(self.padding_scheme())
                    .map_err(|e| GeneratorError::SigningError(e.into()))?;
                signer
                    .update(message.as_bytes())
                    .map_err(|e| GeneratorError::SigningError(e.into()))?;
                signer
                    .sign_to_vec()
                    .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.as_bytes()))
            }
        }
    };
}

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)
    }
}