pmcp-code-mode 0.5.0

Code Mode validation and execution framework for MCP servers
Documentation
//! Property-based tests for pmcp-code-mode security invariants.
//!
//! Covers:
//! - HMAC round-trip integrity (generate token -> verify -> same code passes)
//! - Code modification detection (modified code fails verification)
//! - Bit-flipped token rejection (tampered tokens fail verification)
//! - Wrong-secret rejection (token from one secret fails with another)
//! - Hash determinism (same input -> same hash)
//! - Canonicalize idempotency (canonicalize(canonicalize(x)) == canonicalize(x))
//! - Validation pipeline determinism (same GraphQL input -> same result)
//! - Negative trait assertions for TokenSecret (no Debug, Clone, Serialize)
//! - Default-deny policy test

use pmcp_code_mode::{
    canonicalize_code, hash_code, ApprovalToken, CodeModeConfig, GraphQLValidator,
    HmacTokenGenerator, NoopPolicyEvaluator, PolicyEvaluator, RiskLevel, TokenGenerator,
    TokenSecret, ValidationContext, ValidationPipeline,
};
use proptest::prelude::*;
use std::sync::Arc;

// ---------------------------------------------------------------------------
// HMAC Round-Trip Tests
// ---------------------------------------------------------------------------

proptest! {
    /// For any secret and any code string, a token generated by
    /// HmacTokenGenerator for that code can be verified successfully
    /// against the same code.
    #[test]
    fn hmac_roundtrip_any_code(
        secret in prop::collection::vec(any::<u8>(), 16..64),
        code in "[ -~]{1,500}"
    ) {
        let token_secret = TokenSecret::new(secret);
        let generator = HmacTokenGenerator::new(token_secret).unwrap();

        let token = generator.generate(
            &code,
            "user-prop",
            "session-prop",
            "server-prop",
            "context-prop",
            RiskLevel::Low,
            300,
        );

        // Signature and expiry verification must succeed
        prop_assert!(generator.verify(&token).is_ok(),
            "verify() failed for generated token");

        // Code hash verification must succeed
        prop_assert!(generator.verify_code(&code, &token).is_ok(),
            "verify_code() failed for original code");
    }

    /// Modifying the code after token generation causes verification failure.
    #[test]
    fn hmac_detects_code_modification(
        secret in prop::collection::vec(any::<u8>(), 16..64),
        code in "[ -~]{1,200}",
        suffix in "[ -~]{1,50}"
    ) {
        let token_secret = TokenSecret::new(secret);
        let generator = HmacTokenGenerator::new(token_secret).unwrap();

        let token = generator.generate(
            &code,
            "user-prop",
            "session-prop",
            "server-prop",
            "context-prop",
            RiskLevel::Low,
            300,
        );

        let modified_code = format!("{}{}", code, suffix);

        // Only assert failure when the modification actually changes the
        // canonicalized hash (appending whitespace-only suffixes may not).
        if hash_code(&code) != hash_code(&modified_code) {
            prop_assert!(generator.verify_code(&modified_code, &token).is_err(),
                "verify_code() should fail for modified code");
        }
    }

    /// hash_code is deterministic: same input -> same hash.
    #[test]
    fn hash_code_deterministic(code in "[ -~]{0,1000}") {
        let hash1 = hash_code(&code);
        let hash2 = hash_code(&code);
        prop_assert_eq!(hash1, hash2);
    }

    /// canonicalize_code is deterministic and idempotent.
    #[test]
    fn canonicalize_idempotent(code in "[ -~]{0,500}") {
        let c1 = canonicalize_code(&code);
        let c2 = canonicalize_code(&c1);
        prop_assert_eq!(c1, c2);
    }
}

// ---------------------------------------------------------------------------
// Negative Token Tests (bit-flipping and wrong-secret)
// ---------------------------------------------------------------------------

