cedarling 0.0.8

// This software is available under the Apache-2.0 license.
// See https://www.apache.org/licenses/LICENSE-2.0.txt for full text.
//
// Copyright (c) 2024, Gluu, Inc.

#![cfg(test)]

use std::sync::{Arc, LazyLock};

use super::http_utils::OpenIdConfig;
use super::status_list::{self, StatusBitSize};
use super::validation::RefJwtStatusList;
use crate::common::issuer_utils::IssClaim;
use crate::common::policy_store::TrustedIssuer;
use jsonwebtoken::DecodingKey;
use mockito::{Mock, Server, ServerGuard};
use reqwest::Client;
use serde::Serialize;
use serde_json::{json, Value};
use url::Url;

use {jsonwebkey as jwk, jsonwebtoken as jwt};

/// A pair of encoding and decoding keys.
#[derive(Clone)]
pub(crate) struct KeyPair {
    kid: Option<String>,
    encoding_key: jwt::EncodingKey,
    decoding_key: jwt::jwk::Jwk,
    alg: jwt::Algorithm,
}

impl KeyPair {
    pub(crate) fn decoding_key(&self) -> Result<Arc<DecodingKey>, jsonwebtoken::errors::Error> {
        DecodingKey::from_jwk(&self.decoding_key).map(Arc::new)
    }
}

#[derive(Debug, thiserror::Error)]
pub(crate) enum KeyGenerationError {
    #[error("Failed to serialize the decoding key onto the right struct")]
    SerializeDecodingKey(#[from] serde_json::Error),
    #[error("The given key was generated with the wrong algorithm")]
    KeyMismatch,
}

/// Generates a HS256-signed token using the given claims.
pub(crate) fn generate_keypair_hs256(
    kid: Option<impl ToString>,
) -> Result<KeyPair, KeyGenerationError> {
    let mut jwk = jwk::JsonWebKey::new(
        jwk::Key::try_generate_symmetric(256).expect("invalid symmetric key size"),
    );
    jwk.set_algorithm(jwk::Algorithm::HS256)
        .expect("should set encryption algorithm");
    jwk.key_id = Some("some_id".to_string());

    // since this is a symmetric key, the public key is the same as the private
    let mut decoding_key = serde_json::to_value(jwk.key.clone())?;

    // set the key parameters
    if let Some(kid) = &kid {
        decoding_key["kid"] = serde_json::Value::String(kid.to_string());
    }
    let mut decoding_key: jwt::jwk::Jwk = serde_json::from_value(decoding_key)?;
    decoding_key.common.key_algorithm = Some(jwt::jwk::KeyAlgorithm::HS256);

    let encoding_key = match *jwk.key {
        jsonwebkey::Key::Symmetric { key } => jwt::EncodingKey::from_secret(&key),
        _ => Err(KeyGenerationError::KeyMismatch)?,
    };

    Ok(KeyPair {
        kid: kid.map(|s| s.to_string()),
        encoding_key,
        decoding_key,
        alg: jwt::Algorithm::HS256,
    })
}

#[derive(Debug, thiserror::Error)]
pub(crate) enum TokenGenerationError {
    #[error("Failed to encode token into a JWT string")]
    Encode(#[from] jwt::errors::Error),
}

/// Generates a token string in the given format: `"header.claim.signature"`
pub(crate) fn generate_token_using_claims(
    claims: &impl Serialize,
    keypair: &KeyPair,
) -> Result<String, TokenGenerationError> {
    let header = jwt::Header {
        alg: keypair.alg,
        kid: keypair.kid.clone(),
        ..Default::default()
    };

    // serialize token to a string
    Ok(jwt::encode(&header, &claims, &keypair.encoding_key)?)
}

/// Generates a `JwkSet` from the given keys
pub(crate) fn generate_jwks(keys: &[KeyPair]) -> jwt::jwk::JwkSet {
    let keys = keys
        .iter()
        .map(|key_pair| key_pair.decoding_key.clone())
        .collect::<Vec<jwt::jwk::Jwk>>();
    jwt::jwk::JwkSet { keys }
}

pub(crate) struct MockServer {
    pub endpoints: MockEndpoints,
    server: ServerGuard,
    keys: KeyPair,
}

pub(crate) struct MockEndpoints {
    pub oidc: Option<Mock>,
    pub jwks: Option<Mock>,
    pub status_list: Option<Mock>,
}

impl MockEndpoints {
    pub(crate) fn new_with_defaults(server: &mut Server, keys: &KeyPair) -> Self {
        let oidc = Some(
            server
                .mock("GET", "/.well-known/openid-configuration")
                .with_status(200)
                .with_header("content-type", "application/json")
                .with_body(
                    json!({
                        "issuer": server.url(),
                        "jwks_uri": server.url() + MOCK_JWKS_URI,
                    })
                    .to_string(),
                )
                .expect(1)
                .create(),
        );

        let jwks = Some(
            server
                .mock("GET", MOCK_JWKS_URI)
                .with_status(200)
                .with_header("content-type", "application/json")
                .with_body(
                    json!({"keys": generate_jwks(std::slice::from_ref(keys)).keys}).to_string(),
                )
                .expect(1)
                .create(),
        );

        Self {
            oidc,
            jwks,
            status_list: None,
        }
    }
}

#[derive(Clone, Copy)]
pub enum TokenTypeHeader {
    /// "typ": "statuslist+jwt"
    StatusListJwt,
    /// "typ": "JWT"
    Jwt,
}

impl From<TokenTypeHeader> for &str {
    fn from(val: TokenTypeHeader) -> Self {
        match val {
            TokenTypeHeader::StatusListJwt => "statuslist+jwt",
            TokenTypeHeader::Jwt => "JWT",
        }
    }
}

impl From<TokenTypeHeader> for Option<String> {
    fn from(val: TokenTypeHeader) -> Self {
        let typ_str: &str = val.into();
        Some(typ_str.into())
    }
}

const MOCK_OIDC_ENDPOINT: &str = "/.well-known/openid-configuration";
const MOCK_STATUS_LIST_ENDPOINT: &str = "/jans-auth/restv1/status_list";
const MOCK_JWKS_URI: &str = "/jans-auth/restv1/jwks";

impl MockServer {
    pub(crate) async fn new_with_defaults() -> Result<Self, KeyGenerationError> {
        let mut server = Server::new_async().await;

        let keys = generate_keypair_hs256(Some("some_hs256_key"))?;

        let endpoints = MockEndpoints::new_with_defaults(&mut server, &keys);

        Ok(Self {
            endpoints,
            server,
            keys,
        })
    }

