codex-memory 3.0.15

use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::process::Command;

/// Test that the MCP server is protected against buffer exhaustion attacks
#[tokio::test]
async fn test_buffer_exhaustion_protection() {
    // Create a large JSON payload (exceeding 10MB buffer limit)
    let large_payload_size = 11 * 1024 * 1024; // 11MB
    let mut large_payload = String::with_capacity(large_payload_size + 1000);
    large_payload.push_str(r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"store_memory","arguments":{"content":""#);
    
    // Fill with 'A' characters to exceed buffer limit
    for _ in 0..large_payload_size {
        large_payload.push('A');
    }
    large_payload.push_str(r#"","context":"test","summary":"test","tags":["test"]}}}"#);

    // Start MCP server
    let mut child = Command::new("cargo")
        .args(&["run", "--", "mcp-stdio"])
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .expect("Failed to start MCP server");

    let stdin = child.stdin.as_mut().expect("Failed to get stdin");
    let mut stdout = child.stdout.take().expect("Failed to get stdout");

    // Send the oversized payload
    stdin.write_all(large_payload.as_bytes()).await.ok();
    stdin.write_all(b"\n").await.ok();

    // Read response with timeout
    let mut response = Vec::new();
    let result = tokio::time::timeout(Duration::from_secs(5), async {
        let mut buffer = [0u8; 1024];
        while let Ok(n) = stdout.read(&mut buffer).await {
            if n == 0 { break; }
            response.extend_from_slice(&buffer[..n]);
            if response.len() > 2048 { break; } // Prevent test from hanging
        }
    }).await;

    // Verify that server sends error response instead of crashing
    assert!(result.is_ok(), "Server should respond within timeout");
    let response_str = String::from_utf8_lossy(&response);
    
    // Should contain parse error or buffer limit error
    assert!(
        response_str.contains("Parse error") || response_str.contains("Buffer size limit exceeded"),
        "Server should return buffer limit error, got: {}",
        response_str
    );

    // Clean up
    child.kill().await.ok();
}

/// Test protection against deeply nested JSON stack overflow attacks
#[tokio::test]
async fn test_json_stack_overflow_protection() {
    // Create deeply nested JSON (over 100 levels to trigger limit)
    let mut nested_json = String::new();
    nested_json.push_str(r#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"store_memory","arguments":{"content":"#);
    
    // Create 150 levels of nesting (should exceed 100 level limit)
    for _ in 0..150 {
        nested_json.push_str(r#"{"nested":"#);
    }
    nested_json.push_str("\"value\"");
    for _ in 0..150 {
        nested_json.push('}');
    }
    nested_json.push_str(r#"","context":"test","summary":"test","tags":["test"]}}}"#);

    // Start MCP server
    let mut child = Command::new("cargo")
        .args(&["run", "--", "mcp-stdio"])
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .expect("Failed to start MCP server");

    let stdin = child.stdin.as_mut().expect("Failed to get stdin");
    let mut stdout = child.stdout.take().expect("Failed to get stdout");

    // Send the deeply nested JSON
    stdin.write_all(nested_json.as_bytes()).await.ok();
    stdin.write_all(b"\n").await.ok();

    // Read response with timeout
    let mut response = Vec::new();
    let result = tokio::time::timeout(Duration::from_secs(5), async {
        let mut buffer = [0u8; 1024];
        while let Ok(n) = stdout.read(&mut buffer).await {
            if n == 0 { break; }
            response.extend_from_slice(&buffer[..n]);
            if response.len() > 2048 { break; }
        }
    }).await;

    assert!(result.is_ok(), "Server should respond within timeout");
    let response_str = String::from_utf8_lossy(&response);
    
    // Should contain nesting depth error
    assert!(
        response_str.contains("nesting depth exceeded") || response_str.contains("Parse error"),
        "Server should return nesting depth error, got: {}",
        response_str
    );

    // Clean up
    child.kill().await.ok();
}

/// Test protection against malformed UTF-8 attacks
#[tokio::test]
async fn test_utf8_validation_protection() {
    // Create malformed UTF-8 sequence
    let mut malformed_json = Vec::new();
    malformed_json.extend_from_slice(br#"{"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"name":"store_memory","arguments":{"content":""#);
    
    // Insert invalid UTF-8 bytes
    malformed_json.push(0xFF); // Invalid UTF-8 start byte
    malformed_json.push(0xFE); // Invalid UTF-8 start byte
    malformed_json.push(0xFD); // Invalid UTF-8 start byte
    
    malformed_json.extend_from_slice(br#"","context":"test","summary":"test","tags":["test"]}}}"#);

    // Start MCP server
    let mut child = Command::new("cargo")
        .args(&["run", "--", "mcp-stdio"])
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .expect("Failed to start MCP server");

    let stdin = child.stdin.as_mut().expect("Failed to get stdin");
    let mut stdout = child.stdout.take().expect("Failed to get stdout");

    // Send the malformed UTF-8
    stdin.write_all(&malformed_json).await.ok();
    stdin.write_all(b"\n").await.ok();

    // Read response with timeout
    let mut response = Vec::new();
    let result = tokio::time::timeout(Duration::from_secs(5), async {
        let mut buffer = [0u8; 1024];
        while let Ok(n) = stdout.read(&mut buffer).await {
            if n == 0 { break; }
            response.extend_from_slice(&buffer[..n]);
            if response.len() > 2048 { break; }
        }
    }).await;

    assert!(result.is_ok(), "Server should respond within timeout");
    let response_str = String::from_utf8_lossy(&response);
    
    // Should contain UTF-8 encoding error
    assert!(
        response_str.contains("Invalid UTF-8") || response_str.contains("Parse error"),
        "Server should return UTF-8 validation error, got: {}",
        response_str
    );

    // Clean up
    child.kill().await.ok();
}

/// Test that normal valid JSON still works after security fixes
#[tokio::test]
async fn test_normal_json_still_works() {
    let valid_json = r#"{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}"#;

    // Start MCP server
    let mut child = Command::new("cargo")
        .args(&["run", "--", "mcp-stdio"])
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .expect("Failed to start MCP server");

    let stdin = child.stdin.as_mut().expect("Failed to get stdin");
    let mut stdout = child.stdout.take().expect("Failed to get stdout");

    // Send valid JSON
    stdin.write_all(valid_json.as_bytes()).await.ok();
    stdin.write_all(b"\n").await.ok();

    // Read response
    let mut response = Vec::new();
    let result = tokio::time::timeout(Duration::from_secs(5), async {
        let mut buffer = [0u8; 1024];
        while let Ok(n) = stdout.read(&mut buffer).await {
            if n == 0 { break; }
            response.extend_from_slice(&buffer[..n]);
            if response.len() > 2048 { break; }
        }
    }).await;

    assert!(result.is_ok(), "Server should respond to valid JSON");
    let response_str = String::from_utf8_lossy(&response);
    
    // Should contain successful response
    assert!(
        response_str.contains("protocolVersion") && response_str.contains("capabilities"),
        "Server should return successful initialize response, got: {}",
        response_str
    );

    // Clean up
    child.kill().await.ok();
}

/// Test request timeout protection
#[tokio::test]
async fn test_request_timeout_protection() {
    // This test verifies that the server properly handles timeouts
    // We'll send a valid request and verify it responds within reasonable time
    
    let valid_json = r#"{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}"#;

    let mut child = Command::new("cargo")
        .args(&["run", "--", "mcp-stdio"])
        .stdin(std::process::Stdio::piped())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .expect("Failed to start MCP server");

    let stdin = child.stdin.as_mut().expect("Failed to get stdin");
    let mut stdout = child.stdout.take().expect("Failed to get stdout");

    stdin.write_all(valid_json.as_bytes()).await.ok();
    stdin.write_all(b"\n").await.ok();

    // Verify server responds within 5 seconds (well under 60s timeout)
    let mut response = Vec::new();
    let result = tokio::time::timeout(Duration::from_secs(5), async {
        let mut buffer = [0u8; 1024];
        while let Ok(n) = stdout.read(&mut buffer).await {
            if n == 0 { break; }
            response.extend_from_slice(&buffer[..n]);
            if response.len() > 100 { break; } // Got some response
        }
    }).await;

    assert!(result.is_ok(), "Server should respond within timeout");
    let response_str = String::from_utf8_lossy(&response);
    assert!(!response_str.is_empty(), "Server should return some response");

    child.kill().await.ok();
}