perimeterx-fastly-enforcer 2.2.2

#[cfg(feature = "e2e_testing")]
use super::handlers::pxe2etesting;
use super::handlers::pxfilter::should_filter_request;
use super::handlers::pxlogger::{PX_LOG, PXLogger};
use super::handlers::pxtemplate;
use super::px_debug;
use crate::handlers::{
    pxactivities, pxagentic_trust, pxapi, pxcookievalidator, pxcors, pxcredentials_intelligence,
    pxfirstparty, pxgraphql, pxhd, pxtelemetry, pxuser_identifiers,
};
use crate::modules::pxconstants::*;
use crate::pxconfig::{
    PXAdditionalActivityHandlerFn, PXConfig, PXCorsCustomBlockResponseHeadersFn,
    PXCorsCustomPreflightHandlerFn, PXEnrichCustomParamsFn, PXIsEnforcedRequestFn,
    PXIsFilteredRequestFn, PXIsMonitoredRequestFn, PXIsSensitiveRequestFn,
};
use crate::pxcontext::{BlockReason, CallReason, PXContext, PXModuleMode, PassReason};
use fastly::http::{Method, StatusCode};
use fastly::{Body, Error, Request, Response};

/// Enforcer structure
#[derive(Default)]
pub struct PXEnforcer {
    pub(crate) conf: PXConfig,
    pub(crate) ctx: PXContext,
}

impl PXEnforcer {
    pub fn new(config_store_name: &str, secret_store_name: &str) -> Self {
        let conf = PXConfig::new(config_store_name, secret_store_name);

        // initialize logger
        if let Ok(mut log) = PX_LOG.lock() {
            *log = Some(PXLogger::new(&conf));
        }

        PXEnforcer {
            conf,
            ..Default::default()
        }
    }