    /// Generates an HS256 signed token string in the given format: `"header.claim.signature"`
    ///
    /// If `jwt_status_idx` is [`Some`], the jwt will also havae a `status` claim.
    #[track_caller]
    pub(crate) fn generate_token_with_hs256sig(
        &mut self,
        claims: &mut Value,
        jwt_status_idx: Option<usize>,
    ) -> Result<String, TokenGenerationError> {
        let header = jwt::Header {
            alg: self.keys.alg,
            kid: self.keys.kid.clone(),
            typ: TokenTypeHeader::StatusListJwt.into(),
            ..Default::default()
        };

        claims["iss"] = json!(self.server.url());

        if let Some(idx) = jwt_status_idx {
            claims["status"] = json!({"status_list": RefJwtStatusList {
                idx,
                uri: self
                    .status_list_endpoint()
                    .expect("the status list endpoint hasn't been generated yet. call `generate_status_list_endpoint` first")
                    .to_string(),
                ttl: None,
            }});
        }

        // serialize token to a string
        let jwt = jwt::encode(&header, &claims, &self.keys.encoding_key)?;

        Ok(jwt)
    }

    /// Generates a [`TrustedIssuer`] for this instance of the [`MockServer`].
    pub(crate) fn trusted_issuer(&self) -> TrustedIssuer {
        let mut issuer = TrustedIssuer::default();
        issuer.set_oidc_endpoint(
            Url::parse(&(self.server.url() + MOCK_OIDC_ENDPOINT)).expect("should be a valid url"),
        );
        issuer
    }

    /// Use this to generate the status list JWT
    #[track_caller]
    pub(crate) fn generate_status_list_endpoint(
        &mut self,
        status_list_bits: StatusBitSize,
        status_list: &[u8],
        ttl: Option<u64>,
    ) {
        // building the status list JWT
        let bits: u8 = status_list_bits.into();
        let lst = status_list::compress_and_encode(status_list);
        let sub = format!("{}{}", self.server.url(), MOCK_STATUS_LIST_ENDPOINT);
        let iss = self.server.url();
        let header = jwt::Header {
            alg: self.keys.alg,
            kid: self.keys.kid.clone(),
            typ: TokenTypeHeader::StatusListJwt.into(),
            ..Default::default()
        };
        let encoding_key = self.keys.encoding_key.clone();
        let build_jwt_claims = move || {
            let now = chrono::Utc::now().timestamp();
            let exp = now + 3600; // defaults to 1 hour
            let ttl_secs = ttl.unwrap_or(300); // defaults to 5 mins if the ttl is None
            let claims = json!({
                "sub": sub,
                "status_list": {
                  "bits": bits,
                  "lst": lst,
                },
                "iss": iss,
                "exp": exp,
                "ttl": ttl_secs,
                "iat": now,
            });

            jwt::encode(&header, &claims, &encoding_key)
                .expect("encode status list JWT")
                .as_bytes()
                .to_vec()
        };

        let endpoint = Some(
            self.server
                .mock("GET", MOCK_STATUS_LIST_ENDPOINT)
                .with_status(200)
                .with_header("content-type", "application/statuslist+jwt")
                .with_body_from_request(move |_| build_jwt_claims())
                .expect(1)
                .create(),
        );
        self.endpoints.status_list = endpoint;
    }

