wltProxy/SProxy/ProcessManager.hpp

#pragma once
# ifndef PROCESS_H
# define PROCESS_H


#define OBJPREFIX LR"(Local\)"
constexpr size_t guidSize = std::size(L"{00000000-0000-0000-0000-000000000000}"); // including the null terminator
constexpr size_t prefixSize = std::size(OBJPREFIX) - 1;
constexpr size_t objNameSize = prefixSize + guidSize;

struct ProcessInfo
{
	DWORD processId;
	DWORD threadId;
	HWND hWndConsole;
};

namespace ProcessManager
{
	enum struct State
	{
		Stop,
		Running,
		WaitStop
	};

	namespace {
		State _state = State::Stop;
		std::filesystem::path _exePath;
		std::filesystem::path _homeDir;
		std::filesystem::path _configFile;
		std::wstring _ctlAddr, _ctlSecret, _clashCmd;
		wil::unique_handle _hJob;
		wil::unique_process_information _subProcInfo, _clashProcInfo;
		HWND _hWndConsole = nullptr;
		wil::unique_handle _hEvent;
		HINSTANCE mInstance;
 
		void _UpdateClashCmd()
		{
			_clashCmd.assign(LR"(")");
			_clashCmd.append(_exePath.filename());
			_clashCmd.append(LR"(")");
			if (!_homeDir.empty())
			{
				_clashCmd.append(LR"( -d ")");
				_clashCmd.append(_homeDir);
				_clashCmd.append(LR"(")");
			}
			if (!_configFile.empty())
			{
				_clashCmd.append(LR"( -f ")");
				_clashCmd.append(_configFile);
				_clashCmd.append(LR"(")");
			}
			//if (!_uiDir.empty())
			//{
			//	_clashCmd.append(LR"( -ext-ui ")");
			//	_clashCmd.append(_uiDir);
			//	_clashCmd.append(LR"(")");
			//}
			//if (!_ctlAddr.empty())
			//{
			//	_clashCmd.append(LR"( -ext-ctl ")");
			//	_clashCmd.append(_ctlAddr);
			//	_clashCmd.append(LR"(")");
			//}

			//// Override secret even if empty
			//_clashCmd.append(LR"( -secret ")");
			//_clashCmd.append(_ctlSecret);
			//_clashCmd.append(LR"(")");
		}

	}

	inline void SetInsTanCe(HINSTANCE m) {
		mInstance = m;
	}

	 void SetArgs(std::filesystem::path exePath, std::filesystem::path homeDir, std::filesystem::path cconfigFile)
	{

	 

		_exePath = exePath;
		_homeDir = homeDir;
		_configFile = cconfigFile;
		//_uiDir = "";
		//_ctlAddr = L"";
		//_ctlSecret = L"";

		_UpdateClashCmd();
	}

	 void SetConfigFile(std::filesystem::path configFile)
	{
		_configFile = configFile;

		_UpdateClashCmd();
	}

	inline State IsRunning() { return _state; }
	inline const PROCESS_INFORMATION& GetSubProcessInfo() { return _subProcInfo; }
	inline const PROCESS_INFORMATION& GetClashProcessInfo() { return _clashProcInfo; }
	inline HWND GetConsoleWindow() { return _hWndConsole; }

