/* radare - LGPL - Copyright 2009-2024 - pancake */

#include <r_userconf.h>
#include <sys/sysctl.h>
#if DEBUGGER

// TODO much work remains to be done
#include "xnu_debug.h"
#include "xnu_threads.h"

static void xnu_thread_free(xnu_thread_t *thread) {
	kern_return_t kr;
	if (!thread) {
		return;
	}
	free (thread->name);
	// if we free our thread from the list we need to decrement the ref
	// count
	kr = mach_port_deallocate (mach_task_self (), thread->port);
	if (kr != KERN_SUCCESS) {
		R_LOG_ERROR ("failed to deallocate thread port");
	}
	free (thread);
}

static kern_return_t xnu_convert_thread_state(mach_port_t thread, int direction, thread_state_flavor_t flavor, const void * in_state, mach_msg_type_number_t in_state_count, void * out_state, mach_msg_type_number_t * out_state_count)
{
	static bool initialized = false;
	static kern_return_t (* convert)(thread_act_t thread, int direction, thread_state_flavor_t flavor,
		thread_state_t in_state, mach_msg_type_number_t in_state_count,
		thread_state_t out_state, mach_msg_type_number_t * out_state_count);
	if (!initialized) {
		void * module;
		module = dlopen ("/usr/lib/system/libsystem_kernel.dylib", RTLD_GLOBAL | RTLD_LAZY);
		convert = dlsym (module, "thread_convert_thread_state");
		dlclose (module);
		initialized = true;
	}
	if (convert == NULL) {
		goto fallback;
	}
	return convert (thread, direction, flavor, (thread_state_t) in_state, in_state_count, out_state, out_state_count);
fallback:
	{
		const mach_msg_type_number_t n = MIN (in_state_count, *out_state_count);
		if (out_state != in_state) {
			memmove (out_state, in_state, n * sizeof (integer_t));
		}
		*out_state_count = n;
		return KERN_SUCCESS;
	}
}

// XXX this should work as long as in arm trace bit relies on this
bool xnu_thread_get_drx(RDebug *dbg, xnu_thread_t *thread) {
	R_RETURN_VAL_IF_FAIL (dbg && thread, false);
	kern_return_t rc;
#if __x86_64__ || __i386__
	thread->flavor = x86_DEBUG_STATE;
	thread->count = x86_DEBUG_STATE_COUNT;
	thread->state_size = R_SYS_BITS_CHECK (dbg->bits, 64)
		? sizeof (x86_debug_state64_t)
		: sizeof (x86_debug_state32_t);
	thread->state = &thread->drx.uds;
	rc = thread_get_state (thread->port, thread->flavor,
	       (thread_state_t)&thread->drx, &thread->count);
#elif __arm64__ || __arm64 || __aarch64 || __aarch64__ || __arm64e__
	if (R_SYS_BITS_CHECK (dbg->bits, 64)) {
		thread->count = ARM_DEBUG_STATE64_COUNT;
		thread->flavor = ARM_DEBUG_STATE64;
		rc = thread_get_state (thread->port, thread->flavor,
				       (thread_state_t)&thread->debug.drx64,
				       &thread->count);
	} else {
		thread->count = ARM_DEBUG_STATE32_COUNT;
		thread->flavor = ARM_DEBUG_STATE32;
		rc = thread_get_state (thread->port, thread->flavor,
				       (thread_state_t)&thread->debug.drx32,
				       &thread->count);
	}
#elif __arm__ || __arm || __armv7__
	thread->count = ARM_DEBUG_STATE_COUNT;
	thread->flavor = ARM_DEBUG_STATE;
	rc = thread_get_state (thread->port, thread->flavor,
			       (thread_state_t)&thread->debug.drx,
			       &thread->count);
#else
#warning xnu_thread_get_drx: Unsupported architecture
	rc = KERN_FAILURE;
#endif
	if (rc != KERN_SUCCESS) {
		r_sys_perror (__FUNCTION__);
		thread->count = 0;
		return false;
	}
	return true;
}

