From c88b65726f727524d4188c80b23af99b3115b2f2 Mon Sep 17 00:00:00 2001
From: midipix <writeonce@midipix.org>
Date: Thu, 29 Aug 2024 02:02:20 +0000
Subject: __ntapi_tt_get_argv_envp_utf16(): guard against repeated vector
 relocations.

---
 src/argv/ntapi_tt_argv_envp.c                   | 58 +++++++++++++------------
 src/process/ntapi_tt_create_native_process_v1.c |  6 +++
 src/process/ntapi_tt_create_native_process_v2.c |  4 ++
 3 files changed, 41 insertions(+), 27 deletions(-)

(limited to 'src')

diff --git a/src/argv/ntapi_tt_argv_envp.c b/src/argv/ntapi_tt_argv_envp.c
index 11a9f67..9c769ac 100644
--- a/src/argv/ntapi_tt_argv_envp.c
+++ b/src/argv/ntapi_tt_argv_envp.c
@@ -517,44 +517,48 @@ int32_t __stdcall __ntapi_tt_get_argv_envp_utf16(
 		rtdata->peb_wenvp = main_params->wenvp_buffer;
 
 		/* integral wargv, wenvp, argv, envp */
-		if (rtdata->wargv) {
-			rtdata->wargv += (uintptr_t)rtdata / sizeof(wchar16_t *);
+		if (!(rtdata->flags & NT_RUNTIME_DATA_ARGV_ENVP_RELOCATED)) {
+			if (rtdata->wargv) {
+				rtdata->wargv += (uintptr_t)rtdata / sizeof(wchar16_t *);
 
-			for (wch_p=rtdata->wargv; *wch_p; wch_p++)
-				*wch_p += (uintptr_t)rtdata / sizeof(wchar16_t);
-		};
+				for (wch_p=rtdata->wargv; *wch_p; wch_p++)
+					*wch_p += (uintptr_t)rtdata / sizeof(wchar16_t);
+			};
 
-		if (rtdata->wenvp) {
-			rtdata->wenvp += (uintptr_t)rtdata / sizeof(wchar16_t *);
+			if (rtdata->wenvp) {
+				rtdata->wenvp += (uintptr_t)rtdata / sizeof(wchar16_t *);
 
-			for (wch_p=rtdata->wenvp; *wch_p; wch_p++)
-				*wch_p += (uintptr_t)rtdata / sizeof(wchar16_t);
-		}
+				for (wch_p=rtdata->wenvp; *wch_p; wch_p++)
+					*wch_p += (uintptr_t)rtdata / sizeof(wchar16_t);
+			}
 
-		if (rtdata->argv) {
-			rtdata->argv += (uintptr_t)rtdata / sizeof(char *);
+			if (rtdata->argv) {
+				rtdata->argv += (uintptr_t)rtdata / sizeof(char *);
 
-			for (ch_p=rtdata->argv; *ch_p; ch_p++)
-				*ch_p += (uintptr_t)rtdata;
+				for (ch_p=rtdata->argv; *ch_p; ch_p++)
+					*ch_p += (uintptr_t)rtdata;
 
-			rtdata->argc = (int32_t)(ch_p - rtdata->argv);
-		};
+				rtdata->argc = (int32_t)(ch_p - rtdata->argv);
+			};
 
-		if (rtdata->envp) {
-			rtdata->envp += (uintptr_t)rtdata / sizeof(char *);
+			if (rtdata->envp) {
+				rtdata->envp += (uintptr_t)rtdata / sizeof(char *);
 
-			for (ch_p=rtdata->envp; *ch_p; ch_p++)
-				*ch_p += (uintptr_t)rtdata;
+				for (ch_p=rtdata->envp; *ch_p; ch_p++)
+					*ch_p += (uintptr_t)rtdata;
 
-			rtdata->envc = (int32_t)(ch_p - rtdata->envp);
-		};
+				rtdata->envc = (int32_t)(ch_p - rtdata->envp);
+			};
 
-		if (rtdata->w32_envp) {
-			rtdata->w32_envp += (uintptr_t)rtdata / sizeof(char *);
+			if (rtdata->w32_envp) {
+				rtdata->w32_envp += (uintptr_t)rtdata / sizeof(char *);
 
-			for (ch_p=rtdata->w32_envp; *ch_p; ch_p++)
-				*ch_p += (uintptr_t)rtdata;
-		};
+				for (ch_p=rtdata->w32_envp; *ch_p; ch_p++)
+					*ch_p += (uintptr_t)rtdata;
+			};
+
+			rtdata->flags |= NT_RUNTIME_DATA_ARGV_ENVP_RELOCATED;
+		}
 	}
 
 	/* we're good */
diff --git a/src/process/ntapi_tt_create_native_process_v1.c b/src/process/ntapi_tt_create_native_process_v1.c
index 57ca36f..12f0073 100644
--- a/src/process/ntapi_tt_create_native_process_v1.c
+++ b/src/process/ntapi_tt_create_native_process_v1.c
@@ -238,6 +238,12 @@ int32_t __stdcall __ntapi_tt_create_native_process_v1(nt_create_process_params *
 		return __tt_create_process_cancel(params,hsection,status);
 
 	/* rtdata */
+	if (params->rtblock) {
+		struct _nt_runtime_data * rtdata;
+		rtdata =  (struct _nt_runtime_data *)params->rtblock->addr;
+		rtdata->flags &= ~NT_RUNTIME_DATA_ARGV_ENVP_RELOCATED;
+	}
+
 	if (params->rtblock && (status = __ntapi_tt_create_remote_runtime_data(tparams.hprocess,params->rtblock)))
 		return __tt_create_process_cancel(params,hsection,status);
 
diff --git a/src/process/ntapi_tt_create_native_process_v2.c b/src/process/ntapi_tt_create_native_process_v2.c
index af683de..7214dc5 100644
--- a/src/process/ntapi_tt_create_native_process_v2.c
+++ b/src/process/ntapi_tt_create_native_process_v2.c
@@ -224,6 +224,10 @@ int32_t __stdcall __ntapi_tt_create_native_process_v2(
 		return NT_STATUS_SUCCESS;
 
 	/* rtdata */
+	struct _nt_runtime_data * rtdata;
+	rtdata =  (struct _nt_runtime_data *)params->rtblock->addr;
+	rtdata->flags &= ~NT_RUNTIME_DATA_ARGV_ENVP_RELOCATED;
+
 	if ((status = __ntapi_tt_create_remote_runtime_data(params->hprocess,params->rtblock)))
 		return __tt_create_process_cancel(params,status);
 
-- 
cgit v1.2.3