// https://syzkaller.appspot.com/bug?id=9fb7155e1356de32d373b5cd8fd98407f148a7c4
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/socket.h>
#include <linux/netlink.h>
#include <linux/genetlink.h>
#include <linux/rtnetlink.h>
#include <net/if.h>
#include <arpa/inet.h>
#include <sched.h>
#include <fcntl.h>
#include <time.h>
#include <pthread.h>
#include <errno.h>

#define WG_GENL_NAME "wireguard"
#define WG_CMD_SET_DEVICE 1

enum wgdevice_attribute {
    WGDEVICE_A_UNSPEC,
    WGDEVICE_A_IFINDEX,
    WGDEVICE_A_IFNAME,
    WGDEVICE_A_PRIVATE_KEY,
    WGDEVICE_A_PUBLIC_KEY,
    WGDEVICE_A_FLAGS,
    WGDEVICE_A_LISTEN_PORT,
    WGDEVICE_A_FWMARK,
    WGDEVICE_A_PEERS,
};

enum wgpeer_attribute {
    WGPEER_A_UNSPEC,
    WGPEER_A_PUBLIC_KEY,
    WGPEER_A_PRESHARED_KEY,
    WGPEER_A_FLAGS,
    WGPEER_A_ENDPOINT,
    WGPEER_A_PERSISTENT_KEEPALIVE_INTERVAL,
    WGPEER_A_LAST_HANDSHAKE_TIME,
    WGPEER_A_RX_BYTES,
    WGPEER_A_TX_BYTES,
    WGPEER_A_ALLOWEDIPS,
    WGPEER_A_PROTOCOL_VERSION,
};

enum wgallowedip_attribute {
    WGALLOWEDIP_A_UNSPEC,
    WGALLOWEDIP_A_FAMILY,
    WGALLOWEDIP_A_IPADDR,
    WGALLOWEDIP_A_CIDR_MASK,
};

struct nl_msg {
    struct nlmsghdr n;
    struct genlmsghdr g;
    char buf[2048];
};

#ifndef NLA_ALIGNTO
#define NLA_ALIGNTO 4
#endif
#ifndef NLA_ALIGN
#define NLA_ALIGN(len) (((len) + NLA_ALIGNTO - 1) & ~(NLA_ALIGNTO - 1))
#endif
#ifndef NLA_HDRLEN
#define NLA_HDRLEN ((int) NLA_ALIGN(sizeof(struct nlattr)))
#endif

#ifndef NLA_OK
#define NLA_OK(attr,len) \
	((len) >= (int)sizeof(struct nlattr) && \
	 (attr)->nla_len >= sizeof(struct nlattr) && \
	 (attr)->nla_len <= (len))
#endif
#ifndef NLA_NEXT
#define NLA_NEXT(attr,attrlen)	\
	((attrlen) -= NLA_ALIGN((attr)->nla_len), \
	 (struct nlattr*)(((char*)(attr)) + NLA_ALIGN((attr)->nla_len)))
#endif

void add_attr(struct nl_msg *msg, int type, const void *data, int len) {
    int aligned_len = NLA_ALIGN(msg->n.nlmsg_len);
    int pad_before = aligned_len - msg->n.nlmsg_len;
    if (pad_before > 0) {
        memset(((char *)msg) + msg->n.nlmsg_len, 0, pad_before);
    }
    struct nlattr *attr = (struct nlattr *)(((char *)msg) + aligned_len);
    attr->nla_type = type;
    attr->nla_len = NLA_HDRLEN + len;
    if (len > 0 && data != NULL) {
        memcpy((char *)attr + NLA_HDRLEN, data, len);
        int pad_after = NLA_ALIGN(attr->nla_len) - attr->nla_len;
        if (pad_after > 0) {
            memset((char *)attr + attr->nla_len, 0, pad_after);
        }
    }
    msg->n.nlmsg_len = aligned_len + attr->nla_len;
}