bool xnu_thread_set_drx(RDebug *dbg, xnu_thread_t *thread) {
	R_RETURN_VAL_IF_FAIL (dbg && thread, false);
	kern_return_t rc;
#if __i386__ || __x86_64__
	x86_debug_state_t *regs = &thread->drx;
	if (!regs) {
		return false;
	}
	thread->flavor = x86_DEBUG_STATE;
	thread->count = x86_DEBUG_STATE_COUNT;
	if (R_SYS_BITS_CHECK (dbg->bits, 64)) {
		regs->dsh.flavor = x86_DEBUG_STATE64;
		regs->dsh.count  = x86_DEBUG_STATE64_COUNT;
	} else {
		regs->dsh.flavor = x86_DEBUG_STATE32;
		regs->dsh.count  = x86_DEBUG_STATE32_COUNT;
	}
	rc = thread_set_state (thread->port, thread->flavor,
			       (thread_state_t)regs, thread->count);
#elif __arm64__ || __arm64 || __aarch64 || __aarch64__ || __arm64e__
	if (R_SYS_BITS_CHECK (dbg->bits, 64)) {
		thread->count = ARM_DEBUG_STATE64_COUNT;
		thread->flavor = ARM_DEBUG_STATE64;
		rc = thread_set_state (thread->port, thread->flavor,
						(thread_state_t)&thread->debug.drx64,
						thread->count);
	} else {
		thread->count = ARM_DEBUG_STATE32_COUNT;
		thread->flavor = ARM_DEBUG_STATE32;
		rc = thread_set_state (thread->port, thread->flavor,
				       (thread_state_t)&thread->debug.drx32,
				       thread->count);
	}
#elif __arm__ || __arm || __armv7__
	thread->count = ARM_DEBUG_STATE_COUNT;
	thread->flavor = ARM_DEBUG_STATE;
	rc = thread_set_state (thread->port, thread->flavor,
					(thread_state_t)&thread->debug.drx,
					thread->count);
#elif __POWERPC__
/* not supported */
# ifndef PPC_DEBUG_STATE32
# define PPC_DEBUG_STATE32 1
# endif
	//ppc_debug_state_t *regs;
	//thread->flavor = PPC_DEBUG_STATE32;
	//thread->count  = R_MIN (thread->count, sizeof (regs->uds.ds32));
	return false;
#else
	regs->dsh.flavor = 0;
	thread->count = 0;
#endif
	if (rc != KERN_SUCCESS) {
		r_sys_perror (__FUNCTION__);
		thread->count = 0;
		return false;
	}
	return true;
}

static bool xnu_thread_set_gpr(RDebug *dbg, xnu_thread_t *thread) {
	R_RETURN_VAL_IF_FAIL (dbg && thread, false);
	kern_return_t rc;
	R_REG_T *regs = &thread->gpr;
	if (!regs) {
		return false;
	}
#if __i386__ || __x86_64__
	// thread->flavor is used in a switch+case but in regs->tsh.flavor we
	// specify
	thread->state = &regs->uts;
	thread->flavor = x86_THREAD_STATE;
	thread->count = x86_THREAD_STATE_COUNT;
	if (R_SYS_BITS_CHECK (dbg->bits, 64)) {
		regs->tsh.flavor = x86_THREAD_STATE64;
		regs->tsh.count  = x86_THREAD_STATE64_COUNT;
	} else {
		regs->tsh.flavor = x86_THREAD_STATE32;
		regs->tsh.count  = x86_THREAD_STATE32_COUNT;
	}
#elif __arm64 || __aarch64 || __arm64__ || __aarch64__ || __arm64e__
	if (R_SYS_BITS_CHECK (dbg->bits, 64)) {
		thread_state_t remote_state = alloca (thread->state_size);

		thread->count = ARM_UNIFIED_THREAD_STATE_COUNT;
		thread->flavor = ARM_UNIFIED_THREAD_STATE;
		ut32 count = ARM_UNIFIED_THREAD_STATE_COUNT;
		thread->state = &regs->uts;
		thread->state_size = sizeof (arm_thread_state64_t);

		rc = xnu_convert_thread_state (
			thread->port,
			THREAD_CONVERT_THREAD_STATE_FROM_SELF,
			thread->flavor,
			regs, count,
			remote_state, &count);
		if (rc == KERN_SUCCESS) {
			R_LOG_INFO ("convert from self works");
		} else {
			R_LOG_DEBUG ("failed to convert %d", rc);
		}

		thread->flavor = ARM_UNIFIED_THREAD_STATE;
		rc = thread_set_state (thread->port, thread->flavor,
				       remote_state,
				       thread->count);
		if (rc != KERN_SUCCESS) {
			r_sys_perror (__FUNCTION__);
			thread->count = 0;
			return false;
		}
		return true;
	} else {
		thread->flavor = ARM_THREAD_STATE32;
		thread->count = ARM_THREAD_STATE32_COUNT;
		thread->state_size = sizeof (arm_thread_state32_t);
		thread->state = regs;
	}
#elif __arm || __arm__ || __armv7__
	thread->flavor = ARM_THREAD_STATE;
	thread->count = ARM_THREAD_STATE_COUNT;
	thread->state_size = sizeof (arm_thread_state_t);

#endif
	rc = thread_set_state (thread->port, thread->flavor,
			       (thread_state_t)regs, thread->count);
	if (rc != KERN_SUCCESS) {
		r_sys_perror (__FUNCTION__);
		thread->count = 0;
		return false;
	}
	return true;
}

