pdk-jwt-lib 1.10.0

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

//! # Proxy WASM JWT Generator
//!
//! The following signing relies on a proxy_wasm FFI with the following signature:
//! sign([`pdk_core::ffi::SignArguments`]) -> [`pdk_core::ffi::SignResult`]
//! FIPS compliance is not guaranteed by this generator. FIPS compliance needs to be supported by
//! the host implementing the sign ABI. Envoy with FIPS enabled provides such features.

use crate::api::model::{SigningAlgorithm, SigningKeyLength};
use crate::model::claims::JWTClaims;
use pdk_core::ffi::{self, Message};

use super::{base64, GeneratorError, JwtGeneratorImpl};

mod hostcall {

    #[cfg_attr(test, faux::create)]
    #[derive(Default)]
    pub struct ProxyWasmCall;

    #[cfg_attr(test, faux::methods)]
    impl ProxyWasmCall {
        pub fn call(
            &self,
            function_name: &str,
            arguments: Option<&[u8]>,
        ) -> Result<
            Option<pdk_core::classy::proxy_wasm::types::Bytes>,
            pdk_core::classy::proxy_wasm::types::Status,
        > {
            pdk_core::classy::proxy_wasm::hostcalls::call_foreign_function(function_name, arguments)
        }
    }
}

pub struct ProxyWasmJwtGenerator {
    key_length: SigningKeyLength,
    private_key_pem: String,
    alg_name: &'static str,
    ffi: hostcall::ProxyWasmCall,
}

impl ProxyWasmJwtGenerator {
    #[cfg_attr(test, allow(dead_code))]
    pub(crate) fn new_with_ffi(
        algorithm: SigningAlgorithm,
        key_length: SigningKeyLength,
        key: &str,
        ffi: hostcall::ProxyWasmCall,
    ) -> Result<Self, GeneratorError> {
        let alg_name = match (&algorithm, &key_length) {
            (SigningAlgorithm::Rsa, SigningKeyLength::Len256) => "RS256",
            (SigningAlgorithm::Rsa, SigningKeyLength::Len384) => "RS384",
            (SigningAlgorithm::Rsa, SigningKeyLength::Len512) => "RS512",
            (SigningAlgorithm::Es, SigningKeyLength::Len256) => "ES256",
            (SigningAlgorithm::Es, SigningKeyLength::Len384) => "ES384",
            (SigningAlgorithm::Es, SigningKeyLength::Len512) => {
                return Err(GeneratorError::InvalidAlgorithm("ES512".to_string()));
            }
            #[allow(unreachable_patterns)]
            _ => {
                return Err(GeneratorError::InvalidAlgorithm(
                    "Algorithm not supported in FIPS mode".to_string(),
                ));
            }
        };

        Ok(Self {
            key_length,
            private_key_pem: key.to_string(),
            alg_name,
            ffi,
        })
    }

    pub fn new(
        algorithm: SigningAlgorithm,
        key_length: SigningKeyLength,
        key: &str,
    ) -> Result<Self, GeneratorError> {
        // `default()`, not the unit literal: under `#[cfg(test)]` faux replaces
        // ProxyWasmCall with a generated (non-unit) type whose Default is a real
        // constructor. clippy's unit-struct lint is a false positive here.
        #[allow(clippy::default_constructed_unit_structs)]
        Self::new_with_ffi(
            algorithm,
            key_length,
            key,
            hostcall::ProxyWasmCall::default(),
        )
    }

    fn hash_function(&self) -> &'static str {
        match self.key_length {
            SigningKeyLength::Len256 => "sha256",
            SigningKeyLength::Len384 => "sha384",
            SigningKeyLength::Len512 => "sha512",
        }
    }

    fn sign(&self, message: &str) -> Result<Vec<u8>, GeneratorError> {
        let args = ffi::SignArguments {
            hash_function: self.hash_function().to_string(),
            private_key: vec![],
            text: message.as_bytes().to_vec(),
            private_key_pem: self.private_key_pem.clone(),
            ..Default::default()
        };

        let args = args.write_to_bytes().map_err(|e| {
            GeneratorError::SigningError(anyhow::anyhow!("Failed to serialize: {e}"))
        })?;

        match self.ffi.call("sign", Some(&args)) {
            Ok(Some(bytes)) => match ffi::SignResult::parse_from_bytes(&bytes) {
                Ok(result) if result.result => Ok(result.signature.to_vec()),
                Ok(result) => Err(GeneratorError::SigningError(anyhow::anyhow!(
                    "Sign failed: {}",
                    result.error
                ))),
                Err(e) => Err(GeneratorError::SigningError(anyhow::anyhow!(
                    "Invalid deserialization: {e}"
                ))),
            },
            Ok(None) => Err(GeneratorError::SigningError(anyhow::anyhow!(
                "Sign returned no result"
            ))),
            Err(e) => Err(GeneratorError::SigningError(anyhow::anyhow!(
                "Sign hostcall failed: {e:?}"
            ))),
        }
    }
}

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

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

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

#[cfg(test)]
mod tests {
    use super::*;
    use crate::api::model::JWTClaims;
    use chrono::{TimeZone, Utc};
    use pdk_core::ffi::{Message, SignResult};
    use serde_json::Value;
    use std::collections::HashMap;

    // @@Credential-Scanning-Service-Mock@@
    const TEST_EC_PRIVATE_KEY: &str = "-----BEGIN PRIVATE KEY-----
MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQgyzC5ONgTRTCKUko8
elZbv5YjMntLt/exJgdjmm1l7IWhRANCAARLDDmGlhxW5yzgu7SrO86MtQcKtm1j
2ydR/7bsAOVDGeQB6kVBqvOGZXKmVaT8X2mBc/VZrZkYh5PVapnMngmC
-----END PRIVATE KEY-----";