struct nlattr *add_nested(struct nl_msg *msg, int type) {
    int aligned_len = NLA_ALIGN(msg->n.nlmsg_len);
    int pad_before = aligned_len - msg->n.nlmsg_len;
    if (pad_before > 0) {
        memset(((char *)msg) + msg->n.nlmsg_len, 0, pad_before);
    }
    struct nlattr *attr = (struct nlattr *)(((char *)msg) + aligned_len);
    attr->nla_type = type | NLA_F_NESTED;
    attr->nla_len = NLA_HDRLEN;
    msg->n.nlmsg_len = aligned_len + attr->nla_len;
    return attr;
}

void end_nested(struct nl_msg *msg, struct nlattr *attr) {
    attr->nla_len = (char *)msg + msg->n.nlmsg_len - (char *)attr;
}

void nl_send_and_recv(int fd, void *req, int len, const char *desc) {
    if (send(fd, req, len, 0) < 0) {
        printf("[-] send %s failed: %s\n", desc, strerror(errno));
        exit(1);
    }
    char buf[4096];
    int rlen = recv(fd, buf, sizeof(buf), 0);
    if (rlen < 0) {
        printf("[-] recv %s failed: %s\n", desc, strerror(errno));
        exit(1);
    }
    struct nlmsghdr *nh = (struct nlmsghdr *)buf;
    if (nh->nlmsg_type == NLMSG_ERROR) {
        struct nlmsgerr *err = (struct nlmsgerr *)NLMSG_DATA(nh);
        if (err->error != 0) {
            printf("[-] %s NLMSG_ERROR: %s\n", desc, strerror(-err->error));
            exit(1);
        }
    }
    printf("[+] %s successful.\n", desc);
}

void nl_send_and_recv_silent(int fd, void *req, int len, const char *desc) {
    if (send(fd, req, len, 0) < 0) {
        printf("[-] send %s failed: %s\n", desc, strerror(errno));
        exit(1);
    }
    char buf[4096];
    int rlen = recv(fd, buf, sizeof(buf), 0);
    if (rlen < 0) {
        printf("[-] recv %s failed: %s\n", desc, strerror(errno));
        exit(1);
    }
    struct nlmsghdr *nh = (struct nlmsghdr *)buf;
    if (nh->nlmsg_type == NLMSG_ERROR) {
        struct nlmsgerr *err = (struct nlmsgerr *)NLMSG_DATA(nh);
        if (err->error != 0) {
            printf("[-] %s NLMSG_ERROR: %s\n", desc, strerror(-err->error));
            exit(1);
        }
    }
}

int get_family_id(int fd, const char *name) {
    struct nl_msg msg = {0};
    msg.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
    msg.n.nlmsg_type = GENL_ID_CTRL;
    msg.n.nlmsg_flags = NLM_F_REQUEST;
    msg.g.cmd = CTRL_CMD_GETFAMILY;
    msg.g.version = 1;
    
    add_attr(&msg, CTRL_ATTR_FAMILY_NAME, name, strlen(name) + 1);
    
    if (send(fd, &msg, msg.n.nlmsg_len, 0) < 0) {
        printf("[-] send CTRL_CMD_GETFAMILY failed: %s\n", strerror(errno));
        exit(1);
    }
    
    char buf[4096];
    int len = recv(fd, buf, sizeof(buf), 0);
    if (len <= 0) {
        printf("[-] recv CTRL_CMD_GETFAMILY failed: %s\n", strerror(errno));
        exit(1);
    }
    
    struct nlmsghdr *nh = (struct nlmsghdr *)buf;
    if (nh->nlmsg_type == NLMSG_ERROR) {
        printf("[-] NLMSG_ERROR in get_family_id\n");
        exit(1);
    }
    
    struct genlmsghdr *gh = (struct genlmsghdr *)NLMSG_DATA(nh);
    struct nlattr *attr = (struct nlattr *)((char *)gh + GENL_HDRLEN);
    int attr_len = nh->nlmsg_len - NLMSG_LENGTH(GENL_HDRLEN);
    
    while (NLA_OK(attr, attr_len)) {
        if (attr->nla_type == CTRL_ATTR_FAMILY_ID) {
            printf("[+] get_family_id successful.\n");
            return *(uint16_t *)((char *)attr + NLA_HDRLEN);
        }
        attr = NLA_NEXT(attr, attr_len);
    }
    printf("[-] Family ID not found\n");
    exit(1);
}

