perimeterx-fastly-enforcer 2.2.2

use std::collections::HashMap;

use aes::cipher::{BlockModeEncrypt, KeyIvInit, block_padding::Pkcs7};
use base64::{Engine as _, engine::general_purpose};
use cbc::Encryptor;
use pbkdf2::pbkdf2_hmac_array;
use serde_json::json;
use sha2::Sha256;

use crate::handlers::pxcookievalidator::verify_cookie;
use crate::handlers::pxcrypto;
use crate::modules::pxconstants::{IV_LEN, KEY_LEN};
use crate::pxconfig::PXConfig;
use crate::pxcontext::{
    CallReason, CookieOrigin, CookieVersion, PXContext, PassReason, TokenVersion,
};

type Aes256CbcEnc = Encryptor<aes::Aes256>;

fn mobile_error_ctx(s2s_call_reason: CallReason, original_token: Option<String>) -> PXContext {
    PXContext {
        cookie_origin: Some(CookieOrigin::Header),
        original_token,
        s2s_call_reason: Some(s2s_call_reason),
        ..Default::default()
    }
}

fn test_conf(cookie_secret: &str) -> PXConfig {
    PXConfig {
        cookie_secret: cookie_secret.to_owned(),
        ..Default::default()
    }
}

fn test_conf_v2(cookie_secret: &str) -> PXConfig {
    PXConfig {
        cookie_secret: cookie_secret.to_owned(),
        token_version: TokenVersion::V2,
        ..Default::default()
    }
}

fn build_valid_original_token_v2(cookie_secret: &str) -> String {
    let mut token_json = json!({
        "t": 1_900_000_000_000_i64,
        "u": "550e8400-e29b-41d4-a716-446655440000",
        "v": "550e8400-e29b-41d4-a716-446655440001",
    });

    let Some(hmac) = pxcrypto::get_cookie_hmac(&token_json, "0", cookie_secret) else {
        return String::new();
    };
    if let Some(obj) = token_json.as_object_mut() {
        obj.insert("h".to_owned(), json!(hmac));
    }

    general_purpose::STANDARD.encode(token_json.to_string())
}

fn encrypt_cookie_v3_payload(
    cookie_secret: &str,
    salt: &[u8],
    iterations: u32,
    plaintext: &[u8],
) -> Option<Vec<u8>> {
    const LEN: usize = KEY_LEN + IV_LEN;
    let decoded_key = pbkdf2_hmac_array::<Sha256, LEN>(cookie_secret.as_bytes(), salt, iterations);
    let key: &[u8; KEY_LEN] = decoded_key.get(0..KEY_LEN)?.try_into().ok()?;
    let iv: &[u8; IV_LEN] = decoded_key.get(KEY_LEN..)?.try_into().ok()?;

    let mut buf = plaintext.to_vec();
    let msg_len = buf.len();
    buf.resize(msg_len + IV_LEN, 0);
    let cipher = Aes256CbcEnc::new(key.into(), iv.into());
    let ciphertext = cipher.encrypt_padded::<Pkcs7>(&mut buf, msg_len).ok()?;
    Some(ciphertext.to_vec())
}

fn build_valid_original_token_v3(cookie_secret: &str) -> String {
    let salt = vec![7_u8, 8, 9, 10, 11, 12, 13, 14];
    let iterations = 500_u32;
    let cookie_json = json!({
        "t": 1_900_000_000_000_i64,
        "u": "550e8400-e29b-41d4-a716-446655440000",
        "v": "550e8400-e29b-41d4-a716-446655440001",
        "s": 0,
        "a": "c",
        "x": "",
    });
    let plaintext = cookie_json.to_string();
    let encrypted =
        encrypt_cookie_v3_payload(cookie_secret, &salt, iterations, plaintext.as_bytes())
            .unwrap_or_default();
    if encrypted.is_empty() {
        return String::new();
    }

    let salt_b64 = general_purpose::STANDARD.encode(&salt);
    let payload_b64 = general_purpose::STANDARD.encode(&encrypted);
    let cookie_to_sign = format!("{salt_b64}:{iterations}:{payload_b64}");
    let signing_fields = String::new();
    let ehmac = pxcrypto::create_hmac(&(cookie_to_sign.clone() + &signing_fields), cookie_secret)
        .unwrap_or_default();

    format!("{ehmac}:{cookie_to_sign}")
}

#[test]
fn mobile_error_validates_original_token_and_preserves_reason() {
    let cookie_secret = "secret";
    let original_token = build_valid_original_token_v2(cookie_secret);
    assert!(!original_token.is_empty());

    let mut ctx = mobile_error_ctx(
        CallReason::MobileError1,
        Some(format!("2:{original_token}")),
    );
    let conf = test_conf(cookie_secret);

    assert!(!verify_cookie(&mut ctx, &conf));
    assert_eq!(ctx.s2s_call_reason, Some(CallReason::MobileError1));
    assert_eq!(ctx.original_token_error, None);
    assert!(ctx.cookie_version.is_none());
}

#[test]
fn mobile_error_invalid_v2_original_token_sets_validation_failed() {
    let cookie_secret = "secret";
    let mut original_token = build_valid_original_token_v2(cookie_secret);
    assert!(!original_token.is_empty());
    original_token.push('x');

    let mut ctx = mobile_error_ctx(
        CallReason::MobileError1,
        Some(format!("2:{original_token}")),
    );
    let conf = test_conf(cookie_secret);

    assert!(!verify_cookie(&mut ctx, &conf));
    assert_eq!(ctx.s2s_call_reason, Some(CallReason::MobileError1));
    assert_eq!(
        ctx.original_token_error,
        Some(CallReason::CookieValidationFailed)
    );
}