	inline bool Start()
	{
		if (_state != State::Stop)
			return false;

		try
		{
			_hJob.reset(CreateJobObjectW(nullptr, nullptr));
			THROW_LAST_ERROR_IF_NULL(_hJob);
			JOBOBJECT_BASIC_LIMIT_INFORMATION bli;
			bli.LimitFlags = JOB_OBJECT_LIMIT_DIE_ON_UNHANDLED_EXCEPTION | JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;
			JOBOBJECT_EXTENDED_LIMIT_INFORMATION eli;
			eli.BasicLimitInformation = bli;
			THROW_IF_WIN32_BOOL_FALSE(SetInformationJobObject(_hJob.get(), JobObjectExtendedLimitInformation, &eli, sizeof(eli)));
			
			wchar_t objName[objNameSize + 1] = OBJPREFIX;
			GUID guid = {};
			THROW_IF_FAILED(CoCreateGuid(&guid));
			THROW_HR_IF(E_OUTOFMEMORY, StringFromGUID2(guid, objName + prefixSize, guidSize) != guidSize);

			size_t s = (_exePath.native().size() + 1) + (_clashCmd.size() + 1) * sizeof(wchar_t);
			size_t l = sizeof(ProcessInfo);
			size_t size = max(s, l);

			objName[objNameSize - 1] = L'F';
			wil::unique_handle hFileMapping(CreateFileMappingW(INVALID_HANDLE_VALUE, nullptr, PAGE_READWRITE, 0, static_cast<DWORD>(size), objName));
			THROW_LAST_ERROR_IF_NULL(hFileMapping);
			auto error = GetLastError();
			if (error == ERROR_ALREADY_EXISTS) THROW_WIN32(error);

			wil::unique_mapview_ptr<wchar_t> buffer(reinterpret_cast<wchar_t*>(MapViewOfFile(hFileMapping.get(), FILE_MAP_WRITE | FILE_MAP_READ, 0, 0, 0)));
			THROW_LAST_ERROR_IF_NULL(buffer);


			objName[objNameSize - 1] = L'E';
			wil::unique_handle hEvent(CreateEventW(nullptr, FALSE, FALSE, objName));
			THROW_LAST_ERROR_IF_NULL(hEvent);
			error = GetLastError();
			if (error == ERROR_ALREADY_EXISTS) THROW_WIN32(error);

			auto exePathPtr = buffer.get();
			size_t i = _exePath.native().copy(exePathPtr, std::filesystem::path::string_type::npos);
			exePathPtr[i] = 0;

			auto cmdPtr = buffer.get() + i + 1;
			i = _clashCmd.copy(cmdPtr, std::wstring::npos);
			cmdPtr[i] = 0;


			auto selfPath = GetModuleFsPath(mInstance);
			auto guidStr = objName + prefixSize;
			std::wstring cmd(LR"(")");
			cmd.append(selfPath);
			cmd.append(LR"(" --pm=)");
			cmd.append(guidStr, guidSize - 1);

			{
				// Disable Windows Error Reporting for subprocess
				auto lastErrorMode = SetErrorMode(SEM_FAILCRITICALERRORS | SEM_NOGPFAULTERRORBOX);
				auto restoreErrorMode = wil::scope_exit([=]() { // Restore error mode after CreateProcessW
					SetErrorMode(lastErrorMode);
					});
				wil::unique_cotaskmem_string appId;
				THROW_IF_FAILED(GetCurrentProcessExplicitAppUserModelID(&appId));

				/*wil::unique_cotaskmem_string appId;
				THROW_IF_FAILED(GetCurrentProcessExplicitAppUserModelID(&appId));*/
				/*STARTUPINFOW si = {
					.cb = sizeof(si),
					.lpTitle = appId.get(),
					.dwFlags = STARTF_FORCEOFFFEEDBACK | STARTF_PREVENTPINNING | STARTF_TITLEISAPPID | STARTF_USESHOWWINDOW,
					.wShowWindow = SW_HIDE
				};*/

				STARTUPINFOW si;
				si.cb = sizeof(si);
				si.lpTitle = appId.get();
				si.dwFlags = STARTF_FORCEOFFFEEDBACK | STARTF_PREVENTPINNING | STARTF_TITLEISAPPID | STARTF_USESHOWWINDOW;
				si.wShowWindow = SW_HIDE;
				THROW_IF_WIN32_BOOL_FALSE(CreateProcessW(selfPath.c_str(), cmd.data(), nullptr, nullptr, FALSE, CREATE_SUSPENDED, nullptr, nullptr, &si, &_subProcInfo));
			}

			// Ensure process in job before start
			THROW_IF_WIN32_BOOL_FALSE(AssignProcessToJobObject(_hJob.get(), _subProcInfo.hProcess));
			THROW_LAST_ERROR_IF(ResumeThread(_subProcInfo.hThread) == static_cast<DWORD>(-1));

			HANDLE handles[] = { hEvent.get(), _subProcInfo.hProcess };
			auto ret = WaitForMultipleObjects(static_cast<DWORD>(std::size(handles)), handles, FALSE, INFINITE);
			error = ERROR_TIMEOUT;

			switch (ret)
			{
			case WAIT_OBJECT_0: // Event signaled, clash process started suspended
			{
				auto info = reinterpret_cast<ProcessInfo*>(buffer.get());
				_clashProcInfo.dwProcessId = info->processId;
				_clashProcInfo.hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, _clashProcInfo.dwProcessId);
				THROW_LAST_ERROR_IF_NULL(_clashProcInfo.hProcess);

				_clashProcInfo.dwThreadId = info->threadId;
				_clashProcInfo.hThread = OpenThread(THREAD_ALL_ACCESS, FALSE, _clashProcInfo.dwThreadId);
				THROW_LAST_ERROR_IF_NULL(_clashProcInfo.hThread);

				_hWndConsole = info->hWndConsole;
			}
			break;
			case WAIT_OBJECT_0 + 1: // Sub process exited before event signaled
			{
				DWORD exitCode;
				THROW_IF_WIN32_BOOL_FALSE(GetExitCodeProcess(_subProcInfo.hProcess, &exitCode));
				HRESULT hr = static_cast<HRESULT>(exitCode); // Treat exit code as hresult
				if (SUCCEEDED(hr)) hr = E_UNEXPECTED;
				THROW_HR(hr);
			}
			return false;
			case WAIT_FAILED:
				error = GetLastError();
				[[fallthrough]];
			case WAIT_TIMEOUT:
				THROW_WIN32(error);
				return false;
			}
			_hEvent = std::move(hEvent);
			ResumeThread(_clashProcInfo.hThread);
		}
		catch (...)
		{
			_hJob.reset();
			LOG_CAUGHT_EXCEPTION();
			return false;
		}
		_state = State::Running;
		return true;
	}


	inline void Stop()
	{
		if (_state != State::Stop)
		{
			_state = State::Stop;
			_hJob.reset();
			_subProcInfo.reset();
			_clashProcInfo.reset();
			_hWndConsole = nullptr;
			_hEvent.reset();
		}
	}

	inline void SendStopSignal()
	{
		if (_state == State::Running)
		{
			_state = State::WaitStop;
			SetEvent(_hEvent.get());
		}
	}


	inline int SubProcess(std::wstring_view guid)
	{
		try
		{
			wchar_t objName[objNameSize + 1] = OBJPREFIX;
			size_t i = guid.copy(objName + prefixSize, guidSize - 1);
			objName[prefixSize + i] = L'F';

			wil::unique_handle hFileMapping(OpenFileMappingW(FILE_MAP_WRITE | FILE_MAP_READ, FALSE, objName));
			THROW_LAST_ERROR_IF_NULL(hFileMapping);

			wil::unique_mapview_ptr<wchar_t> buffer(reinterpret_cast<wchar_t*>(MapViewOfFile(hFileMapping.get(), FILE_MAP_WRITE | FILE_MAP_READ, 0, 0, 0)));
			THROW_LAST_ERROR_IF_NULL(buffer);

			objName[prefixSize + i] = L'E';
			wil::unique_handle hEvent(OpenEventW(EVENT_MODIFY_STATE | SYNCHRONIZE, FALSE, objName));
			THROW_LAST_ERROR_IF_NULL(hEvent);

			THROW_IF_WIN32_BOOL_FALSE(AllocConsole());
			HWND hWndConsole = ::GetConsoleWindow();
			THROW_LAST_ERROR_IF_NULL(hWndConsole);
			ShowWindow(hWndConsole, SW_HIDE);

			auto exePath = buffer.get();
			auto clashCmd = buffer.get() + wcslen(exePath) + 1;
		 
			STARTUPINFOW si;
			si.cb = sizeof(si);
			wil::unique_process_information procInfo;
			THROW_IF_WIN32_BOOL_FALSE(CreateProcessW(exePath, clashCmd, nullptr, nullptr, FALSE, CREATE_SUSPENDED, nullptr, nullptr, &si, &procInfo));

			auto info = reinterpret_cast<ProcessInfo*>(buffer.get());
			/**info = {
				.processId = procInfo.dwProcessId,
				.threadId = procInfo.dwThreadId,
				.hWndConsole = hWndConsole
			};*/

			info->processId = procInfo.dwProcessId;
			info->threadId = procInfo.dwThreadId;
			info->hWndConsole = hWndConsole;


			SetConsoleCtrlHandler(nullptr, TRUE); // Ignores Ctrl+C

			THROW_IF_WIN32_BOOL_FALSE(SetEvent(hEvent.get()));

			while (true)
			{
				HANDLE handles[] = { hEvent.get(), procInfo.hProcess };
				auto ret = WaitForMultipleObjects(static_cast<DWORD>(std::size(handles)), handles, FALSE, INFINITE);
				if (ret == WAIT_OBJECT_0) // Event signaled
				{
					GenerateConsoleCtrlEvent(CTRL_C_EVENT, 0);
				}
				else if (ret == WAIT_OBJECT_0 + 1) // Clash process exited
				{
					break;
				}
				else if (ret == WAIT_FAILED)
				{
					THROW_LAST_ERROR();
				}
				else if (ret == WAIT_TIMEOUT)
				{
					THROW_WIN32(ERROR_TIMEOUT);
				}
			}

			const std::wstring_view msg = (
				L"\n"
				L"[Process completed]\n");
			THROW_IF_WIN32_BOOL_FALSE(WriteConsoleW(GetStdHandle(STD_OUTPUT_HANDLE), msg.data(), static_cast<DWORD>(msg.size()), nullptr, nullptr));
			static_cast<void>(_getch());
		}
		CATCH_RETURN();
		return S_OK;
	}
}


#undef OBJPREFIX

# endif