void setup_wg_dev(int fd, int family_id, const char *ifname, const char *priv_key, uint16_t listen_port) {
    struct nl_msg msg = {0};
    msg.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
    msg.n.nlmsg_type = family_id;
    msg.n.nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    msg.g.cmd = WG_CMD_SET_DEVICE;
    msg.g.version = 1;

    int ifindex = if_nametoindex(ifname);
    if (ifindex == 0) {
        printf("[-] if_nametoindex failed for %s: %s\n", ifname, strerror(errno));
        exit(1);
    }
    add_attr(&msg, WGDEVICE_A_IFINDEX, &ifindex, sizeof(ifindex));
    add_attr(&msg, WGDEVICE_A_PRIVATE_KEY, priv_key, 32);
    add_attr(&msg, WGDEVICE_A_LISTEN_PORT, &listen_port, sizeof(listen_port));

    char desc[64];
    sprintf(desc, "setup_wg_dev %s", ifname);
    nl_send_and_recv(fd, &msg, msg.n.nlmsg_len, desc);
}

void add_peer(int fd, int family_id, const char *ifname, const char *pub_key, const char *endpoint_ip, uint16_t endpoint_port, const char *allowed_ip, uint16_t keepalive) {
    struct nl_msg msg = {0};
    msg.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
    msg.n.nlmsg_type = family_id;
    msg.n.nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    msg.g.cmd = WG_CMD_SET_DEVICE;
    msg.g.version = 1;

    int ifindex = if_nametoindex(ifname);
    if (ifindex == 0) {
        printf("[-] if_nametoindex failed for %s: %s\n", ifname, strerror(errno));
        exit(1);
    }
    add_attr(&msg, WGDEVICE_A_IFINDEX, &ifindex, sizeof(ifindex));

    struct nlattr *peers = add_nested(&msg, WGDEVICE_A_PEERS);
    struct nlattr *peer = add_nested(&msg, 0);
    
    add_attr(&msg, WGPEER_A_PUBLIC_KEY, pub_key, 32);
    
    struct sockaddr_in endpoint = {0};
    endpoint.sin_family = AF_INET;
    endpoint.sin_port = htons(endpoint_port);
    inet_pton(AF_INET, endpoint_ip, &endpoint.sin_addr);
    add_attr(&msg, WGPEER_A_ENDPOINT, &endpoint, sizeof(endpoint));
    
    add_attr(&msg, WGPEER_A_PERSISTENT_KEEPALIVE_INTERVAL, &keepalive, sizeof(keepalive));
    
    struct nlattr *allowedips = add_nested(&msg, WGPEER_A_ALLOWEDIPS);
    struct nlattr *allowedip = add_nested(&msg, 0);
    
    uint16_t family = AF_INET;
    add_attr(&msg, WGALLOWEDIP_A_FAMILY, &family, sizeof(family));
    struct in_addr ip;
    inet_pton(AF_INET, allowed_ip, &ip);
    add_attr(&msg, WGALLOWEDIP_A_IPADDR, &ip, sizeof(ip));
    uint8_t cidr = 32;
    add_attr(&msg, WGALLOWEDIP_A_CIDR_MASK, &cidr, sizeof(cidr));
    
    end_nested(&msg, allowedip);
    end_nested(&msg, allowedips);
    end_nested(&msg, peer);
    end_nested(&msg, peers);

    char desc[64];
    sprintf(desc, "add_peer to %s", ifname);
    nl_send_and_recv(fd, &msg, msg.n.nlmsg_len, desc);
}

