#!/bin/bash

# Memory Corruption Detection Test Suite for C-Relay
# Tests for buffer overflows, use-after-free, and memory safety issues

set -e

# Configuration
RELAY_HOST="127.0.0.1"
RELAY_PORT="8888"
TEST_TIMEOUT=15

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Test counters
TOTAL_TESTS=0
PASSED_TESTS=0
FAILED_TESTS=0

# Function to test for memory corruption (buffer overflows, crashes, etc.)
test_memory_safety() {
    local description="$1"
    local message="$2"
    local expect_error="${3:-false}"

    TOTAL_TESTS=$((TOTAL_TESTS + 1))

    echo -n "Testing $description... "

    # Send message and monitor for crashes or memory issues
    local start_time=$(date +%s%N)
    local response
    response=$(echo "$message" | timeout $TEST_TIMEOUT websocat ws://$RELAY_HOST:$RELAY_PORT 2>/dev/null | head -1 || echo 'CONNECTION_FAILED')
    local end_time=$(date +%s%N)

    # Check if relay is still responsive after the test
    local relay_status
    local ping_response=$(echo '["REQ","ping_test_'$RANDOM'",{}]' | timeout 2 websocat ws://$RELAY_HOST:$RELAY_PORT 2>/dev/null | head -1)
    if [[ -n "$ping_response" ]]; then
        relay_status="OK"
    else
        relay_status="DOWN"
    fi

    # Calculate response time (rough indicator of processing issues)
    local response_time=$(( (end_time - start_time) / 1000000 )) # Convert to milliseconds

    if [[ "$response" == "CONNECTION_FAILED" ]]; then
        if [[ "$expect_error" == "true" ]]; then
            echo -e "${GREEN}PASSED${NC} - Expected connection failure"
            PASSED_TESTS=$((PASSED_TESTS + 1))
            return 0
        else
            echo -e "${RED}FAILED${NC} - Unexpected connection failure"
            FAILED_TESTS=$((FAILED_TESTS + 1))
            return 1
        fi
    elif [[ "$relay_status" != "OK" ]]; then
        echo -e "${RED}FAILED${NC} - Relay crashed or became unresponsive after test"
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    elif [[ $response_time -gt 5000 ]]; then # More than 5 seconds
        echo -e "${YELLOW}SUSPICIOUS${NC} - Very slow response (${response_time}ms), possible DoS"
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    else
        if [[ "$expect_error" == "true" ]]; then
            echo -e "${YELLOW}UNCERTAIN${NC} - Expected error but got normal response"
            PASSED_TESTS=$((PASSED_TESTS + 1)) # Count as passed since no crash
            return 0
        else
            echo -e "${GREEN}PASSED${NC} - No memory corruption detected"
            PASSED_TESTS=$((PASSED_TESTS + 1))
            return 0
        fi
    fi
}

# Function to test concurrent access patterns
test_concurrent_access() {
    local description="$1"
    local message="$2"
    local concurrent_count="${3:-5}"

    TOTAL_TESTS=$((TOTAL_TESTS + 1))

    echo -n "Testing $description... "

    # Launch multiple concurrent connections
    local pids=()
    local results=()

    for i in $(seq 1 $concurrent_count); do
        (
            local response
            response=$(echo "$message" | timeout $TEST_TIMEOUT websocat ws://$RELAY_HOST:$RELAY_PORT 2>/dev/null | head -1 || echo 'FAILED')
            echo "$response"
        ) &
        pids+=($!)
    done

    # Wait for all to complete
    local failed_count=0
    for pid in "${pids[@]}"; do
        wait "$pid" 2>/dev/null || failed_count=$((failed_count + 1))
    done

    # Check if relay is still responsive
    local relay_status
    local ping_response=$(echo '["REQ","ping_test_'$RANDOM'",{}]' | timeout 2 websocat ws://$RELAY_HOST:$RELAY_PORT 2>/dev/null | head -1)
    if [[ -n "$ping_response" ]]; then
        relay_status="OK"
    else
        relay_status="DOWN"
    fi

    if [[ "$relay_status" != "OK" ]]; then
        echo -e "${RED}FAILED${NC} - Relay crashed during concurrent access"
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    elif [[ $failed_count -gt 0 ]]; then
        echo -e "${YELLOW}PARTIAL${NC} - Some concurrent requests failed ($failed_count/$concurrent_count)"
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    else
        echo -e "${GREEN}PASSED${NC} - Concurrent access handled safely"
        PASSED_TESTS=$((PASSED_TESTS + 1))
        return 0
    fi
}

echo "=========================================="
echo "C-Relay Memory Corruption Test Suite"
echo "=========================================="
echo "Testing against relay at ws://$RELAY_HOST:$RELAY_PORT"
echo "Note: These tests may cause the relay to crash if vulnerabilities exist"
echo

# Test basic connectivity first
echo "=== Basic Connectivity Test ==="
test_memory_safety "Basic connectivity" '["REQ","basic_test",{}]'
echo

echo "=== Subscription ID Memory Corruption Tests ==="
# Test malformed subscription IDs that could cause buffer overflows
test_memory_safety "Empty subscription ID" '["REQ","",{}]' true
test_memory_safety "Very long subscription ID (1KB)" '["REQ","'$(printf 'a%.0s' {1..1024})'",{}]' true
test_memory_safety "Very long subscription ID (10KB)" '["REQ","'$(printf 'a%.0s' {1..10240})'",{}]' true
test_memory_safety "Subscription ID with null bytes" '["REQ","test\x00injection",{}]' true
test_memory_safety "Subscription ID with special chars" '["REQ","test@#$%^&*()",{}]' true
test_memory_safety "Unicode subscription ID" '["REQ","test🚀💣🔥",{}]' true
test_memory_safety "Subscription ID with path traversal" '["REQ","../../../etc/passwd",{}]' true
echo

echo "=== Filter Array Memory Corruption Tests ==="
# Test oversized filter arrays (limited to avoid extremely long output)
test_memory_safety "Too many filters (50)" '["REQ","test_many_filters",{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{} ,{}]' true
echo

echo "=== Concurrent Access Memory Tests ==="
# Test concurrent access patterns that might cause race conditions
test_concurrent_access "Concurrent subscription creation" '["REQ","concurrent_'$(date +%s%N)'",{}]' 10
test_concurrent_access "Concurrent CLOSE operations" '["CLOSE","test_sub"]' 10
echo

echo "=== Malformed JSON Memory Tests ==="
# Test malformed JSON that might cause parsing issues
test_memory_safety "Unclosed JSON object" '["REQ","test",{' true
test_memory_safety "Mismatched brackets" '["REQ","test"]' true
test_memory_safety "Extra closing brackets" '["REQ","test",{}]]' true
test_memory_safety "Null bytes in JSON" '["REQ","test\x00",{}]' true
echo

echo "=== Large Message Memory Tests ==="
# Test very large messages that might cause buffer issues
test_memory_safety "Very large filter array" '["REQ","large_test",{"authors":['$(printf '"test%.0s",' {1..1000})'"test"]}]' true
test_memory_safety "Very long search term" '["REQ","search_test",{"search":"'$(printf 'a%.0s' {1..10000})'"}]' true
echo

echo "=== Test Results ==="
echo "Total tests: $TOTAL_TESTS"
echo -e "Passed: ${GREEN}$PASSED_TESTS${NC}"
echo -e "Failed: ${RED}$FAILED_TESTS${NC}"

if [[ $FAILED_TESTS -eq 0 ]]; then
    echo -e "${GREEN}✓ All memory corruption tests passed!${NC}"
    echo "The relay appears to handle memory safely."
    exit 0
else
    echo -e "${RED}✗ Memory corruption vulnerabilities detected!${NC}"
    echo "The relay may be vulnerable to memory corruption attacks."
    echo "Failed tests: $FAILED_TESTS"
    exit 1
fi