    /// Same as [`Self::generate_status_list_endpoint`] but produces a Status List JWT
    /// that intentionally omits the `ttl` claim. Used to exercise the bootstrap-config
    /// fallback refresh interval.
    #[track_caller]
    pub(crate) fn generate_status_list_endpoint_without_ttl(
        &mut self,
        status_list_bits: StatusBitSize,
        status_list: &[u8],
    ) {
        let bits: u8 = status_list_bits.into();
        let lst = status_list::compress_and_encode(status_list);
        let sub = format!("{}{}", self.server.url(), MOCK_STATUS_LIST_ENDPOINT);
        let iss = self.server.url();
        let header = jwt::Header {
            alg: self.keys.alg,
            kid: self.keys.kid.clone(),
            typ: TokenTypeHeader::StatusListJwt.into(),
            ..Default::default()
        };
        let encoding_key = self.keys.encoding_key.clone();
        let build_jwt_claims = move || {
            let now = chrono::Utc::now().timestamp();
            let exp = now + 3600;
            let claims = json!({
                "sub": sub,
                "status_list": {
                  "bits": bits,
                  "lst": lst,
                },
                "iss": iss,
                "exp": exp,
                "iat": now,
            });

            jwt::encode(&header, &claims, &encoding_key)
                .expect("encode status list JWT")
                .as_bytes()
                .to_vec()
        };

        let endpoint = Some(
            self.server
                .mock("GET", MOCK_STATUS_LIST_ENDPOINT)
                .with_status(200)
                .with_header("content-type", "application/statuslist+jwt")
                .with_body_from_request(move |_| build_jwt_claims())
                .expect(1)
                .create(),
        );
        self.endpoints.status_list = endpoint;
    }

    /// Replaces the status list endpoint mock with one that always returns 500,
    /// simulating an unavailable status endpoint so the background refresh fails.
    /// Dropping the previous `Mock` removes it from the server.
    pub(crate) fn fail_status_list_endpoint(&mut self) {
        let endpoint = Some(
            self.server
                .mock("GET", MOCK_STATUS_LIST_ENDPOINT)
                .with_status(500)
                .create(),
        );
        self.endpoints.status_list = endpoint;
    }

    /// Updates the `OpenID` configuration mock to include the status list endpoint.
    /// This should be called after `generate_status_list_endpoint` to ensure the
    /// OIDC configuration returned via HTTP includes the `status_list_endpoint` field.
    pub(crate) fn update_openid_config_with_status_list_endpoint(&mut self) {
        // Remove the existing OIDC mock (drop the Mock)
        let old = self.endpoints.oidc.take();
        drop(old); // explicit drop for clarity

        let status_list_endpoint = self
            .status_list_endpoint()
            .expect("status list endpoint not generated");
        let body = json!({
            "issuer": self.server.url(),
            "jwks_uri": self.server.url() + MOCK_JWKS_URI,
            "status_list_endpoint": status_list_endpoint.to_string(),
        });
        let oidc = self
            .server
            .mock("GET", MOCK_OIDC_ENDPOINT)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(body.to_string())
            .expect(1)
            .create();
        self.endpoints.oidc = Some(oidc);
    }

    /// Rotates the active HS256 signing key and updates mocked OIDC/JWKS responses.
    pub(crate) fn rotate_signing_key_hs256(
        &mut self,
        kid: impl ToString,
    ) -> Result<(), KeyGenerationError> {
        self.keys = generate_keypair_hs256(Some(kid))?;

        // Replace JWKS endpoint with the new key set.
        let old_jwks = self.endpoints.jwks.take();
        drop(old_jwks);
        let jwks = self
            .server
            .mock("GET", MOCK_JWKS_URI)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(
                json!({"keys": generate_jwks(std::slice::from_ref(&self.keys)).keys}).to_string(),
            )
            .expect_at_least(1)
            .create();
        self.endpoints.jwks = Some(jwks);

        let old_oidc = self.endpoints.oidc.take();
        drop(old_oidc);

        let mut body = json!({
            "issuer": self.server.url(),
            "jwks_uri": self.server.url() + MOCK_JWKS_URI,
        });

        if let Some(status_list_endpoint) = self.status_list_endpoint() {
            body["status_list_endpoint"] = json!(status_list_endpoint.to_string());
        }

        let oidc = self
            .server
            .mock("GET", MOCK_OIDC_ENDPOINT)
            .with_status(200)
            .with_header("content-type", "application/json")
            .with_body(body.to_string())
            .expect_at_least(0)
            .create();
        self.endpoints.oidc = Some(oidc);

        Ok(())
    }