void update_keepalive(int fd, int family_id, const char *ifname, const char *pub_key, uint16_t keepalive) {
    struct nl_msg msg = {0};
    msg.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
    msg.n.nlmsg_type = family_id;
    msg.n.nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    msg.g.cmd = WG_CMD_SET_DEVICE;
    msg.g.version = 1;

    int ifindex = if_nametoindex(ifname);
    if (ifindex == 0) {
        printf("[-] if_nametoindex failed for %s: %s\n", ifname, strerror(errno));
        exit(1);
    }
    add_attr(&msg, WGDEVICE_A_IFINDEX, &ifindex, sizeof(ifindex));

    struct nlattr *peers = add_nested(&msg, WGDEVICE_A_PEERS);
    struct nlattr *peer = add_nested(&msg, 0);
    
    add_attr(&msg, WGPEER_A_PUBLIC_KEY, pub_key, 32);
    add_attr(&msg, WGPEER_A_PERSISTENT_KEEPALIVE_INTERVAL, &keepalive, sizeof(keepalive));
    
    end_nested(&msg, peer);
    end_nested(&msg, peers);

    char desc[64];
    sprintf(desc, "update_keepalive %s", ifname);
    nl_send_and_recv_silent(fd, &msg, msg.n.nlmsg_len, desc);
}

#define NLMSG_TAIL(nmsg) \
    ((struct rtattr *) (((void *) (nmsg)) + NLMSG_ALIGN((nmsg)->nlmsg_len)))

void addattr_l(struct nlmsghdr *n, int maxlen, int type, const void *data, int alen) {
    int len = RTA_LENGTH(alen);
    struct rtattr *rta;
    if (NLMSG_ALIGN(n->nlmsg_len) + RTA_ALIGN(len) > maxlen) return;
    rta = NLMSG_TAIL(n);
    rta->rta_type = type;
    rta->rta_len = len;
    if (alen) memcpy(RTA_DATA(rta), data, alen);
    n->nlmsg_len = NLMSG_ALIGN(n->nlmsg_len) + RTA_ALIGN(len);
}

struct rtattr *addattr_nest(struct nlmsghdr *n, int maxlen, int type) {
    struct rtattr *nest = NLMSG_TAIL(n);
    addattr_l(n, maxlen, type, NULL, 0);
    return nest;
}

void addattr_nest_end(struct nlmsghdr *n, struct rtattr *nest) {
    nest->rta_len = (void *)NLMSG_TAIL(n) - (void *)nest;
}

int create_wg_iface(const char *ifname) {
    int fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
    if (fd < 0) {
        printf("[-] socket(AF_NETLINK) failed: %s\n", strerror(errno));
        exit(1);
    }

    struct {
        struct nlmsghdr n;
        struct ifinfomsg i;
        char buf[1024];
    } req = {0};

    req.n.nlmsg_len = NLMSG_LENGTH(sizeof(struct ifinfomsg));
    req.n.nlmsg_flags = NLM_F_REQUEST | NLM_F_CREATE | NLM_F_EXCL | NLM_F_ACK;
    req.n.nlmsg_type = RTM_NEWLINK;
    req.i.ifi_family = AF_UNSPEC;

    addattr_l(&req.n, sizeof(req), IFLA_IFNAME, ifname, strlen(ifname) + 1);
    
    struct rtattr *linkinfo = addattr_nest(&req.n, sizeof(req), IFLA_LINKINFO);
    addattr_l(&req.n, sizeof(req), IFLA_INFO_KIND, "wireguard", 10);
    addattr_nest_end(&req.n, linkinfo);

    char desc[64];
    sprintf(desc, "create_wg_iface %s", ifname);
    nl_send_and_recv(fd, &req, req.n.nlmsg_len, desc);
    close(fd);
    return 0;
}

