Use attached executable environment and working directory for compile…

… process (crosire#39)
CyberFlameGO · Feb 6, 2022 · 1a9365f · 1a9365f
1 parent 87a9178
commit 1a9365f
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 4 deletions.
diff --git a/source/blink.cpp b/source/blink.cpp
@@ -57,7 +57,7 @@ blink::application::application()
 	_symbols.insert({ "__ImageBase", _image_base });
 }
 
-void blink::application::run(HANDLE blink_handle)
+void blink::application::run(HANDLE blink_handle, wchar_t* blink_environment, wchar_t* blink_working_directory)
 {
 	std::vector<const BYTE *> dlls;
 
@@ -135,7 +135,12 @@ void blink::application::run(HANDLE blink_handle)
 		TCHAR cmdline[] = TEXT("cmd.exe /q /d /k @echo off");
 		PROCESS_INFORMATION pi;
 
-		if (!CreateProcess(nullptr, cmdline, nullptr, nullptr, TRUE, CREATE_NO_WINDOW, nullptr, nullptr, &si, &pi))
+		// Use blink.exe's environment and working directory for our compiler process.
+		// This way, the user can run blink.exe from their build prompt and our compiler
+		// process will run compiles similar to how they run from the user's prompt.
+		if (!CreateProcess(nullptr, cmdline, nullptr, nullptr, TRUE,
+			CREATE_UNICODE_ENVIRONMENT | CREATE_NO_WINDOW, reinterpret_cast<LPVOID>(blink_environment),
+			blink_working_directory, &si, &pi))
 		{
 			print("  Error: Could not create process.");
 

diff --git a/source/blink.hpp b/source/blink.hpp
@@ -28,7 +28,7 @@ namespace blink
 	public:
 		application();
 
-		void run(HANDLE blink_handle);
+		void run(HANDLE blink_handle, wchar_t* blink_environment, wchar_t* blink_working_directory);
 		bool link(const std::filesystem::path &object_file);
 
 		template <typename T>

diff --git a/source/main.cpp b/source/main.cpp
@@ -27,6 +27,9 @@ extern "C" void _initterm(_PVFV *beg, _PVFV *end);
 
 HANDLE console = INVALID_HANDLE_VALUE;
 char blink_pipe_name[MAX_PATH];
+#define MAX_ENVIRONMENT_LENGTH 32767 // max for Windows
+wchar_t blink_environment[MAX_ENVIRONMENT_LENGTH];
+wchar_t blink_working_directory[MAX_PATH];
 
 void print(const char *message, size_t length)
 {
@@ -135,7 +138,7 @@ DWORD CALLBACK remote_main(BYTE *image_base)
 	scoped_handle blink_handle = CreateFileA(blink_pipe_name, GENERIC_READ | GENERIC_WRITE, 0, NULL, OPEN_ALWAYS, FILE_FLAG_WRITE_THROUGH, NULL);
 
 	// Run main loop
-	blink::application().run(blink_handle);
+	blink::application().run(blink_handle, blink_environment, blink_working_directory);
 
 	CloseHandle(console);
 
@@ -299,6 +302,26 @@ int main(int argc, char *argv[])
 		return GetLastError();
 	}
 
+	// Get blink.exe's environment, so we can use it for compiles in the remote thread.
+	const LPWCH environment = GetEnvironmentStringsW();
+	LPWCH environment_end = environment;
+	for (; *environment_end != '\0'; environment_end += wcslen(environment_end) + 1)
+	{
+		// continue until we find the double null end
+	}
+	environment_end += 1; // get position after the double null end
+	const size_t environment_length = environment_end - environment;
+	if (environment_length > MAX_ENVIRONMENT_LENGTH)
+	{
+		std::cout << "Environment string exceeded max length!" << std::endl;
+		return ERROR_NOT_ENOUGH_MEMORY;
+	}
+	memcpy(blink_environment, environment, environment_length * sizeof(wchar_t));
+	FreeEnvironmentStringsW(environment);
+
+	// Get blink.exe's working directory, so we can use it for compiles in the remote thread.
+	wcscpy_s(blink_working_directory, std::filesystem::current_path().c_str());
+
 	// Copy current module image to target application (including the IAT and value of the global variables)
 	if (remote_baseaddress == nullptr || !WriteProcessMemory(remote_process, remote_baseaddress, module_info.lpBaseOfDll, module_info.SizeOfImage, nullptr))
 	{