DSProxy/SProxy/CProcessManager.cpp

#include "stdafx.h"
#include "CProcessManager.h"
CProcessManager* SSingleton<CProcessManager>::ms_Singleton = NULL;

CProcessManager::CProcessManager()
{

}

CProcessManager::~CProcessManager()
{

}

void CProcessManager::_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"(")");
}

void CProcessManager::SetInsTanCe(HINSTANCE m) {
	mInstance = m;
}

 void CProcessManager::SetArgs(std::filesystem::path exePath, std::filesystem::path homeDir, std::filesystem::path cconfigFile)
{
	_exePath = exePath;
	_homeDir = homeDir;
	_configFile = cconfigFile;
	_UpdateClashCmd();
}

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

 bool CProcessManager::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 size = (std::max)((_exePath.native().size() + 1) + (_clashCmd.size() + 1) * sizeof(wchar_t), sizeof(ProcessInfo));

		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;
			ZeroMemory(&si, sizeof(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;
}

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

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

 int CProcessManager::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;
}