int set_iface_up(const char *ifname) {
    int fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
    if (fd < 0) {
        printf("[-] socket(AF_NETLINK) failed: %s\n", strerror(errno));
        exit(1);
    }
    
    struct {
        struct nlmsghdr n;
        struct ifinfomsg i;
        char buf[1024];
    } req = {0};

    req.n.nlmsg_len = NLMSG_LENGTH(sizeof(struct ifinfomsg));
    req.n.nlmsg_flags = NLM_F_REQUEST | NLM_F_ACK;
    req.n.nlmsg_type = RTM_NEWLINK;
    req.i.ifi_family = AF_UNSPEC;
    req.i.ifi_index = if_nametoindex(ifname);
    if (req.i.ifi_index == 0) {
        printf("[-] if_nametoindex failed for %s: %s\n", ifname, strerror(errno));
        exit(1);
    }
    req.i.ifi_flags = IFF_UP;
    req.i.ifi_change = IFF_UP;

    char desc[64];
    sprintf(desc, "set_iface_up %s", ifname);
    nl_send_and_recv(fd, &req, req.n.nlmsg_len, desc);
    close(fd);
    return 0;
}

int assign_ip(const char *ifname, const char *ip) {
    int fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_ROUTE);
    if (fd < 0) {
        printf("[-] socket(AF_NETLINK) failed: %s\n", strerror(errno));
        exit(1);
    }
    
    struct {
        struct nlmsghdr n;
        struct ifaddrmsg i;
        char buf[1024];
    } req = {0};

    req.n.nlmsg_len = NLMSG_LENGTH(sizeof(struct ifaddrmsg));
    req.n.nlmsg_flags = NLM_F_REQUEST | NLM_F_CREATE | NLM_F_EXCL | NLM_F_ACK;
    req.n.nlmsg_type = RTM_NEWADDR;
    req.i.ifa_family = AF_INET;
    req.i.ifa_prefixlen = 24;
    req.i.ifa_index = if_nametoindex(ifname);
    if (req.i.ifa_index == 0) {
        printf("[-] if_nametoindex failed for %s: %s\n", ifname, strerror(errno));
        exit(1);
    }

    struct in_addr addr;
    inet_pton(AF_INET, ip, &addr);
    addattr_l(&req.n, sizeof(req), IFA_LOCAL, &addr, sizeof(addr));
    addattr_l(&req.n, sizeof(req), IFA_ADDRESS, &addr, sizeof(addr));

    char desc[64];
    sprintf(desc, "assign_ip %s to %s", ip, ifname);
    nl_send_and_recv(fd, &req, req.n.nlmsg_len, desc);
    close(fd);
    return 0;
}

int sock0, sock1;
struct sockaddr_in dst0, dst1;

void *send_packets_wg0(void *arg) {
    char buf[64] = "hello";
    while (1) {
        if (sendto(sock0, buf, sizeof(buf), 0, (struct sockaddr *)&dst0, sizeof(dst0)) < 0) {
            // Ignore ENOBUFS and EAGAIN
        }
        usleep(100);
    }
    return NULL;
}

void *send_packets_wg1(void *arg) {
    char buf[64] = "hello";
    while (1) {
        if (sendto(sock1, buf, sizeof(buf), 0, (struct sockaddr *)&dst1, sizeof(dst1)) < 0) {
            // Ignore ENOBUFS and EAGAIN
        }
        usleep(100);
    }
    return NULL;
}

