#include <stdio.h>
#include <err.h>
#include <string.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <inttypes.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <unistd.h>
#include <fcntl.h>
#include <linux/types.h>
#include <linux/kvm.h>
#include <linux/kcov.h>
#include <linux/ptrace.h>
#include <sys/sysinfo.h>

#include <sys/un.h>
#include <sys/wait.h>
#include <sys/socket.h>
#include <errno.h>

/* ++ uapi/pkvm_shmem.h */
#define KVM_SHMEM_ALLOC_TYPE 1

enum kvm_shmem_alloc_type { 
    KVM_SHMEM_ALLOCTYPE_VMALLOC,
    KVM_SHMEM_ALLOCTYPE_PAGES_EXACT };

#define KVM_SHMEM_ALLOC(alloc) _IO(KVM_SHMEM_ALLOC_TYPE, alloc)
#define KVM_SHMEM_ALLOC_PAGES KVM_SHMEM_ALLOC(KVM_SHMEM_ALLOCTYPE_PAGES_EXACT)
#define KVM_SHMEM_VMALLOC KVM_SHMEM_ALLOC(KVM_SHMEM_ALLOCTYPE_VMALLOC)

// ioctl on the mmapable fd from the KVM_SHMEM_ALLOC ioctl
#define KVM_SHMEM_AREA_KADDR _IOR(1, 0, void *)
#define KVM_SHMEM_AREA_PHYS _IOR(1, 1, void *)
#define KVM_SHMEM_AREA_MAKE_READONLY _IO(1, 2)
/* -- uapi/pkvm_shmem.h */

#define PAGE_SHIFT 12
#define PAGE_SIZE (1 << PAGE_SHIFT)

struct shmem_area {
    int size;
    int fd;
    void* mmap;
    ulong phys;
};

void shmem_area_alloc(int fd, struct shmem_area *area)
{
    area->fd = ioctl(fd, KVM_SHMEM_ALLOC_PAGES, area->size);
    if (area->fd < 0)
        err(1, "Can't allocate shmem_area");

    if (ioctl(area->fd, KVM_SHMEM_AREA_PHYS, &area->phys))
        err(1, "Can't get kernel area physical address");
}

void shmem_area_mmap(struct shmem_area *area)
{
    area->mmap = mmap(NULL, area->size, PROT_READ | PROT_WRITE,
                      MAP_SHARED, area->fd, 0);
    if ((void*)area->mmap == MAP_FAILED)
        err(1, "Can't mmap kernel area");
}

void shmem_area_make_readonly(struct shmem_area *area)
{
    if (ioctl(area->fd, KVM_SHMEM_AREA_MAKE_READONLY))
        err(1, "Can't make area r/o");
}

static
void send_fd(int socket, int *fds, int n)  // send fd by socket
{
    struct msghdr msg = {0};
    struct cmsghdr *cmsg;
    char buf[CMSG_SPACE(n * sizeof(int))], dup[256];
    struct iovec io = { .iov_base = &dup, .iov_len = sizeof(dup) };

    memset(buf, 0, sizeof(buf));

    msg.msg_iov = &io;
    msg.msg_iovlen = 1;
    msg.msg_control = buf;
    msg.msg_controllen = sizeof(buf);

    cmsg = CMSG_FIRSTHDR(&msg);
    cmsg->cmsg_level = SOL_SOCKET;
    cmsg->cmsg_type = SCM_RIGHTS;
    cmsg->cmsg_len = CMSG_LEN(n * sizeof(int));

    memcpy ((int *) CMSG_DATA(cmsg), fds, n * sizeof (int));

    if (sendmsg(socket, &msg, 0) < 0)
        err(1, "Failed to send message");
}


int mk_uds(const char *path)
{
    int sfd;
    struct sockaddr_un addr;

    sfd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (sfd == -1)
        err(1, "Failed to create socket");

    if (unlink(path) == -1 && errno != ENOENT)
        err(1, "Removing socket file failed");

    memset(&addr, 0, sizeof(struct sockaddr_un));
    addr.sun_family = AF_UNIX;
    strncpy(addr.sun_path, path, sizeof(addr.sun_path) - 1);

    if (bind(sfd, (struct sockaddr *)&addr,
             sizeof(struct sockaddr_un)) == -1)
        err(1, "Failed to bind to socket");

    if (listen(sfd, 5) == -1)
        err(1, NULL);

    return sfd;
}

int main(int argc, char** argv)
{
    int fd, sock_fd;
    struct shmem_area area = { .size = PAGE_SIZE };

    fd = open("/sys/kernel/debug/pkvm_shmem", O_RDWR);
    if (fd == -1) err(1, "pkvm_shmem");

    shmem_area_alloc(fd, &area);
    shmem_area_mmap(&area);
    printf("Area: pa=%lx va=%p size=%d\n",
           area.phys, area.mmap, area.size);
    *(volatile uint32_t *)area.mmap = 0xdeadf00d;
    shmem_area_make_readonly(&area);
    printf("Written %08"PRIx32" to %lx\n",
           *(volatile uint32_t *)area.mmap, area.phys);
    printf("Area %lx-%lx is now read-only\n",
           area.phys, area.phys+area.size-1);

    sock_fd = mk_uds("shmem.sock");
    for (;;) {
        int cfd = accept(sock_fd, NULL, NULL);
        if (cfd == -1)
            err(1, "Failed to accept incoming connection");
        send_fd(cfd, &area.fd, 1);
        close(cfd);
    }

    return 0;
}

/*
 * Local variables:
 * mode: C
 * c-file-style: "Linux"
 * c-basic-offset: 4
 * tab-width: 4
 * indent-tabs-mode: nil
 * End:
 */