__ntapi_tt_get_argv_envp_utf16(): guard against repeated vector relocations.

4 files changed, 42 insertions, 27 deletions

diff --git a/include/ntapi/nt_process.h b/include/ntapi/nt_process.h
index 5abec1a..4e7bd92 100644
--- a/include/ntapi/nt_process.h
+++ b/include/ntapi/nt_process.h
@@ -237,6 +237,7 @@ typedef enum _nt_process_opcode_idx {
 #define NT_RUNTIME_DATA_DEBUG_NATIVE_CHILDREN			(0x04)
 #define NT_RUNTIME_DATA_DEBUG_FOREIGN_CHILDREN			(0x08)
 #define NT_RUNTIME_DATA_TTY_TOP_LEVEL				(0x10)
+#define NT_RUNTIME_DATA_ARGV_ENVP_RELOCATED			(0x20)
 
 /* runtime data convenience storage */
 #define NT_RUNTIME_DATA_SYNC_OPCODES				(0x08)
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);