int main() {
    if (unshare(CLONE_NEWNET) < 0) {
        printf("[-] unshare failed: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] unshare successful.\n");

    set_iface_up("lo");

    create_wg_iface("wg0");
    create_wg_iface("wg1");

    assign_ip("wg0", "10.0.0.1");
    assign_ip("wg1", "10.0.0.2");

    set_iface_up("wg0");
    set_iface_up("wg1");

    // RFC 7748 Curve25519 Test Vectors
    char priv1[] = "\x77\x07\x6d\x0a\x73\x18\xa5\x7d\x3c\x16\xc1\x72\x51\xb2\x66\x45\xdf\x4c\x2f\x87\xeb\xc0\x99\x2a\xb1\x77\xfb\xa5\x1d\xb9\x2c\x2a";
    char pub1[] = "\x85\x20\xf0\x09\x89\x30\xa7\x54\x74\x8b\x7d\xdc\xb4\x3e\xf7\x5a\x0d\xbf\x3a\x0d\x26\x38\x1a\xf4\xeb\xa4\xa9\x8e\xaa\x9b\x4e\x6a";

    char priv2[] = "\x5d\xab\x08\x7e\x62\x4a\x8a\x4b\x79\xe1\x7f\x8b\x83\x80\x0e\xe6\x6f\x3b\xb1\x29\x26\x18\xb6\xfd\x1c\x2f\x8b\x27\xff\x88\xe0\xeb";
    char pub2[] = "\xde\x9e\xdb\x7d\x7b\x7d\xc1\xb4\xd3\x5b\x61\xc2\xec\xe4\x35\x37\x3f\x83\x43\xc8\x5b\x78\x67\x4d\xad\xfc\x7e\x14\x6f\x88\x2b\x4f";

    int nl_fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC);
    if (nl_fd < 0) {
        printf("[-] socket(AF_NETLINK) failed: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] socket(AF_NETLINK) successful.\n");

    int family_id = get_family_id(nl_fd, WG_GENL_NAME);

    setup_wg_dev(nl_fd, family_id, "wg0", priv1, 51820);
    setup_wg_dev(nl_fd, family_id, "wg1", priv2, 51821);

    add_peer(nl_fd, family_id, "wg0", pub2, "127.0.0.1", 51821, "10.0.0.2", 65535);
    add_peer(nl_fd, family_id, "wg1", pub1, "127.0.0.1", 51820, "10.0.0.1", 65535);

    sock0 = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock0 < 0) {
        printf("[-] socket(AF_INET) failed: %s\n", strerror(errno));
        exit(1);
    }
    if (setsockopt(sock0, SOL_SOCKET, SO_BINDTODEVICE, "wg0", 4) < 0) {
        printf("[-] setsockopt SO_BINDTODEVICE wg0 failed: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] sock0 setup successful.\n");

    dst0.sin_family = AF_INET;
    dst0.sin_port = htons(12345);
    inet_pton(AF_INET, "10.0.0.2", &dst0.sin_addr);

    sock1 = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock1 < 0) {
        printf("[-] socket(AF_INET) failed: %s\n", strerror(errno));
        exit(1);
    }
    if (setsockopt(sock1, SOL_SOCKET, SO_BINDTODEVICE, "wg1", 4) < 0) {
        printf("[-] setsockopt SO_BINDTODEVICE wg1 failed: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] sock1 setup successful.\n");

    dst1.sin_family = AF_INET;
    dst1.sin_port = htons(12345);
    inet_pton(AF_INET, "10.0.0.1", &dst1.sin_addr);

    pthread_t t1, t2;
    if (pthread_create(&t1, NULL, send_packets_wg0, NULL) != 0) {
        printf("[-] pthread_create t1 failed: %s\n", strerror(errno));
        exit(1);
    }
    if (pthread_create(&t2, NULL, send_packets_wg1, NULL) != 0) {
        printf("[-] pthread_create t2 failed: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] threads created successfully.\n");

    printf("[+] Starting race loop...\n");
    for (int i = 0; i < 500; i++) {
        update_keepalive(nl_fd, family_id, "wg0", pub2, 65535);
        update_keepalive(nl_fd, family_id, "wg1", pub1, 65535);
        usleep(5000);
        update_keepalive(nl_fd, family_id, "wg0", pub2, 1);
        update_keepalive(nl_fd, family_id, "wg1", pub1, 1);
        usleep(5000);
    }
    printf("[+] Race loop finished.\n");

    // Sleep forever to avoid task exit KCSAN bugs.
    // The executor will kill the process after a timeout.
    while (1) {
        sleep(1000);
    }

    return 0;
}