    // @@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-----";

    // @@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-----";

    fn create_test_claims() -> JWTClaims {
        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"));

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

    fn create_mock_sign_success() -> hostcall::ProxyWasmCall {
        let mut mock = hostcall::ProxyWasmCall::faux();
        faux::when!(mock.call("sign", _)).then(|(_, _)| {
            let result = SignResult {
                result: true,
                signature: vec![0x01, 0x02, 0x03, 0x04, 0x05], // Mock signature
                error: String::new(),
                ..Default::default()
            };
            Ok(Some(result.write_to_bytes().unwrap()))
        });
        mock
    }

    fn assert_jwt_structure(jwt: &str, expected_alg: &str) {
        let parts: Vec<&str> = jwt.split('.').collect();
        assert_eq!(
            parts.len(),
            3,
            "JWT must have 3 parts: header.body.signature"
        );

        let header_bytes = base64::decode_config(parts[0], base64::URL_SAFE_NO_PAD).unwrap();
        let header: HashMap<String, Value> = serde_json::from_slice(&header_bytes).unwrap();
        assert_eq!(
            header.get("alg").and_then(|v| v.as_str()),
            Some(expected_alg)
        );
        assert_eq!(header.get("typ").and_then(|v| v.as_str()), Some("JWT"));
        assert_eq!(
            header.get("custom_header").and_then(|v| v.as_str()),
            Some("custom_header")
        );

        let claims_bytes = base64::decode_config(parts[1], base64::URL_SAFE_NO_PAD).unwrap();
        let claims: HashMap<String, Value> = serde_json::from_slice(&claims_bytes).unwrap();
        assert_eq!(
            claims.get("custom_claim").and_then(|v| v.as_str()),
            Some("custom_claim")
        );
    }

    #[test]
    fn rsa256_signs_correctly() {
        let mock = create_mock_sign_success();
        let generator = ProxyWasmJwtGenerator::new_with_ffi(
            SigningAlgorithm::Rsa,
            SigningKeyLength::Len256,
            TEST_RSA_PRIVATE_KEY,
            mock,
        )
        .unwrap();

        let jwt = generator.jwt(create_test_claims()).unwrap();
        assert_jwt_structure(&jwt, "RS256");
    }

    #[test]
    fn rsa384_signs_correctly() {
        let mock = create_mock_sign_success();
        let generator = ProxyWasmJwtGenerator::new_with_ffi(
            SigningAlgorithm::Rsa,
            SigningKeyLength::Len384,
            TEST_RSA_PRIVATE_KEY,
            mock,
        )
        .unwrap();

        let jwt = generator.jwt(create_test_claims()).unwrap();
        assert_jwt_structure(&jwt, "RS384");
    }

    #[test]
    fn rsa512_signs_correctly() {
        let mock = create_mock_sign_success();
        let generator = ProxyWasmJwtGenerator::new_with_ffi(
            SigningAlgorithm::Rsa,
            SigningKeyLength::Len512,
            TEST_RSA_PRIVATE_KEY,
            mock,
        )
        .unwrap();

        let jwt = generator.jwt(create_test_claims()).unwrap();
        assert_jwt_structure(&jwt, "RS512");
    }

    #[test]
    fn es256_signs_correctly() {
        let mock = create_mock_sign_success();
        let generator = ProxyWasmJwtGenerator::new_with_ffi(
            SigningAlgorithm::Es,
            SigningKeyLength::Len256,
            TEST_EC_PRIVATE_KEY,
            mock,
        )
        .unwrap();

        let jwt = generator.jwt(create_test_claims()).unwrap();
        assert_jwt_structure(&jwt, "ES256");
    }

    #[test]
    fn es384_signs_correctly() {
        let mock = create_mock_sign_success();
        let generator = ProxyWasmJwtGenerator::new_with_ffi(
            SigningAlgorithm::Es,
            SigningKeyLength::Len384,
            TEST_EC_PRIVATE_KEY_384,
            mock,
        )
        .unwrap();

        let jwt = generator.jwt(create_test_claims()).unwrap();
        assert_jwt_structure(&jwt, "ES384");
    }

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

        assert!(matches!(result, Err(GeneratorError::InvalidAlgorithm(_))));
    }

    #[test]
    fn sign_failed_returns_error() {
        let mut mock = hostcall::ProxyWasmCall::faux();
        faux::when!(mock.call("sign", _)).then(|(_, _)| {
            let result = SignResult {
                result: false,
                signature: vec![],
                error: "signature verification failed".to_string(),
                ..Default::default()
            };
            Ok(Some(result.write_to_bytes().unwrap()))
        });

        let generator = ProxyWasmJwtGenerator::new_with_ffi(
            SigningAlgorithm::Rsa,
            SigningKeyLength::Len256,
            TEST_RSA_PRIVATE_KEY,
            mock,
        )
        .unwrap();

        let result = generator.jwt(create_test_claims());
        assert!(result.is_err());
    }

    #[test]
    fn sign_returns_none() {
        let mut mock = hostcall::ProxyWasmCall::faux();
        faux::when!(mock.call("sign", _)).then(|(_, _)| Ok(None));

        let generator = ProxyWasmJwtGenerator::new_with_ffi(
            SigningAlgorithm::Rsa,
            SigningKeyLength::Len256,
            TEST_RSA_PRIVATE_KEY,
            mock,
        )
        .unwrap();

        let result = generator.jwt(create_test_claims());
        assert!(result.is_err());
    }
}