    /// Helper function for generating a status list JWT for this mock server
    pub(crate) async fn status_list_jwt(&self) -> Result<String, reqwest::Error> {
        static CLIENT: LazyLock<Client> = LazyLock::new(Client::new);
        let url = self.status_list_endpoint().expect("the status list endpoint hasn't been generated yet. call `generate_status_list_endpoint` first");
        CLIENT
            .get(url.to_string())
            .send()
            .await?
            .error_for_status()?
            .text()
            .await
    }

    pub(crate) fn status_list_endpoint(&self) -> Option<Url> {
        self.endpoints.status_list.as_ref()?;

        Some(
            Url::parse(&(self.server.url() + MOCK_STATUS_LIST_ENDPOINT))
                .expect("invalid status list url"),
        )
    }

    pub(crate) fn openid_config_endpoint(&self) -> Option<Url> {
        self.endpoints.oidc.as_ref()?;

        Some(
            Url::parse(&(self.server.url() + MOCK_OIDC_ENDPOINT)).expect("invalid status list url"),
        )
    }

    pub(crate) fn jwks_endpoint(&self) -> Option<Url> {
        self.endpoints.jwks.as_ref()?;

        Some(Url::parse(&(self.server.url() + MOCK_JWKS_URI)).expect("invalid status list url"))
    }

    pub(crate) fn jwt_decoding_key(&self) -> Result<Arc<DecodingKey>, jsonwebtoken::errors::Error> {
        self.keys.decoding_key()
    }

    pub(crate) fn jwt_decoding_key_and_id(
        &self,
    ) -> Result<(Arc<DecodingKey>, Option<String>), jsonwebtoken::errors::Error> {
        Ok((self.keys.decoding_key()?, self.keys.kid.clone()))
    }

    pub(crate) fn issuer(&self) -> IssClaim {
        IssClaim::new(&self.server.url())
    }

    #[track_caller]
    pub(super) fn openid_config(&self) -> OpenIdConfig {
        OpenIdConfig {
            issuer: self.issuer(),
            jwks_uri: self.jwks_endpoint().unwrap(),
            status_list_endpoint: self.status_list_endpoint(),
        }
    }

    /// Creates a mock server that will fail on OIDC configuration fetch.
    /// This is useful for testing error handling when trusted issuers fail to load.
    pub(crate) async fn new_with_failing_oidc() -> Result<Self, KeyGenerationError> {
        let mut server = Server::new_async().await;

        let keys = generate_keypair_hs256(Some("some_hs256_key"))?;

        // Create a mock that returns 500 Internal Server Error for OIDC config
        let oidc = Some(
            server
                .mock("GET", "/.well-known/openid-configuration")
                .with_status(500)
                .with_header("content-type", "application/json")
                .with_body(r#"{"error": "Internal Server Error"}"#)
                .expect(1)
                .create(),
        );

        // JWKS endpoint may still be created but won't be used since OIDC fails
        let jwks = Some(
            server
                .mock("GET", MOCK_JWKS_URI)
                .with_status(200)
                .with_header("content-type", "application/json")
                .with_body(
                    json!({"keys": generate_jwks(std::slice::from_ref(&keys)).keys}).to_string(),
                )
                .expect(0) // Expect 0 calls since OIDC will fail first
                .create(),
        );

        let endpoints = MockEndpoints {
            oidc,
            jwks,
            status_list: None,
        };

        Ok(Self {
            endpoints,
            server,
            keys,
        })
    }
}

/// Creates a trusted issuer with an invalid URL that will fail to fetch.
/// This is useful for testing error handling in issuer loading.
pub(crate) fn create_failing_trusted_issuer(issuer_id: &str) -> TrustedIssuer {
    let mut issuer = TrustedIssuer::default();
    issuer.set_oidc_endpoint(
        Url::parse(&format!(
            "invalid://{issuer_id}/.well-known/openid-configuration"
        ))
        .expect("should be a valid URL format"),
    );
    issuer
}

/// Creates a trusted issuer with a URL that points to a non-existent server.
/// This will fail with network connection errors.
pub(crate) fn create_unreachable_trusted_issuer(issuer_id: &str) -> TrustedIssuer {
    let mut issuer = TrustedIssuer::default();
    // Use a localhost port that's unlikely to be in use
    issuer.set_oidc_endpoint(
        Url::parse(&format!(
            "http://localhost:65535/{issuer_id}/.well-known/openid-configuration"
        ))
        .expect("should be a valid URL format"),
    );
    issuer
}