    pub fn enforce(&mut self, req: &mut Request) -> Result<Option<Response>, Error> {
        // starting PerimeterX
        // verify module enabled
        if !self.conf.module_enabled {
            px_debug!("Request will not be verified, module is disabled.");
            return Ok(None);
        }

        if self.conf.cors_support_enabled {
            if let Some(resp) = pxcors::run_preflight_custom_handler(req, &self.conf) {
                return Ok(Some(resp));
            }

            if pxcors::should_filter_preflight_request(req, &self.conf) {
                return Ok(None);
            }
        }

        if let Some(is_filtered_request_fn) = &self.conf.is_filtered_request_fn {
            if is_filtered_request_fn(req, &self.conf) {
                px_debug!("Request will not be verified, custom filter matched.");
                return Ok(None);
            }
        }

        // get the app id without the "PX" prefix
        let px_fp_appid = match self.conf.app_id.get(..2) {
            Some(prefix) if prefix.eq_ignore_ascii_case("px") => {
                self.conf.app_id.get(2..).unwrap_or(&self.conf.app_id)
            }
            _ => &self.conf.app_id,
        };

        // check and handle first party. response with 200 code for HEAD requests
        match (req.get_method(), req.get_url().path()) {
            // sensor
            (&Method::GET, path) if path.starts_with(&format!("/{px_fp_appid}/init.js")) => {
                let mut req1 = req.clone_without_body();
                pxfirstparty::req_redirect_client(&mut req1, &self.conf);
                return Ok(Some(req1.send(&self.conf.human_client_backend)?));
            }
            (&Method::HEAD, path) if path.starts_with(&format!("/{px_fp_appid}/init.js")) => {
                let mut resp = Response::new();
                resp.set_status(StatusCode::OK);
                return Ok(Some(resp));
            }

            // captcha
            (&Method::GET, path) if path.starts_with(&format!("/{px_fp_appid}/captcha")) => {
                let mut req1 = req.clone_without_body();
                pxfirstparty::req_redirect_captcha(&mut req1, &self.conf);
                return Ok(Some(req1.send(&self.conf.human_captcha_backend)?));
            }
            (&Method::HEAD, path) if path.starts_with(&format!("/{px_fp_appid}/captcha")) => {
                let mut resp = Response::new();
                resp.set_status(StatusCode::OK);
                return Ok(Some(resp));
            }

            // XHR
            (&Method::POST | &Method::GET, path)
                if path.starts_with(&format!("/{px_fp_appid}/xhr")) =>
            {
                let mut req1 = req.clone_with_body();
                pxfirstparty::req_redirect_xhr(&mut req1, &self.conf);
                return Ok(Some(req1.send(&self.conf.human_collector_backend)?));
            }
            (&Method::HEAD, path) if path.starts_with(&format!("/{px_fp_appid}/xhr")) => {
                let mut resp = Response::new();
                resp.set_status(StatusCode::OK);
                return Ok(Some(resp));
            }

            // handle e2e endpoints (if enabled)
            #[cfg(feature = "e2e_testing")]
            (&Method::POST, path) if path.eq("/pushdata") => {
                return pxe2etesting::handle_pushdata(req, &mut self.conf);
            }
            _ => {}
        }

        // handle telemetry
        if req.contains_header(TELEMETRY_HEADER) {
            px_debug!("found telemetry header");
            pxtelemetry::handle_telemetry(req, &self.conf, &self.ctx);
            return Ok(None);
        }

        px_debug!(
            "{} starting request verification for {}",
            PX_MODULE_VERSION,
            req.get_path()
        );

        // check request filtering
        if should_filter_request(req, &self.conf) {
            return Ok(None);
        }

        // create context
        self.ctx = PXContext::new(req, &self.conf);
        self.ctx.extract_pxde_cookie(&self.conf);
        px_debug!("Request context created successfully");

        // handle graphql
        if self.conf.graphql_enabled
            && req.get_method() == Method::POST
            && req
                .get_header_str_lossy("content-type")
                .map(|v| v.into_owned())
                .unwrap_or_default()
                .contains(APPLICATION_JSON)
        {
            // check if URL matches any of the configured GraphQL patterns
            for pattern in &self.conf.graphql_routes {
                if pattern.is_match(req.get_path().trim()) {
                    if req.get_content_length().unwrap_or(0) > self.conf.graphql_body_max_length {
                        px_debug!(
                            "GraphQL request body exceeds maximum configured length, skipping GraphQL processing"
                        );
                        break;
                    }
                    // Read body bytes up to the maximum configured length and attempt to parse as JSON
                    let body = req.get_body_prefix_mut(self.conf.graphql_body_max_length);

                    match serde_json::from_slice::<serde_json::Value>(body.as_slice()) {
                        Ok(v) => {
                            self.ctx.graphql_extracted_items =
                                pxgraphql::px_graphql_extract(&v, &self.conf);
                        }
                        Err(e) => {
                            px_debug!("Failed to parse JSON GraphQL body: {}", e);
                        }
                    }
                    break;
                }
            }

            if pxgraphql::px_graphql_is_sensitive(&self.ctx.graphql_extracted_items) {
                self.ctx.is_sensitive_route = true;
            }
        }

        pxagentic_trust::enrich_context_from_request(req, &self.conf, &mut self.ctx);
        pxuser_identifiers::enrich_context_from_request(req, &self.conf, &mut self.ctx);
        pxcredentials_intelligence::enrich_context_from_request(req, &self.conf, &mut self.ctx);
        if pxagentic_trust::should_postpone_activities(&self.ctx) {
            self.ctx.postpone_activities = true;
        }

        if let Some(is_sensitive_request_fn) = &self.conf.is_sensitive_request_fn {
            self.ctx.is_sensitive_route =
                is_sensitive_request_fn(req, &self.conf) || self.ctx.is_sensitive_route;
        }

        // if token validation failed, verify request with Risk API
        // otherwise, verify request with cookie validator
        let should_verify_cookie = self.ctx.s2s_call_reason.as_ref().is_none()
            || self
                .ctx
                .s2s_call_reason
                .as_ref()
                .is_some_and(|r| r.is_mobile_sdk_error());

        let is_cookie_verified = if should_verify_cookie {
            pxcookievalidator::verify_cookie(&mut self.ctx, &self.conf)
        } else {
            false
        };

        if let Some(custom_params_fn) = &self.conf.enrich_params_fn {
            custom_params_fn(req, &self.conf, &mut self.ctx.custom_params);
        }

        if !is_cookie_verified
            && self
                .ctx
                .s2s_call_reason
                .as_ref()
                .is_some_and(|r| r != &CallReason::None)
        {
            pxapi::verify_request(&mut self.ctx, &self.conf);
        }

        if self.ctx.telemetry_requested {
            let telemetry_result = pxtelemetry::send_telemetry_activity(
                &self.conf,
                &self.ctx,
                pxtelemetry::TelemetryUpdateReason::RiskApi,
            );
            match telemetry_result {
                Ok(response) => {
                    px_debug!("telemetry status: {}", response.get_status());
                }
                Err(e) => {
                    px_debug!("error sending telemetry: {}", e);
                }
            }
        }

        self.handle_verification(req)
    }