static bool xnu_thread_get_gpr(RDebug *dbg, xnu_thread_t *thread) {
	R_RETURN_VAL_IF_FAIL (dbg && thread, false);
	R_REG_T *regs = &thread->gpr;
	if (!regs) {
		return false;
	}
	kern_return_t rc;
#if __POWERPC__
	thread->state = regs;
#elif  __arm64 || __aarch64 || __arch64__ || __arm64__ || __arm64e__
	if (R_SYS_BITS_CHECK (dbg->bits, 64)) {
		thread->flavor = ARM_UNIFIED_THREAD_STATE;
		thread->count = ARM_UNIFIED_THREAD_STATE_COUNT;
		thread->state_size = sizeof (arm_thread_state64_t);
		thread->state = &regs->uts;
	} else {
		thread->flavor = ARM_THREAD_STATE;
		thread->count = ARM_THREAD_STATE_COUNT;
		thread->state_size = sizeof (arm_thread_state32_t);
		thread->state = regs;
	}
#elif  __arm || __arm__ || __armv7__
	thread->state = regs;
	thread->flavor = ARM_THREAD_STATE;
	thread->count = ARM_THREAD_STATE_COUNT;
	thread->state_size = sizeof (arm_thread_state_t);
#elif __x86_64__ || __i386__
	thread->state = &regs->uts;
	thread->flavor = x86_THREAD_STATE;
	thread->count = x86_THREAD_STATE_COUNT;
	thread->state_size = R_SYS_BITS_CHECK (dbg->bits, 64)
		     ? sizeof (x86_thread_state64_t)
		     : sizeof (x86_thread_state32_t);
#endif
#if defined(THREAD_CONVERT_THREAD_STATE_TO_SELF)
#if !__POWERPC__
	rc = thread_get_state (thread->port, thread->flavor,
			       (thread_state_t)regs, &thread->count);

	if (rc == KERN_SUCCESS) {
		ut32 count = thread->count;
		kern_return_t rc;
		rc = xnu_convert_thread_state (
			thread->port,
			THREAD_CONVERT_THREAD_STATE_TO_SELF,
			thread->flavor,
			regs, count,
			regs, &count);
		if (rc != KERN_SUCCESS) {
			R_LOG_DEBUG ("failed to convert %d", rc);
		}
#endif
#if  __arm64e__
		else {
			if (R_SYS_BITS_CHECK (dbg->bits, 64)) {
				arm_thread_state64_ptrauth_strip (regs->ts_64);
			}
		}
#endif
#if !__POWERPC__
	}
#endif
#endif
	if (rc != KERN_SUCCESS) {
		r_sys_perror (__FUNCTION__);
		thread->count = 0;
		return false;
	}
	return true;
}

static bool xnu_fill_info_thread(RDebug *dbg, xnu_thread_t *thread) {
#if __POWERPC__
	thread->name = strdup ("unknown");
	return false;
#else
	mach_msg_type_number_t count = THREAD_BASIC_INFO_COUNT;
	thread_identifier_info_data_t identifier_info;
	kern_return_t kr = thread_info (thread->port, THREAD_BASIC_INFO,
		(thread_info_t)&thread->basic_info, &count);
	if (kr != KERN_SUCCESS) {
		R_LOG_ERROR ("Fail to get thread_basic_info");
		return false;
	}
	count = THREAD_IDENTIFIER_INFO_COUNT;
	kr = thread_info (thread->port, THREAD_IDENTIFIER_INFO,
			  (thread_info_t)&identifier_info, &count);
	if (kr != KERN_SUCCESS) {
		R_LOG_ERROR ("Fail to get thread_identifier_info");
		return false;
	}
	R_FREE (thread->name);
#if TARGET_OS_IPHONE
	// TODO proc_pidinfo here
	thread->name = strdup ("unknown");
#else
	struct proc_threadinfo proc_threadinfo;
	int ret_proc = proc_pidinfo (dbg->pid, PROC_PIDTHREADINFO,
				 identifier_info.thread_handle,
				 &proc_threadinfo, PROC_PIDTHREADINFO_SIZE);
	if (ret_proc && proc_threadinfo.pth_name[0]) {
		thread->name = strdup (proc_threadinfo.pth_name);
	} else {
		thread->name = strdup ("unknown");
	}
#endif
#endif
	return true;
}