#[test]
fn mobile_error_valid_v3_original_token_preserves_reason() {
    let cookie_secret = "secret";
    let original_token = build_valid_original_token_v3(cookie_secret);
    assert!(!original_token.is_empty());

    let mut ctx = mobile_error_ctx(
        CallReason::MobileError1,
        Some(format!("3:{original_token}")),
    );
    let conf = test_conf(cookie_secret);

    assert!(!verify_cookie(&mut ctx, &conf));
    assert_eq!(ctx.s2s_call_reason, Some(CallReason::MobileError1));
    assert_eq!(ctx.original_token_error, None);
    assert!(ctx.cookie_version.is_none());
}

#[test]
fn mobile_error_invalid_v3_original_token_sets_decryption_failed() {
    let cookie_secret = "secret";
    let mut original_token = build_valid_original_token_v3(cookie_secret);
    assert!(!original_token.is_empty());
    original_token.push('x');

    let mut ctx = mobile_error_ctx(
        CallReason::MobileError1,
        Some(format!("3:{original_token}")),
    );
    let conf = test_conf(cookie_secret);

    assert!(!verify_cookie(&mut ctx, &conf));
    assert_eq!(ctx.s2s_call_reason, Some(CallReason::MobileError1));
    assert_eq!(
        ctx.original_token_error,
        Some(CallReason::CookieDecryptionFailed)
    );
}

#[test]
fn mobile_error_invalid_v3_original_token_hmac_sets_validation_failed() {
    let cookie_secret = "secret";
    let original_token = build_valid_original_token_v3(cookie_secret);
    assert!(!original_token.is_empty());
    let tampered = if let Some((_, rest)) = original_token.split_once(':') {
        format!("deadbeef:{rest}")
    } else {
        original_token
    };

    let mut ctx = mobile_error_ctx(CallReason::MobileError1, Some(format!("3:{tampered}")));
    let conf = test_conf(cookie_secret);

    assert!(!verify_cookie(&mut ctx, &conf));
    assert_eq!(ctx.s2s_call_reason, Some(CallReason::MobileError1));
    assert_eq!(
        ctx.original_token_error,
        Some(CallReason::CookieValidationFailed)
    );
}

#[test]
fn v3_config_validates_v3_cookie_when_present() {
    let cookie_secret = "secret";
    let px3 = build_valid_original_token_v3(cookie_secret);
    assert!(!px3.is_empty());

    let mut ctx = PXContext {
        cookies: HashMap::from([("_px3".to_owned(), px3)]),
        ..Default::default()
    };
    let conf = test_conf(cookie_secret);

    assert!(verify_cookie(&mut ctx, &conf));
    assert!(ctx.cookie_version == Some(CookieVersion::V3));
    assert!(ctx.pass_reason == Some(PassReason::Cookie));
}

#[test]
fn v2_sensitive_route_without_cookie_uses_no_cookie_reason() {
    let mut ctx = PXContext {
        is_sensitive_route: true,
        ..Default::default()
    };
    let conf = test_conf_v2("secret");

    assert!(!verify_cookie(&mut ctx, &conf));
    assert_eq!(ctx.s2s_call_reason, Some(CallReason::NoCookie));
}

#[test]
fn v2_sensitive_route_without_cookie_with_pxhd_vid_uses_no_cookie_w_vid() {
    let mut ctx = PXContext {
        is_sensitive_route: true,
        pxhd_cookie: Some("some-vid".to_owned()),
        ..Default::default()
    };
    let conf = test_conf_v2("secret");

    assert!(!verify_cookie(&mut ctx, &conf));
    assert_eq!(ctx.s2s_call_reason, Some(CallReason::NoCookieWVid));
}

#[test]
fn v3_sensitive_route_without_cookie_uses_no_cookie_reason() {
    let mut ctx = PXContext {
        is_sensitive_route: true,
        ..Default::default()
    };
    let conf = test_conf("secret");

    assert!(!verify_cookie(&mut ctx, &conf));
    assert_eq!(ctx.s2s_call_reason, Some(CallReason::NoCookie));
}

#[test]
fn v3_config_rejects_v2_when_v3_absent() {
    let cookie_secret = "secret";
    let px2 = build_valid_original_token_v2(cookie_secret);
    assert!(!px2.is_empty());

    let mut ctx = PXContext {
        cookies: HashMap::from([("_px2".to_owned(), px2)]),
        ..Default::default()
    };
    let conf = test_conf(cookie_secret);

    assert!(!verify_cookie(&mut ctx, &conf));
    assert_eq!(ctx.s2s_call_reason, Some(CallReason::NoCookie));
    assert!(ctx.pass_reason != Some(PassReason::Cookie));
}

#[test]
fn v3_config_does_not_fall_back_to_v2_when_v3_present_and_invalid() {
    let cookie_secret = "secret";
    let px2 = build_valid_original_token_v2(cookie_secret);
    let mut invalid_px3 = build_valid_original_token_v3(cookie_secret);
    assert!(!px2.is_empty());
    assert!(!invalid_px3.is_empty());
    invalid_px3.push('x');

    let mut ctx = PXContext {
        cookies: HashMap::from([("_px3".to_owned(), invalid_px3), ("_px2".to_owned(), px2)]),
        ..Default::default()
    };
    let conf = test_conf(cookie_secret);

    assert!(!verify_cookie(&mut ctx, &conf));
    assert!(ctx.cookie_version == Some(CookieVersion::V3));
    assert!(ctx.pass_reason != Some(PassReason::Cookie));
    assert_eq!(
        ctx.s2s_call_reason,
        Some(CallReason::CookieDecryptionFailed)
    );
}