proptest! {
    /// Bit-flipping a valid token must cause verification failure.
    /// This tests resistance to token forgery via bitwise manipulation.
    #[test]
    fn hmac_rejects_bitflipped_tokens(
        secret in prop::collection::vec(any::<u8>(), 32..64),
        code in "[ -~]{10,200}",
        flip_position in 0usize..200,
    ) {
        let token_secret = TokenSecret::new(secret);
        let generator = HmacTokenGenerator::new(token_secret).unwrap();

        let token = generator.generate(
            &code,
            "user-prop",
            "session-prop",
            "server-prop",
            "context-prop",
            RiskLevel::Low,
            300,
        );

        // Encode the token, flip a bit, then decode and verify
        let encoded = token.encode().unwrap();
        let mut bytes = encoded.into_bytes();

        if !bytes.is_empty() {
            let idx = flip_position % bytes.len();
            bytes[idx] ^= 0x01; // Flip one bit in one byte
            let flipped_str = String::from_utf8_lossy(&bytes);

            // Try to decode the flipped token -- it may fail to decode
            // (invalid base64/JSON) or decode but fail verification.
            // Either outcome is acceptable; it must NOT verify successfully.
            if let Ok(flipped_token) = ApprovalToken::decode(&flipped_str) {
                // If decode succeeded, verify must fail (tampered signature)
                let verify_result = generator.verify(&flipped_token);
                let code_result = generator.verify_code(&code, &flipped_token);
                // At least one of these checks must fail
                prop_assert!(
                    verify_result.is_err() || code_result.is_err(),
                    "Bit-flipped token should not pass both verify and verify_code"
                );
            }
            // If decode fails, that's fine -- corrupted token rejected
        }
    }

    /// Wrong secret must fail verification even with correct token format.
    #[test]
    fn hmac_rejects_wrong_secret(
        secret1 in prop::collection::vec(any::<u8>(), 32..64),
        secret2 in prop::collection::vec(any::<u8>(), 32..64),
        code in "[ -~]{10,200}",
    ) {
        prop_assume!(secret1 != secret2);

        let gen1 = HmacTokenGenerator::new(TokenSecret::new(secret1)).unwrap();
        let gen2 = HmacTokenGenerator::new(TokenSecret::new(secret2)).unwrap();

        // Generate token with gen1
        let token = gen1.generate(
            &code,
            "user-prop",
            "session-prop",
            "server-prop",
            "context-prop",
            RiskLevel::Low,
            300,
        );

        // Verify with gen2 -- must fail (different secret)
        prop_assert!(gen2.verify(&token).is_err(),
            "Token generated with one secret should not verify with a different secret");
    }
}

// ---------------------------------------------------------------------------
// Validation Pipeline Determinism
// ---------------------------------------------------------------------------

proptest! {
    /// Same GraphQL input to the validator produces the same result.
    #[test]
    fn graphql_validation_deterministic(
        field_name in "[a-z]{1,15}"
    ) {
        let query = format!("query {{ {} {{ id }} }}", field_name);
        let validator = GraphQLValidator::default();

        let result1 = validator.validate(&query);
        let result2 = validator.validate(&query);

        match (result1, result2) {
            (Ok(info1), Ok(info2)) => {
                // Operation type must be identical
                prop_assert_eq!(info1.operation_type, info2.operation_type);
                // Root fields must be identical
                prop_assert_eq!(info1.root_fields, info2.root_fields);
                // Max depth must be identical
                prop_assert_eq!(info1.max_depth, info2.max_depth);
            }
            (Err(_), Err(_)) => {
                // Both errored -- deterministic
            }
            _ => {
                prop_assert!(false, "Non-deterministic: one call succeeded and the other failed");
            }
        }
    }
}

// ---------------------------------------------------------------------------
// TokenSecret Negative Trait Assertions
// ---------------------------------------------------------------------------

/// Trait denial tests for TokenSecret.
///
/// TokenSecret intentionally does NOT implement Debug, Display, Clone,
/// Serialize, or Deserialize. The primary enforcement is:
/// 1. The absence of derives/impls in token.rs (source-level)
/// 2. trybuild compile-fail tests in Plan 04 (compile-time)
///
/// These runtime tests provide a complementary canary: they verify that
/// the type exists, is move-only, and doesn't accidentally gain traits
/// through transitive derives or blanket impls.
#[cfg(test)]
mod token_secret_trait_denials {
    use super::*;

    /// TokenSecret can be constructed and consumed (move semantics).
    /// This test is a canary: if TokenSecret gains Clone, it would
    /// still pass, but the trybuild compile-fail test catches that.
    #[test]
    fn token_secret_is_consumed_on_move() {
        let secret = TokenSecret::new(b"test".to_vec());
        let _moved = secret;
        // Uncommenting the next line would fail to compile
        // because TokenSecret is NOT Clone:
        // let _copy = secret;  // ERROR: use of moved value
    }