    fn send_async_activity(&self, req: Option<&Request>) {
        if self
            .ctx
            .block_reason
            .as_ref()
            .is_some_and(|r| r != &BlockReason::None)
        {
            px_debug!("request blocked, sending block activity");
            pxactivities::send_block_activity(&self.ctx, &self.conf);
        } else {
            px_debug!("request is valid, sending page_requested activity");
            pxactivities::send_page_requested_activity(&self.ctx, &self.conf);
        };

        if let Some(req) = req {
            if let Some(additional_activity_handler_fn) = &self.conf.additional_activity_handler_fn
            {
                additional_activity_handler_fn(req, &self.conf, &self.ctx);
            }
        }
    }

    // when a request should be passed:
    // - if the module mode is active blocking, and the request contains a valid cookie or RiskAPI with a low score
    // - if the module mode is monitor
    // - if the module mode is active blocking, and the request is marked as monitored (custom callback or monitor route)
    fn should_pass_as_monitored(&self, _req: &Request) -> bool {
        // When the module is in active blocking mode or the request is marked as enforced (custom callback or enforced route)
        // pass ONLY if the request is passed by cookie or RiskAPI or S2s error
        let pass_request = if self.ctx.risk_mode == PXModuleMode::Blocking {
            self.ctx.pass_reason.as_ref().is_some_and(|r| {
                matches!(r, PassReason::Cookie | PassReason::S2s | PassReason::Error)
            })
        } else {
            true
        };

        // if the module mode is monitor
        let is_monitor_mode = self.ctx.risk_mode == PXModuleMode::Monitor;

        self.ctx.is_monitored_request || is_monitor_mode || pass_request
    }

    /// Mark Block activities as simulated when the Enforcer passes the request through
    /// instead of returning a block response (monitor mode, or monitored routes in active blocking).
    pub(crate) fn apply_simulated_block_flag(&mut self) {
        if !self.ctx.is_enforced_request
            && (self.ctx.risk_mode == PXModuleMode::Monitor || self.ctx.is_monitored_request)
        {
            self.ctx.is_simulated_block = true;
        }
    }

    fn handle_verification(&mut self, req: &mut Request) -> Result<Option<Response>, Error> {
        if self.ctx.pxde_verified && !self.conf.data_enrichment_header_name.is_empty() {
            req.set_header(
                &self.conf.data_enrichment_header_name,
                self.ctx.pxde.clone().unwrap_or_default(),
            );
        }
        pxcredentials_intelligence::modify_incoming_request(req, &self.conf, &self.ctx);

        let pass_request = self.should_pass_as_monitored(req);

        self.apply_simulated_block_flag();

        // send async activity at once if the request is blocked
        if !pass_request {
            // do not postpone activities
            self.ctx.postpone_activities = false;
            self.send_async_activity(Some(req));

        // request is not blocked, postpone activities if needed
        } else {
            // send async activity if the request is not blocked, postpone activities if needed
            if !self.ctx.postpone_activities {
                self.send_async_activity(Some(req));
            } else {
                px_debug!("Postponing async activities until post_enforce");
            }
        }

        if pass_request {
            return Ok(None);
        };

        let block_response = pxtemplate::get_block_response(&self.conf, &self.ctx);
        let mut response = Response::new()
            .with_status(block_response.status_code)
            .with_header("cache-control", "private, no-store")
            .with_header("content-type", block_response.content_type)
            .with_body(Body::from(block_response.body));

        if pxcors::should_add_cors_headers(req, &self.conf) {
            pxcors::apply_cors_headers_to_block_response(req, &self.conf, &mut response);
        }

        pxhd::apply_pxhd_cookie(&mut response, &self.ctx, &self.conf);

        Ok(Some(response))
    }

