// https://syzkaller.appspot.com/bug?id=301fa557ddbf731875e95df827f770495590ce32
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <linux/nbd.h>
#include <linux/netlink.h>
#include <linux/genetlink.h>
#include <linux/uhid.h>
#include <stdint.h>
#include <sys/wait.h>
#include <errno.h>
#include <signal.h>

#define NLA_ALIGN(len) (((len) + 3) & ~3)
#define NLA_HDRLEN ((int)NLA_ALIGN(sizeof(struct nlattr)))
#define NLA_F_NESTED (1 << 15)
#define NLA_TYPE_MASK 0x3fff

int get_mcast_group_id(int fd, const char *family_name, const char *group_name) {
    struct {
        struct nlmsghdr nlh;
        struct genlmsghdr gnlh;
        struct nlattr nla;
        char name[32];
    } req = {
        .nlh.nlmsg_len = NLMSG_HDRLEN + GENL_HDRLEN + NLA_HDRLEN + strlen(family_name) + 1,
        .nlh.nlmsg_type = GENL_ID_CTRL,
        .nlh.nlmsg_flags = NLM_F_REQUEST,
        .gnlh.cmd = CTRL_CMD_GETFAMILY,
        .nla.nla_type = 2, // CTRL_ATTR_FAMILY_NAME
        .nla.nla_len = NLA_HDRLEN + strlen(family_name) + 1,
    };
    strcpy(req.name, family_name);
    req.nlh.nlmsg_len = NLA_ALIGN(req.nlh.nlmsg_len);
    
    int res = send(fd, &req, req.nlh.nlmsg_len, 0);
    if (res < 0) {
        printf("[-] Failed to send get_mcast_group_id: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] send get_mcast_group_id successful.\n");

    char buf[4096];
    int len = recv(fd, buf, sizeof(buf), 0);
    if (len < 0) {
        printf("[-] Failed to recv get_mcast_group_id: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] recv get_mcast_group_id successful.\n");

    struct nlmsghdr *nlh = (struct nlmsghdr *)buf;
    if (nlh->nlmsg_type == NLMSG_ERROR) {
        printf("[-] NLMSG_ERROR in get_mcast_group_id\n");
        return -1;
    }

    struct genlmsghdr *gnlh = (struct genlmsghdr *)((char *)nlh + NLMSG_HDRLEN);
    struct nlattr *nla = (struct nlattr *)((char *)gnlh + GENL_HDRLEN);
    int rem = nlh->nlmsg_len - NLMSG_HDRLEN - GENL_HDRLEN;

    while (rem >= NLA_HDRLEN && rem >= nla->nla_len) {
        if ((nla->nla_type & NLA_TYPE_MASK) == 7) { // CTRL_ATTR_MCAST_GROUPS
            struct nlattr *mcast = (struct nlattr *)((char *)nla + NLA_HDRLEN);
            int mcast_rem = nla->nla_len - NLA_HDRLEN;
            while (mcast_rem >= NLA_HDRLEN && mcast_rem >= mcast->nla_len) {
                struct nlattr *attr = (struct nlattr *)((char *)mcast + NLA_HDRLEN);
                int attr_rem = mcast->nla_len - NLA_HDRLEN;
                int id = -1;
                char name[32] = {0};
                while (attr_rem >= NLA_HDRLEN && attr_rem >= attr->nla_len) {
                    if ((attr->nla_type & NLA_TYPE_MASK) == 1) { // CTRL_ATTR_MCAST_GRP_NAME
                        strncpy(name, (char *)attr + NLA_HDRLEN, sizeof(name) - 1);
                    } else if ((attr->nla_type & NLA_TYPE_MASK) == 2) { // CTRL_ATTR_MCAST_GRP_ID
                        id = *(uint32_t *)((char *)attr + NLA_HDRLEN);
                    }
                    attr_rem -= NLA_ALIGN(attr->nla_len);
                    attr = (struct nlattr *)((char *)attr + NLA_ALIGN(attr->nla_len));
                }
                if (id != -1 && strcmp(name, group_name) == 0) {
                    return id;
                }
                mcast_rem -= NLA_ALIGN(mcast->nla_len);
                mcast = (struct nlattr *)((char *)mcast + NLA_ALIGN(mcast->nla_len));
            }
        }
        rem -= NLA_ALIGN(nla->nla_len);
        nla = (struct nlattr *)((char *)nla + NLA_ALIGN(nla->nla_len));
    }
    return -1;
}

int get_family_id(int fd, const char *name) {
    struct {
        struct nlmsghdr nlh;
        struct genlmsghdr gnlh;
        struct nlattr nla;
        char name[32];
    } req = {
        .nlh.nlmsg_len = NLMSG_HDRLEN + GENL_HDRLEN + NLA_HDRLEN + strlen(name) + 1,
        .nlh.nlmsg_type = GENL_ID_CTRL,
        .nlh.nlmsg_flags = NLM_F_REQUEST,
        .gnlh.cmd = CTRL_CMD_GETFAMILY,
        .nla.nla_type = 2, // CTRL_ATTR_FAMILY_NAME
        .nla.nla_len = NLA_HDRLEN + strlen(name) + 1,
    };
    strcpy(req.name, name);
    req.nlh.nlmsg_len = NLA_ALIGN(req.nlh.nlmsg_len);
    
    int res = send(fd, &req, req.nlh.nlmsg_len, 0);
    if (res < 0) {
        printf("[-] Failed to send get_family_id: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] send get_family_id successful.\n");

    char buf[4096];
    int len = recv(fd, buf, sizeof(buf), 0);
    if (len < 0) {
        printf("[-] Failed to recv get_family_id: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] recv get_family_id successful.\n");

    struct nlmsghdr *nlh = (struct nlmsghdr *)buf;
    if (nlh->nlmsg_type == NLMSG_ERROR) {
        printf("[-] NLMSG_ERROR in get_family_id\n");
        return -1;
    }

    struct genlmsghdr *gnlh = (struct genlmsghdr *)((char *)nlh + NLMSG_HDRLEN);
    struct nlattr *nla = (struct nlattr *)((char *)gnlh + GENL_HDRLEN);
    int rem = nlh->nlmsg_len - NLMSG_HDRLEN - GENL_HDRLEN;

    while (rem >= NLA_HDRLEN && rem >= nla->nla_len) {
        if ((nla->nla_type & NLA_TYPE_MASK) == 1) { // CTRL_ATTR_FAMILY_ID
            return *(uint16_t *)((char *)nla + NLA_HDRLEN);
        }
        rem -= NLA_ALIGN(nla->nla_len);
        nla = (struct nlattr *)((char *)nla + NLA_ALIGN(nla->nla_len));
    }
    return -1;
}

void nbd_netlink_cmd(int nl_fd, int family_id, int cmd, int nbd_idx, int *sock_fds, int num_socks) {
    char buf[2048] = {0};
    struct nlmsghdr *nlh = (struct nlmsghdr *)buf;
    struct genlmsghdr *gnlh = (struct genlmsghdr *)(buf + NLMSG_HDRLEN);
    
    nlh->nlmsg_len = NLMSG_HDRLEN + GENL_HDRLEN;
    nlh->nlmsg_type = family_id;
    nlh->nlmsg_flags = NLM_F_REQUEST;
    gnlh->cmd = cmd; // 1 = CONNECT, 3 = RECONFIGURE
    
    struct nlattr *nla_idx = (struct nlattr *)(buf + nlh->nlmsg_len);
    nla_idx->nla_type = 1; // NBD_ATTR_INDEX
    nla_idx->nla_len = NLA_HDRLEN + 4;
    *(uint32_t *)((char *)nla_idx + NLA_HDRLEN) = nbd_idx;
    nlh->nlmsg_len += NLA_ALIGN(nla_idx->nla_len);
    
    if (cmd == 1) {
        struct nlattr *nla_sz = (struct nlattr *)(buf + nlh->nlmsg_len);
        nla_sz->nla_type = 2; // NBD_ATTR_SIZE_BYTES
        nla_sz->nla_len = NLA_HDRLEN + 8;
        *(uint64_t *)((char *)nla_sz + NLA_HDRLEN) = 1024 * 1024;
        nlh->nlmsg_len += NLA_ALIGN(nla_sz->nla_len);

        struct nlattr *nla_flags = (struct nlattr *)(buf + nlh->nlmsg_len);
        nla_flags->nla_type = 5; // NBD_ATTR_SERVER_FLAGS
        nla_flags->nla_len = NLA_HDRLEN + 8;
        *(uint64_t *)((char *)nla_flags + NLA_HDRLEN) = 256; // NBD_FLAG_CAN_MULTI_CONN
        nlh->nlmsg_len += NLA_ALIGN(nla_flags->nla_len);
    }
    
    struct nlattr *nla_socks = (struct nlattr *)(buf + nlh->nlmsg_len);
    nla_socks->nla_type = 7 | NLA_F_NESTED; // NBD_ATTR_SOCKETS
    nla_socks->nla_len = NLA_HDRLEN;
    
    for (int i = 0; i < num_socks; i++) {
        struct nlattr *nla_item = (struct nlattr *)((char *)nla_socks + nla_socks->nla_len);
        nla_item->nla_type = 1 | NLA_F_NESTED; // NBD_SOCK_ITEM
        nla_item->nla_len = NLA_HDRLEN;
        
        struct nlattr *nla_fd = (struct nlattr *)((char *)nla_item + nla_item->nla_len);
        nla_fd->nla_type = 1; // NBD_SOCK_FD
        nla_fd->nla_len = NLA_HDRLEN + 4;
        *(uint32_t *)((char *)nla_fd + NLA_HDRLEN) = sock_fds[i];
        
        nla_item->nla_len += NLA_ALIGN(nla_fd->nla_len);
        nla_socks->nla_len += NLA_ALIGN(nla_item->nla_len);
    }
    nlh->nlmsg_len += NLA_ALIGN(nla_socks->nla_len);
    
    int res = send(nl_fd, buf, nlh->nlmsg_len, 0);
    if (res < 0) {
        printf("[-] Failed to send NBD netlink command: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] send NBD netlink command %d successful.\n", cmd);
}

int create_shield() {
    int fd = open("/dev/uhid", O_RDWR);
    if (fd < 0) {
        printf("[-] Failed to open /dev/uhid: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] open /dev/uhid successful.\n");

    struct uhid_event ev = {0};
    ev.type = 0; // UHID_CREATE
    strcpy((char*)ev.u.create.name, "shield");
    ev.u.create.rd_data = (uint8_t *)"\x05\x01\x09\x05\xa1\x01\xc0";
    ev.u.create.rd_size = 7;
    ev.u.create.bus = 3; // BUS_USB
    ev.u.create.vendor = 0x0955;
    ev.u.create.product = 0x7214;
    
    int res = write(fd, &ev, sizeof(ev));
    if (res < 0) {
        printf("[-] Failed to write to /dev/uhid: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] write to /dev/uhid successful.\n");
    
    return fd;
}

int main() {
    printf("[*] Starting reproducer...\n");

    // Ignore SIGCHLD so sleep() is not interrupted by child processes exiting
    if (signal(SIGCHLD, SIG_IGN) == SIG_ERR) {
        printf("[-] Failed to ignore SIGCHLD: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] signal(SIGCHLD, SIG_IGN) successful.\n");

    // 1. Setup thermal netlink listener to force GFP_KERNEL allocation
    int th_nl_fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
    if (th_nl_fd < 0) {
        printf("[-] Failed to socket thermal netlink: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] socket thermal netlink successful.\n");

    int mcast_id = get_mcast_group_id(th_nl_fd, "thermal", "event");
    if (mcast_id > 0) {
        printf("[+] Found thermal event mcast group ID: %d\n", mcast_id);
        int res = setsockopt(th_nl_fd, SOL_NETLINK, NETLINK_ADD_MEMBERSHIP, &mcast_id, sizeof(mcast_id));
        if (res < 0) {
            printf("[-] Failed to setsockopt thermal netlink: %s\n", strerror(errno));
            exit(1);
        }
        printf("[+] setsockopt thermal netlink successful.\n");
    } else {
        printf("[-] Failed to find thermal event mcast group ID. Fallback to 1-255.\n");
        for (int i = 1; i < 256; i++) {
            setsockopt(th_nl_fd, SOL_NETLINK, NETLINK_ADD_MEMBERSHIP, &i, sizeof(i));
        }
    }

    // 2. Setup NBD and trigger reconnect (fs_reclaim -> cpu_hotplug_lock)
    int nl_fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
    if (nl_fd < 0) {
        printf("[-] Failed to socket generic netlink: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] socket generic netlink successful.\n");
    
    int fam = get_family_id(nl_fd, "nbd");
    if (fam < 0) {
        printf("[-] Failed to get NBD family ID.\n");
        exit(1);
    }
    printf("[+] get_family_id successful, fam: %d\n", fam);

    int sv1[2], sv2[2];
    int res = socketpair(AF_UNIX, SOCK_STREAM, 0, sv1);
    if (res < 0) {
        printf("[-] Failed to socketpair 1: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] socketpair 1 successful.\n");

    res = socketpair(AF_UNIX, SOCK_STREAM, 0, sv2);
    if (res < 0) {
        printf("[-] Failed to socketpair 2: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] socketpair 2 successful.\n");

    int fds[2] = {sv1[0], sv2[0]};
    nbd_netlink_cmd(nl_fd, fam, 1, 0, fds, 2); // NBD_CMD_CONNECT
    usleep(100000);
    
    pid_t pid = fork();
    if (pid < 0) {
        printf("[-] Failed to fork: %s\n", strerror(errno));
        exit(1);
    }
    
    if (pid == 0) {
        // CRITICAL: Close inherited sockets so they don't keep the connection alive
        close(sv1[0]);
        close(sv1[1]);
        close(sv2[0]);
        close(sv2[1]);
        close(th_nl_fd);
        close(nl_fd);
        
        // Use O_DIRECT to ensure the read bypasses page cache and hits the block layer
        int fd = open("/dev/nbd0", O_RDONLY | O_DIRECT);
        if (fd >= 0) {
            void *buf;
            if (posix_memalign(&buf, 4096, 4096) == 0) {
                read(fd, buf, 4096);
                free(buf);
            }
            close(fd);
        }
        exit(0);
    }
    printf("[+] fork successful.\n");
    
    usleep(500000); // Wait for read to trigger nbd_queue_rq
    
    close(sv1[1]);
    close(sv2[1]);
    usleep(500000); // Wait for recv_work to mark socket dead
    
    int sv3[2];
    res = socketpair(AF_UNIX, SOCK_STREAM, 0, sv3);
    if (res < 0) {
        printf("[-] Failed to socketpair 3: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] socketpair 3 successful.\n");
    
    int fds2[1] = {sv3[0]};
    nbd_netlink_cmd(nl_fd, fam, 3, 0, fds2, 1); // NBD_CMD_RECONFIGURE
    usleep(100000);
    printf("[+] NBD reconnect triggered successful.\n");

    // 3. Create shield device (thermal_list_lock -> tz->lock -> fs_reclaim)
    int uhid_fd = create_shield();

    printf("[+] Reproducer finished. Waiting for lockdep...\n");
    int rem = 5;
    while (rem > 0) {
        rem = sleep(rem);
    }

    if (uhid_fd >= 0) {
        close(uhid_fd);
    }

    return 0;
}