    /// TokenSecret is NOT Debug -- verified by checking that we CANNOT
    /// format it. This is enforced at compile time: `format!("{:?}", secret)`
    /// would fail to compile. We verify the intent here by constructing
    /// a TokenSecret and confirming it exists without Debug formatting.
    #[test]
    fn token_secret_not_debug_canary() {
        let secret = TokenSecret::new(b"canary-secret".to_vec());
        // If TokenSecret implemented Debug, we could do:
        //   let _ = format!("{:?}", secret);
        // The fact that this file compiles WITHOUT that line is the assertion.
        // We just verify the type exists and can be created.
        let _ = secret;
    }

    /// TokenSecret is NOT Display -- same reasoning as Debug above.
    #[test]
    fn token_secret_not_display_canary() {
        let secret = TokenSecret::new(b"canary-secret".to_vec());
        // If TokenSecret implemented Display, we could do:
        //   let _ = format!("{}", secret);
        // The fact that this file compiles WITHOUT that line is the assertion.
        let _ = secret;
    }

    /// TokenSecret is NOT Serialize -- verified by checking that
    /// serde_json::to_value fails at compile time. This canary ensures
    /// the type can't accidentally be serialized to JSON.
    #[test]
    fn token_secret_not_serialize_canary() {
        let _secret = TokenSecret::new(b"canary-secret".to_vec());
        // If TokenSecret implemented Serialize, we could do:
        //   let _ = serde_json::to_value(&secret);
        // The fact that this file compiles WITHOUT that line is the assertion.
    }

    /// TokenSecret exposes secret bytes only via expose_secret().
    /// This verifies the API exists and returns the correct bytes.
    #[test]
    fn token_secret_expose_secret_returns_correct_bytes() {
        let input = b"test-bytes-123";
        let secret = TokenSecret::new(input.to_vec());
        assert_eq!(secret.expose_secret(), input);
    }
}

// ---------------------------------------------------------------------------
// Default-Deny Policy Test
// ---------------------------------------------------------------------------

/// A policy evaluator that denies all operations.
struct DenyAllEvaluator;

#[pmcp_code_mode::async_trait]
impl PolicyEvaluator for DenyAllEvaluator {
    async fn evaluate_operation(
        &self,
        _operation: &pmcp_code_mode::OperationEntity,
        _server_config: &pmcp_code_mode::ServerConfigEntity,
    ) -> Result<pmcp_code_mode::AuthorizationDecision, pmcp_code_mode::PolicyEvaluationError> {
        Ok(pmcp_code_mode::AuthorizationDecision {
            allowed: false,
            determining_policies: vec!["deny_all_test".to_string()],
            errors: vec![],
        })
    }

    fn name(&self) -> &str {
        "deny_all_test"
    }
}

/// With a DenyAllEvaluator, the async validation pipeline rejects
/// valid GraphQL queries. This proves the default-deny property: without
/// NoopPolicyEvaluator, operations are denied.
#[tokio::test]
async fn default_deny_without_noop() {
    let config = CodeModeConfig::enabled();
    let pipeline = ValidationPipeline::with_policy_evaluator(
        config,
        b"test-secret-key!".to_vec(),
        Arc::new(DenyAllEvaluator),
    )
    .unwrap();

    let ctx = ValidationContext::new("user-123", "session-456", "schema-hash", "perms-hash");

    let result = pipeline
        .validate_graphql_query_async("query { users { id name } }", &ctx)
        .await
        .unwrap();

    // DenyAllEvaluator returned allowed: false, so result must be invalid
    assert!(
        !result.is_valid,
        "Query should be rejected when DenyAllEvaluator is configured"
    );
    assert!(
        !result.violations.is_empty(),
        "There should be policy violations from the deny-all evaluator"
    );
}

/// With NoopPolicyEvaluator, the same query is allowed. This contrast
/// proves NoopPolicyEvaluator is the explicit opt-in to bypass policy.
#[tokio::test]
async fn noop_evaluator_allows_query() {
    let config = CodeModeConfig::enabled();
    let pipeline = ValidationPipeline::with_policy_evaluator(
        config,
        b"test-secret-key!".to_vec(),
        Arc::new(NoopPolicyEvaluator::new()),
    )
    .unwrap();

    let ctx = ValidationContext::new("user-123", "session-456", "schema-hash", "perms-hash");

    let result = pipeline
        .validate_graphql_query_async("query { users { id name } }", &ctx)
        .await
        .unwrap();

    assert!(
        result.is_valid,
        "Query should be allowed when NoopPolicyEvaluator is configured"
    );
}