static xnu_thread_t *xnu_get_thread_with_info(RDebug *dbg, thread_t port) {
	xnu_thread_t *thread = R_NEW0 (xnu_thread_t);
	if (!thread) {
		return NULL;
	}
	thread->port = port;
	if (!xnu_fill_info_thread (dbg, thread)) {
		free (thread->name);
		thread->name = strdup ("unknown");
	}
	return thread;
}

static int xnu_update_thread_info(RDebug *dbg, xnu_thread_t *thread) {
	if (!xnu_fill_info_thread (dbg, thread)) {
		free (thread->name);
		thread->name = strdup ("unknown");
	}
	return true;
}

static int thread_find(thread_t *port, xnu_thread_t *a) {
	return (a && port && (a->port == *port))? 0: 1;
}

static int xnu_update_thread_list(RDebug *dbg) {
	thread_array_t thread_list = NULL;
	unsigned int thread_count  = 0;
	xnu_thread_t *thread;
	kern_return_t kr;
	task_t task;
	int i;

	if (!dbg->threads) {
		dbg->threads = r_list_newf ((RListFree)&xnu_thread_free);
		if (!dbg->threads) {
			return false;
		}
	}
	// ok we have the list that will hold our threads, now is time to get
	// them
	task = pid_to_task (dbg->pid);
	if (!task) {
		return false;
	}
	kr = task_threads (task, &thread_list, &thread_count);
	if (kr != KERN_SUCCESS) {
		// we can get into this when the process has terminated but we
		// still hold the old task because we are caching it
		R_LOG_ERROR ("Failed to get list of task's threads");
		return false;
	}
	if (r_list_empty (dbg->threads)) {
		// it's the first time write all threads inside the list
		for (i = 0; i < thread_count; i++) {
			thread = xnu_get_thread_with_info (dbg, thread_list[i]);
			if (!thread) {
				R_LOG_ERROR ("Failed to fill_thread");
				continue;
			}
			if (!r_list_append (dbg->threads, thread)) {
				R_LOG_ERROR ("Failed to add thread to list");
				xnu_thread_free (thread);
			}
		}
	} else {
		RListIter *iter, *iter2;
		// first pass to get rid of those threads that are not longer
		// alive
		r_list_foreach_safe (dbg->threads, iter, iter2, thread) {
			bool flag = true;  // this flag will denote when delete
					   // a thread
			for (i = 0; i < thread_count; i++) {
				if (thread->port == thread_list[i]) {
					flag = false;
					break;
				}
			}
			if (flag) {
				// it is not longer alive so remove from the
				// list
				r_list_delete (dbg->threads, iter);
			} else {
				// otherwise update the info
				xnu_update_thread_info (dbg, thread);
			}
		}
		// ok now we have to insert those threads that we don't have
		for (i = 0; i < thread_count; i++) {
			xnu_thread_t *t;
			iter = r_list_find (dbg->threads, &thread_list[i],
					    (RListComparator)&thread_find);
			// it means is already in our list
			if (iter) {
				// free the ownership over the thread
				kr = mach_port_deallocate (mach_task_self (),
							   thread_list[i]);
				if (kr != KERN_SUCCESS) {
					R_LOG_ERROR ("Failed to deallocate port");
				}
				continue;
			}
			// otherwise insert it
			t = xnu_get_thread_with_info (dbg, thread_list[i]);
			r_list_append (dbg->threads, t);
		}
	}
	// once that is over we need to free the buffer
	(void)vm_deallocate (mach_task_self (), (mach_vm_address_t)thread_list,
			thread_count * sizeof (thread_t));
	return true;
}
#endif