// https://syzkaller.appspot.com/bug?id=4a8f56e8783cb5f6d9f417c3b1ccf5d969250bd8
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <sys/mount.h>
#include <sys/syscall.h>
#include <pthread.h>
#include <errno.h>
#include <linux/falloc.h>
#include <time.h>

int force_empty_fd = -1;
int reclaim_fd = -1;
volatile int stop = 0;

void setup_cgroup() {
    char mount_path[256] = {0};
    FILE *fp = fopen("/proc/mounts", "r");
    if (!fp) {
        printf("[-] fopen /proc/mounts failed: %s\n", strerror(errno));
    } else {
        char line[1024];
        while (fgets(line, sizeof(line), fp)) {
            if (strstr(line, "cgroup") && strstr(line, "memory")) {
                sscanf(line, "%*s %255s", mount_path);
                break;
            }
        }
        fclose(fp);
    }

    if (mount_path[0] != '\0') {
        char cg_path[512];
        snprintf(cg_path, sizeof(cg_path), "%s/test_cg", mount_path);
        if (mkdir(cg_path, 0755) < 0 && errno != EEXIST) {
            printf("[-] mkdir %s failed: %s\n", cg_path, strerror(errno));
        } else {
            char procs_path[512];
            snprintf(procs_path, sizeof(procs_path), "%s/cgroup.procs", cg_path);
            int fd = open(procs_path, O_WRONLY);
            if (fd < 0) {
                snprintf(procs_path, sizeof(procs_path), "%s/tasks", cg_path);
                fd = open(procs_path, O_WRONLY);
            }
            if (fd >= 0) {
                char pid_str[32];
                snprintf(pid_str, sizeof(pid_str), "%d", getpid());
                if (write(fd, pid_str, strlen(pid_str)) > 0) {
                    close(fd);
                    char force_empty_path[512];
                    snprintf(force_empty_path, sizeof(force_empty_path), "%s/memory.force_empty", cg_path);
                    force_empty_fd = open(force_empty_path, O_WRONLY);
                    if (force_empty_fd >= 0) {
                        printf("[+] cgroup v1 setup successful.\n");
                        return;
                    } else {
                        printf("[-] open %s failed: %s\n", force_empty_path, strerror(errno));
                    }
                } else {
                    printf("[-] write to %s failed: %s\n", procs_path, strerror(errno));
                    close(fd);
                }
            } else {
                printf("[-] open %s failed: %s\n", procs_path, strerror(errno));
            }
        }
    }

    // Fallback to cgroup v2
    if (mkdir("/tmp/cgroup2", 0755) < 0 && errno != EEXIST) {
        printf("[-] mkdir /tmp/cgroup2 failed: %s\n", strerror(errno));
    }
    if (mount("cgroup2", "/tmp/cgroup2", "cgroup2", 0, NULL) == 0 || errno == EBUSY) {
        int fd = open("/tmp/cgroup2/cgroup.subtree_control", O_WRONLY);
        if (fd >= 0) {
            if (write(fd, "+memory", 7) < 0) {
                printf("[-] write +memory failed: %s\n", strerror(errno));
            }
            close(fd);
        } else {
            printf("[-] open subtree_control failed: %s\n", strerror(errno));
        }
        
        if (mkdir("/tmp/cgroup2/test_cg", 0755) < 0 && errno != EEXIST) {
            printf("[-] mkdir test_cg failed: %s\n", strerror(errno));
        }
        
        fd = open("/tmp/cgroup2/test_cg/cgroup.procs", O_WRONLY);
        if (fd >= 0) {
            char pid_str[32];
            snprintf(pid_str, sizeof(pid_str), "%d", getpid());
            if (write(fd, pid_str, strlen(pid_str)) > 0) {
                close(fd);
                reclaim_fd = open("/tmp/cgroup2/test_cg/memory.reclaim", O_WRONLY);
                if (reclaim_fd >= 0) {
                    printf("[+] cgroup v2 setup successful.\n");
                    return;
                } else {
                    printf("[-] open memory.reclaim failed: %s\n", strerror(errno));
                }
            } else {
                printf("[-] write to cgroup.procs failed: %s\n", strerror(errno));
                close(fd);
            }
        } else {
            printf("[-] open cgroup.procs failed: %s\n", strerror(errno));
        }
    } else {
        printf("[-] mount cgroup2 failed: %s\n", strerror(errno));
    }
    printf("[-] Failed to setup cgroup.\n");
    exit(1);
}

struct thread_data {
    int memfd;
};

void *fallocate_thread(void *arg) {
    struct thread_data *data = arg;
    int fd = data->memfd;
    while (!stop) {
        if (fallocate(fd, 0, 0, 1024 * 1024 * 50) < 0) {
            // Ignore ENOSPC
        }
        if (fallocate(fd, FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE, 0, 1024 * 1024 * 50) < 0) {
            // Ignore
        }
    }
    return NULL;
}

void *reclaim_thread(void *arg) {
    while (!stop) {
        if (force_empty_fd >= 0) {
            if (write(force_empty_fd, "0", 1) < 0) {
                // Ignore EINTR/EBUSY
            }
        } else if (reclaim_fd >= 0) {
            if (write(reclaim_fd, "50M", 3) < 0) {
                // Ignore
            }
        }
    }
    return NULL;
}

int main() {
    setvbuf(stdout, NULL, _IONBF, 0);
    printf("[*] Starting reproducer...\n");
    
    setup_cgroup();
    
    pthread_t t1[4], t2[4];
    struct thread_data data[4];
    
    for (int i = 0; i < 4; i++) {
        data[i].memfd = syscall(SYS_memfd_create, "test_shmem", 0);
        if (data[i].memfd < 0) {
            printf("[-] memfd_create failed: %s\n", strerror(errno));
            exit(1);
        }
        if (pthread_create(&t1[i], NULL, fallocate_thread, &data[i]) != 0) {
            printf("[-] pthread_create failed: %s\n", strerror(errno));
            exit(1);
        }
        if (pthread_create(&t2[i], NULL, reclaim_thread, NULL) != 0) {
            printf("[-] pthread_create failed: %s\n", strerror(errno));
            exit(1);
        }
    }
    
    printf("[+] Threads started. Waiting for 5 seconds...\n");
    sleep(5);
    stop = 1;
    
    for (int i = 0; i < 4; i++) {
        pthread_join(t1[i], NULL);
        pthread_join(t2[i], NULL);
        close(data[i].memfd);
    }
    
    if (force_empty_fd >= 0) close(force_empty_fd);
    if (reclaim_fd >= 0) close(reclaim_fd);
    
    printf("[+] Reproducer finished.\n");
    return 0;
}