#include <stdio.h>
#include <err.h>
#include <string.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.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 kaddr;
    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_KADDR, &area->kaddr))
		err(1, "Can't get kernel area address");
	printf("Allocated area at %lx of size 0x%x bytes\n",
	       area->kaddr, area->size);

	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)
{
	printf("Try to mmap kernel address %lx fd %d\n", area->kaddr, area->fd);
	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: ka=%lx pa=%lx va=%p size=%d\n",
	       area.kaddr, area.phys, area.mmap, area.size);
	shmem_area_make_readonly(&area);

	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;
}