pdk_contracts_lib/api/

validator.rs

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

use lock_lib::{Lock, LockBuilder, TryLock};
use pdk_core::classy::extract::context::ConfigureContext;
use pdk_core::classy::extract::{Extract as _, FromContext};
use pdk_core::classy::hl::HttpClientError;
use pdk_core::classy::proxy_wasm::types::Status;
use pdk_core::classy::{Clock, SharedData};
use pdk_core::logger;
use pdk_core::policy_context::api::Metadata;
use std::rc::Rc;
use std::time::Duration;
use thiserror::Error;

use super::error::UpdateError;
use super::error::{AuthenticationError, AuthorizationError};
use crate::api::ClientData;
use crate::implementation::constants::{
    ACCEPT_HASH_ALGORITHM_VALUE, CONTRACTS_FETCHING_INTERVAL, LOCK_TIMEOUT_DURATION,
};
use crate::implementation::model::contracts_storage::{ContractsLocalStorage, ContractsStorage};
use crate::implementation::model::session_storage::{SessionSharedDataStorage, SessionStorage};
use crate::implementation::platform::client::HttpPlatformClient;
use crate::implementation::platform::responses::{
    contract_from_event, ContractsResponse, LoginResponse,
};
use crate::implementation::platform::shared::{AccessToken, ContractsRequestParams};

use super::authentication::authenticate;
use super::authorization::authorize;
use super::credentials::{ClientId, ClientSecret};

/// The object that will collect the contracts and provide the functionality to validate incoming
/// requests.
pub struct ContractValidator {
    api_id: String,
    client: HttpPlatformClient,
    session_lock: TryLock,
    api_lock: TryLock,
    session_storage: SessionSharedDataStorage,
    contract_storage: ContractsLocalStorage,
    clock: Rc<dyn Clock>,
}

