pdk-jwt-lib 1.10.0

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

//! # Public Key Provider
//!
//! This module provides strategies for retrieving public keys used for JWT signature
//! validation from various sources including raw strings and expressions.

use pdk_script::AttributesBinding;

use crate::error::jwt_error::JWTError;
use crate::provider::expression_provider::ExpressionProvider;

/// Strategies to retrieve the public key used for JWT signature from the incoming HTTP requests
/// context.
#[allow(clippy::large_enum_variant)]
pub enum PublicKeyProvider {
    /// Obtains the key as a [`String`] literal.
    Raw { key: String },

    /// Uses an [`ExpressionProvider`] to obtain the key from the request or context.
    Expression { ep: ExpressionProvider },
}

impl PublicKeyProvider {
    pub fn provide(&mut self, binding: Option<&dyn AttributesBinding>) -> Result<String, JWTError> {
        match self {
            PublicKeyProvider::Raw { key } => Ok(key.clone()),
            PublicKeyProvider::Expression { ep } => match binding {
                Some(handler) => {
                    let result: Option<String> = ep.provide(handler);

                    match result {
                        Some(key) => Ok(key),
                        None => Err(JWTError::KidNotPresent),
                    }
                }
                None => Err(JWTError::KeyProvisionMissingHeaderHandler),
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::provider::key_provider::PublicKeyProvider;
    use mockall::mock;
    use pdk_script::HandlerAttributesBinding;

    mock! {
        pub HeadersHandler {}
        impl pdk_core::classy::hl::HeadersHandler for HeadersHandler{
                fn headers(&self) -> Vec<(String, String)>;
                fn header(&self, name: &str) -> Option<String>;
                fn add_header(&self, name: &str, value: &str);
                fn set_header(&self, name: &str, value: &str);
                fn set_headers<'a>(&self, headers: Vec<(&'a str, &'a str)>);
                fn remove_header(&self, name: &str);
        }
    }

    #[test]
    fn raw_public_key_provider_should_return_key() {
        let token = "token";
        let event = MockHeadersHandler::new();

        let mut key_provider = PublicKeyProvider::Raw {
            key: String::from(token),
        };

        assert_eq!(
            key_provider
                .provide(Some(&HandlerAttributesBinding::partial(&event)))
                .unwrap(),
            token
        );
    }
}