    /// Must be called after `px.enforce()` to perform response processing.
    pub fn post_enforce(&mut self, resp: &mut Response) {
        if self.ctx.postpone_activities {
            // try to enrich context from response
            pxagentic_trust::enrich_context_from_response(resp, &self.conf, &mut self.ctx);
            self.send_async_activity(None);
        }

        if self.ctx.credential_intelligence.is_some() && self.conf.additional_s2s_activity_enabled {
            pxcredentials_intelligence::enrich_context_from_response(
                resp,
                &self.conf,
                &mut self.ctx,
            );
            pxcredentials_intelligence::send_additional_s2s(&self.ctx, &self.conf);
        }

        pxhd::apply_pxhd_cookie(resp, &self.ctx, &self.conf);
    }

    // Public API

    /// set a callback function to enrich custom parameters, must be called before `px.enforce()`
    pub fn set_enrich_custom_params_fn(&mut self, f: PXEnrichCustomParamsFn) {
        self.conf.enrich_params_fn = Some(f);
    }

    /// set a callback function to identify "sensitive" requests, must be called before `px.enforce()`
    pub fn set_is_sensitive_request_fn(&mut self, f: PXIsSensitiveRequestFn) {
        self.conf.is_sensitive_request_fn = Some(f);
    }

    /// set a callback function to identify "enforced" requests, must be called before `px.enforce()`
    pub fn set_is_enforced_request_fn(&mut self, f: PXIsEnforcedRequestFn) {
        self.conf.is_enforced_request_fn = Some(f);
    }

    /// set a callback function to identify "monitored" requests, must be called before `px.enforce()`
    pub fn set_is_monitored_request_fn(&mut self, f: PXIsMonitoredRequestFn) {
        self.conf.is_monitored_request_fn = Some(f);
    }

    /// set a callback function to filter out requests that should not be verified by Enforcer, must be called before `px.enforce()`
    pub fn set_is_filtered_request_fn(&mut self, f: PXIsFilteredRequestFn) {
        self.conf.is_filtered_request_fn = Some(f);
    }

    pub fn set_additional_activity_handler_fn(&mut self, f: PXAdditionalActivityHandlerFn) {
        self.conf.additional_activity_handler_fn = Some(f);
    }

    /// Register custom credential extraction callback for CI endpoints with `sent_through: custom`.
    pub fn set_ci_extract_credentials_fn(&mut self, f: crate::pxconfig::PXExtractCredentialsFn) {
        self.conf.set_ci_extract_credentials_fn(f);
    }

    /// Register custom login-success callback for CI endpoints with `login_successful_reporting_method: custom`.
    pub fn set_ci_login_successful_fn(&mut self, f: crate::pxconfig::PXLoginSuccessfulFn) {
        self.conf.set_ci_login_successful_fn(f);
    }

    /// set a callback function to handle CORS preflight requests, must be called before `px.enforce()`
    pub fn set_cors_custom_preflight_handler_fn(&mut self, f: PXCorsCustomPreflightHandlerFn) {
        self.conf.cors_custom_preflight_handler_fn = Some(f);
    }

    /// set a callback function to create custom CORS headers for block responses, must be called before `px.enforce()`
    pub fn set_cors_create_custom_block_response_headers_fn(
        &mut self,
        f: PXCorsCustomBlockResponseHeadersFn,
    ) {
        self.conf.cors_create_custom_block_response_headers_fn = Some(f);
    }

    /// set module mode (Monitor/Blocking), must be called before `px.enforce()`
    pub fn set_module_mode(&mut self, mode: PXModuleMode) {
        self.conf.module_mode = mode;
    }

    /// Return the read-only [`PXContext`] for the current request.
    pub fn ctx(&self) -> &PXContext {
        &self.ctx
    }

    /// return PXConfig read-only structure
    pub fn conf(&self) -> &PXConfig {
        &self.conf
    }
}