#[derive(Error, Debug)]
enum InternalUpdateError {
    #[error("Http client error: {0}")]
    HttpClientError(#[from] HttpClientError),
    #[error("Parsing error: {0}")]
    Serde(#[from] serde_json::error::Error),
    #[error("Lost the lock while executing async function.")]
    LostLock,
    #[error("Upstream returned unexpected status code {0}")]
    UnexpectedResponse(u32),
}

enum PollContractsResponse {
    Continue,
    Renegotiate,
    Finish,
}

impl ContractValidator {
    fn new(
        client: HttpPlatformClient,
        api_id: String,
        clock: Rc<dyn Clock>,
        shared_data: Rc<dyn SharedData>,
        lock_builder: LockBuilder,
    ) -> Self {
        // The session storage and the session lock only use two entries in total despite the amount
        // of Apis/policies deployed. We can safely not remove them without fear of a leak.
        let session_storage =
            SessionSharedDataStorage::new(Rc::clone(&clock), Rc::clone(&shared_data));

        let session_lock = lock_builder
            .new(session_storage.session_lock_key())
            .expiration(LOCK_TIMEOUT_DURATION)
            .shared()
            .build();

        let contract_storage = ContractsLocalStorage::new(&api_id, Rc::clone(&clock), shared_data);
        let api_lock = lock_builder
            .new(contract_storage.api_lock_key())
            .expiration(LOCK_TIMEOUT_DURATION)
            .shared()
            .build();

        Self {
            api_id,
            client,
            session_lock,
            api_lock,
            session_storage,
            contract_storage,
            clock,
        }
    }

    /// Initial updating interval for calling [ContractValidator::update_contracts()].
    pub const INITIALIZATION_PERIOD: Duration = Duration::from_millis(100);

    /// Remaining updating interval for calling [ContractValidator::update_contracts()].
    pub const UPDATE_PERIOD: Duration = CONTRACTS_FETCHING_INTERVAL;

    /// Validates if `client_id` credential can be authorized in the current API contract.
    /// If the validation is succesful, this method returns the [ClientData] of `client_id`.
    /// Otherwise an [AuthorizationError] is returned.
    pub fn authorize(&self, client_id: &ClientId) -> Result<ClientData, AuthorizationError> {
        authorize(&self.contract_storage, client_id)
    }

    /// Validates if `client_id` and `client_secret` credentials can be authenticated
    /// in the current API contract.
    /// If the validation is succesful, this method returns the [ClientData] of `client_id`.
    /// Otherwise an [AuthenticationError] is returned.
    pub fn authenticate(
        &self,
        client_id: &ClientId,
        client_secret: &ClientSecret,
    ) -> Result<ClientData, AuthenticationError> {
        authenticate(&self.contract_storage, client_id, client_secret)
    }

    /// Returns `true` if contracts have been pulled and are available locally.
    /// This method can be used to determine if the [`ContractValidator`] is ready
    /// for authorization and authentication operations.
    pub fn is_ready(&self) -> bool {
        self.contract_storage.last_update().is_some()
    }

    /// Updates local contracts database.
    /// This method is intended to be called periodically in order to keep
    /// the local contracts database up to date.
    /// During initialization time, this method should be invoked in a period of
    /// [ContractValidator::INITIALIZATION_PERIOD].
    /// After initialization, this method should be invoked in a period of
    /// [ContractValidator::UPDATE_PERIOD].
    pub async fn update_contracts(&self) -> Result<(), UpdateError> {
        if !self.should_update() {
            logger::debug!("Contracts update skipped since update period hasn't elapsed.");
            return Ok(());
        }

        logger::debug!("Fetching contracts for API {}", self.api_id);

        let Some(token) = self.session_token().await else {
            return Ok(());
        };

        let Some(api_lock) = self.api_lock.try_lock() else {
            logger::debug!(
                "Other worker has the lock for API {}. Skipping update.",
                self.api_id
            );
            return Ok(());
        };

        let mut token_data = token;
        loop {
            match self.poll_contracts(&token_data, &api_lock).await {
                Ok(PollContractsResponse::Continue) => {
                    logger::debug!("Contract polling request successful. Chaining next request");
                }
                Ok(PollContractsResponse::Renegotiate) => {
                    let Some(token) = self.renegotiate_token(token_data).await else {
                        // Could not renegotiate the token
                        return Ok(());
                    };
                    token_data = token;
                    if !api_lock.refresh_lock() {
                        logger::debug!("Lost the api lock while refreshing the session lock.");
                        return Ok(());
                    }
                }
                Ok(PollContractsResponse::Finish) => {
                    return Ok(());
                }
                Err(error) => {
                    logger::debug!("Error while polling contracts: {error}");
                    return Ok(());
                }
            }
        }
    }

    fn should_update(&self) -> bool {
        self.contract_storage
            .last_update()
            .map(|last| last + CONTRACTS_FETCHING_INTERVAL < self.clock.get_current_time())
            .unwrap_or(true)
    }

    async fn session_token(&self) -> Option<AccessToken> {
        match self.session_storage.get_token() {
            Some(token) => Some(token),
            None => {
                let Some(session_lock) = self.session_lock.try_lock() else {
                    logger::debug!("Other worker has the session lock. Skipping update.");
                    return None;
                };

                // Re check if the token was already set before obtaining the lock.
                // Since the lock is "not blocking" the changes are quite slim, but not null.
                if let Some(token) = self.session_storage.get_token() {
                    return Some(token);
                }

                self.fetch_session_token(&session_lock).await
            }
        }
    }

    async fn renegotiate_token(&self, old_token: AccessToken) -> Option<AccessToken> {
        // Validate that the token was not renegotiated by someone else.
        if let Some(token) = self.session_storage.get_token() {
            if token != old_token {
                return Some(token);
            }
        };

        // Acquire the lock
        let Some(session_lock) = self.session_lock.try_lock() else {
            logger::debug!("Other worker has the session lock. Aborting token renegotiation");
            return None;
        };

        // Validate Again that the token was not renegotiated by someone else.
        // Since the lock is "not blocking" the changes are quite slim, but not null.
        if let Some(token) = self.session_storage.get_token() {
            if token != old_token {
                return Some(token);
            }
        };

        self.fetch_session_token(&session_lock).await
    }

    async fn fetch_session_token(&self, session_lock: &'_ Lock<'_>) -> Option<AccessToken> {
        match self.perform_login_request().await {
            Ok(login) => {
                if !session_lock.refresh_lock() {
                    logger::debug!("Lost the session lock. Aborting update.");
                    return None;
                }
                let token = login.get_token();
                let token_data = AccessToken::new(token.to_string(), login.get_type().to_string());
                logger::debug!("Obtained the session token.");
                self.session_storage.save_token(token_data.clone());
                Some(token_data)
            }
            Err(e) => {
                logger::warn!(
                    "Unexpected error while performing login request {e}. Skipping update."
                );
                None
            }
        }
    }

    async fn perform_login_request(&self) -> Result<LoginResponse, InternalUpdateError> {
        logger::debug!("Getting platform token...");
        match self.client.login().await? {
            r if r.status_code() == 200 => Ok(serde_json::from_slice::<LoginResponse>(r.body())?),
            r => {
                logger::debug!(
                    "Fetching contracts failed with status code: {} and body:\n {}",
                    r.status_code(),
                    String::from_utf8_lossy(r.body())
                );
                Err(InternalUpdateError::UnexpectedResponse(r.status_code()))
            }
        }
    }

    async fn poll_contracts(
        &self,
        access_token: &AccessToken,
        api_lock: &'_ Lock<'_>,
    ) -> Result<PollContractsResponse, InternalUpdateError> {
        let token = access_token.get_access_token();
        let response = self
            .client
            .contracts(
                token,
                self.api_id.as_str(),
                ACCEPT_HASH_ALGORITHM_VALUE,
                self.next_url(),
            )
            .await?;

        if !api_lock.refresh_lock() {
            return Err(InternalUpdateError::LostLock);
        }

        match response.status_code() {
            200 => {
                let contracts: ContractsResponse = serde_json::from_slice(response.body())
                    .map_err(|_| HttpClientError::Status(Status::InternalFailure))?;

                if self.no_updates(&contracts) {
                    self.finish_polling();
                    Ok(PollContractsResponse::Finish)
                } else {
                    self.log_invalid_contracts(&contracts);
                    self.update_data(&contracts);
                    self.update_links(&contracts);
                    Ok(PollContractsResponse::Continue)
                }
            }
            401 => Ok(PollContractsResponse::Renegotiate),
            n => {
                logger::debug!(
                    "Fetching contracts failed with status code: {} and body:\n {}",
                    n,
                    String::from_utf8_lossy(response.body())
                );
                Err(InternalUpdateError::UnexpectedResponse(n))
            }
        }
    }

    fn log_invalid_contracts(&self, response: &ContractsResponse) {
        for invalid_contract_error_msg in response.verify_contracts().err().unwrap_or_default() {
            logger::warn!("{invalid_contract_error_msg}")
        }
    }

    fn no_updates(&self, response: &ContractsResponse) -> bool {
        let links = response.get_links();
        links.self_link() == links.next_link()
    }

    fn finish_polling(&self) {
        self.contract_storage.update_last();
        logger::debug!(
            "No more contracts updates for API {}, polling in next tick.",
            self.api_id
        );
    }

    fn update_data(&self, response: &ContractsResponse) {
        let data = response.get_data();
        for contract_event in data {
            match contract_event.removed.unwrap_or(false) {
                true => self
                    .contract_storage
                    .remove_contract(&contract_event.client_id),
                false => self
                    .contract_storage
                    .save_contract(contract_from_event(contract_event)),
            }
        }
        logger::debug!(
            "{} contract events processed for API {}",
            data.len(),
            self.api_id
        );
    }

    fn update_links(&self, response: &ContractsResponse) {
        let links = response.get_links();
        let params = ContractsRequestParams::new(
            Some(links.next_link().to_string()),
            ACCEPT_HASH_ALGORITHM_VALUE.to_string(),
        );
        self.contract_storage.save_contracts_request_params(params);
    }

    fn next_url(&self) -> Option<String> {
        self.contract_storage
            .get_contracts_request_params()
            .and_then(|x| x.next_url)
    }
}

#[derive(thiserror::Error, Debug)]
pub enum ExtractionError {
    #[error("Api metadata is unavailable.")]
    ApiMetadata,

    #[error("Environment Context is unavailable.")]
    EnvironmentContext,

    #[error("Anypoint Context is unavailable.")]
    AnypointContext,
}

impl FromContext<ConfigureContext> for ContractValidator {
    type Error = ExtractionError;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        let metadata: Metadata = context.extract_always();
        let api_id = metadata
            .api_metadata
            .id
            .ok_or(ExtractionError::ApiMetadata)?;
        let client = context.extract()?;
        let clock = context.extract_always();
        let shared_data = context.extract_always();
        let lock_builder = context.extract_always();

        Ok(Self::new(client, api_id, clock, shared_data, lock_builder))
    }
}