From f85935f8f3038b26d031fc2af919953ee11d56d2 Mon Sep 17 00:00:00 2001
From: EnderIce2 <enderice2@users.noreply.github.com>
Date: Sat, 2 Mar 2024 00:51:37 +0200
Subject: [PATCH] Add GetCurrentProcess function and update KillThread
 signature

---
 include/base.h     | 9 ++++++++-
 include/driver.h   | 3 ++-
 library/task.c     | 9 +++++++--
 misc/vmware/main.c | 4 +++-
 4 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/include/base.h b/include/base.h
index 8806d2a0..6ae5b4cf 100644
--- a/include/base.h
+++ b/include/base.h
@@ -182,6 +182,12 @@ extern "C"
 	pid_t CreateKernelThread(pid_t pId, const char *Name, void *EntryPoint,
 							 void *Argument);
 
+	/**
+	 * @brief Get the PID of the current process
+	 * @return The PID of the current process
+	 */
+	pid_t GetCurrentProcess();
+
 	/**
 	 * @brief Kill a process
 	 * @param pId The PID of the process to kill
@@ -193,10 +199,11 @@ extern "C"
 	/**
 	 * @brief Kill a thread
 	 * @param tId The TID of the thread to kill
+	 * @param pId The PID of the process the thread is in
 	 * @param ExitCode The exit code of the thread
 	 * @return 0 on success, errno on failure
 	 */
-	int KillThread(pid_t tId, int ExitCode);
+	int KillThread(pid_t tId, pid_t pId, int ExitCode);
 
 	/**
 	 * @brief Yield the current thread
diff --git a/include/driver.h b/include/driver.h
index b072c0cc..ad5c26d4 100644
--- a/include/driver.h
+++ b/include/driver.h
@@ -298,8 +298,9 @@ typedef struct
 	/* Scheduling */
 	pid_t (*CreateKernelProcess)(dev_t MajorID, const char *Name);
 	pid_t (*CreateKernelThread)(dev_t MajorID, pid_t pId, const char *Name, void *EntryPoint, void *Argument);
+	pid_t (*GetCurrentProcess)(dev_t MajorID);
 	int (*KillProcess)(dev_t MajorID, pid_t pId, int ExitCode);
-	int (*KillThread)(dev_t MajorID, pid_t tId, int ExitCode);
+	int (*KillThread)(dev_t MajorID, pid_t tId, pid_t pId, int ExitCode);
 	void (*Yield)(dev_t MajorID);
 	void (*Sleep)(dev_t MajorID, uint64_t Milliseconds);
 
diff --git a/library/task.c b/library/task.c
index de7e65df..515416a2 100644
--- a/library/task.c
+++ b/library/task.c
@@ -35,16 +35,21 @@ pid_t CreateKernelThread(pid_t pId, const char *Name, void *EntryPoint, void *Ar
 								   Argument);
 }
 
+pid_t GetCurrentProcess()
+{
+	return API->GetCurrentProcess(API->MajorID);
+}
+
 int KillProcess(pid_t pId, int ExitCode)
 {
 	return API->KillProcess(API->MajorID,
 							pId, ExitCode);
 }
 
-int KillThread(pid_t tId, int ExitCode)
+int KillThread(pid_t tId, pid_t pId, int ExitCode)
 {
 	return API->KillThread(API->MajorID,
-						   tId, ExitCode);
+						   tId, pId, ExitCode);
 }
 
 void Yield()
diff --git a/misc/vmware/main.c b/misc/vmware/main.c
index 515e1a79..c2a0e7a3 100644
--- a/misc/vmware/main.c
+++ b/misc/vmware/main.c
@@ -568,6 +568,7 @@ static int DisplayGetSize(ToolboxContext *ctx)
 }
 
 pid_t dst_id = -1;
+pid_t dst_pid = -1;
 ToolboxContext *tb_ctx = NULL;
 void DisplayScaleThread()
 {
@@ -686,6 +687,7 @@ int DriverEntry()
 		MessageClose(&tb_ctx);
 		dst_id = CreateKernelThread(0, "VMware Display Scale",
 									(void *)DisplayScaleThread, NULL);
+		dst_pid = GetCurrentProcess();
 	}
 
 	PS2WriteCommand(PS2_CMD_ENABLE_PORT_2);
@@ -723,7 +725,7 @@ int DriverFinal()
 
 	if (ToolboxSupported)
 	{
-		KillThread(dst_id, 0);
+		KillThread(dst_id, dst_pid, 0);
 		if (tb_ctx->TCLOChannel != -1)
 			MessageClose(tb_ctx);
 		FreeMemory(tb_ctx, 1);