#include "stdafx.h"

#include "Startup.h"

#include <ShellAPI.h>

Startup::Startup()

{

hJob = NULL;

ZeroMemory(hChildProcesses, sizeof(hChildProcesses));

ZeroMemory(childThreadIds, sizeof(childThreadIds));

ZeroMemory(hookTarget, sizeof(hookTarget));

hookWindow = NULL;

childProcessCount = 0;

flags = NONE;

}

Startup::~Startup()

{

for (DWORD index = 0; index < childProcessCount; index++)

{

PostThreadMessage(childThreadIds[index], MT_APP_MSG_QUIT, 0, 0);

}

for (DWORD index = 0; index < childProcessCount; index++)

{

WaitForSingleObject(hChildProcesses[index], 5000);

CloseHandle(hChildProcesses[index]);

}

if (hJob)

{

CloseHandle(hJob);

}

}

BOOL Startup::ParseFlags(LPCWSTR lpCmdLine)

{

LPWSTR *szArglist;

int nArgs;

szArglist = CommandLineToArgvW(lpCmdLine, &nArgs);

if (szArglist == NULL)

{

return FALSE;

}

// Parse command line

for (int i = 0; i < nArgs; i++)

{

if (wcscmp(szArglist[i], L"-hide") == 0)

{

flags |= Flags::HIDE_TRAY;

}

else if (wcscmp(szArglist[i], L"-hook-host") == 0)

{

flags |= Flags::HOOK_HOST;

}

else if (wcscmp(szArglist[i], L"-hook-target") == 0 && i + 1 < nArgs)

{

_tcsncpy_s(hookTarget, szArglist[++i], _TRUNCATE);

}

else if (wcscmp(szArglist[i], L"-hook-window") == 0 && i + 1 < nArgs)

{

hookWindow = reinterpret_cast<HWND>(static_cast<ULONG_PTR>(

_wcstoui64(szArglist[++i], NULL, 0)));

}

}

// Free memory

LocalFree(szArglist);

return TRUE;

}

BOOL Startup::HasHookTarget() const

{

return hookTarget[0] != _T('\0') || hookWindow != NULL;

}

LPCTSTR Startup::GetHookTarget() const

{

return hookTarget;

}

HWND Startup::GetHookWindow() const

{

return hookWindow;

}

BOOL Startup::IsHookHost() const

{

return (flags & Flags::HOOK_HOST) != 0;

}

BOOL Startup::ShouldShowTray() const

{

return !IsHookHost() && !HasHookTarget();

}

Startup::StartupResult Startup::Start()

{

if (IsHookHost() || HasHookTarget())

{

return STARTUP_RUN;

}

USHORT nativeMachine;

if (!GetNativeMachine(&nativeMachine))

{

return STARTUP_FAILED;

}

#ifdef _M_ARM64

const USHORT processMachine = IMAGE_FILE_MACHINE_ARM64;

#elif defined(_WIN64)

const USHORT processMachine = IMAGE_FILE_MACHINE_AMD64;

#else

const USHORT processMachine = IMAGE_FILE_MACHINE_I386;

#endif

if (processMachine != nativeMachine)

{

return LaunchNativeCoordinator(nativeMachine)

? STARTUP_REDIRECTED

: STARTUP_FAILED;

}

return CreateJob() ? STARTUP_RUN : STARTUP_FAILED;

}

BOOL Startup::GetNativeMachine(USHORT* nativeMachine) const

{

typedef BOOL(WINAPI *IsWow64Process2Proc)(HANDLE, USHORT*, USHORT*);

IsWow64Process2Proc isWow64Process2 = reinterpret_cast<IsWow64Process2Proc>(

GetProcAddress(GetModuleHandle(_T("kernel32.dll")), "IsWow64Process2"));

if (isWow64Process2)

{

USHORT processMachine;

return isWow64Process2(GetCurrentProcess(), &processMachine, nativeMachine);

}

SYSTEM_INFO nativeSystemInfo;

GetNativeSystemInfo(&nativeSystemInfo);

switch (nativeSystemInfo.wProcessorArchitecture)

{

case PROCESSOR_ARCHITECTURE_INTEL:

*nativeMachine = IMAGE_FILE_MACHINE_I386;

return TRUE;

case PROCESSOR_ARCHITECTURE_AMD64:

*nativeMachine = IMAGE_FILE_MACHINE_AMD64;

return TRUE;

case PROCESSOR_ARCHITECTURE_ARM64:

*nativeMachine = IMAGE_FILE_MACHINE_ARM64;

return TRUE;

}

return FALSE;

}

BOOL Startup::GetSiblingExecutablePath(LPCTSTR executableName,

LPTSTR executablePath, DWORD executablePathLength) const

{

DWORD pathLength = GetModuleFileName(NULL, executablePath, executablePathLength);

if (!pathLength || pathLength >= executablePathLength)

{

return FALSE;

}

LPTSTR fileName = _tcsrchr(executablePath, _T('\\'));

fileName = fileName ? fileName + 1 : executablePath;

size_t directoryLength = fileName - executablePath;

if (directoryLength + _tcslen(executableName) + 1 > executablePathLength)

{

return FALSE;

}

_tcscpy_s(fileName, executablePathLength - directoryLength, executableName);

return TRUE;

}

BOOL Startup::LaunchNativeCoordinator(USHORT nativeMachine) const

{

LPCTSTR executableName;

switch (nativeMachine)

{

case IMAGE_FILE_MACHINE_I386:

executableName = MT_EXE_NAME;

break;

case IMAGE_FILE_MACHINE_AMD64:

executableName = MT_EXE_NAME64;

break;

case IMAGE_FILE_MACHINE_ARM64:

executableName = MT_EXE_NAMEARM64;

break;

default:

return FALSE;

}

TCHAR executablePath[MAX_PATH];

if (!GetSiblingExecutablePath(executableName, executablePath, ARRAYSIZE(executablePath)))

{

return FALSE;

}

TCHAR commandLine[MAX_PATH + 16];

_stprintf_s(commandLine, _T("\"%s\"%s"), executablePath,

(flags & Flags::HIDE_TRAY) ? _T(" -hide") : _T(""));

STARTUPINFO startupInfo = { 0 };

startupInfo.cb = sizeof(STARTUPINFO);

PROCESS_INFORMATION processInfo = { 0 };

if (!CreateProcess(executablePath, commandLine, NULL, NULL, FALSE,

0, NULL, NULL, &startupInfo, &processInfo))

{

return FALSE;

}

CloseHandle(processInfo.hThread);

CloseHandle(processInfo.hProcess);

return TRUE;

}

BOOL Startup::CreateJob()

{

// Create a job

hJob = CreateJobObject(NULL, MT_JOB_NAME);

if (!hJob)

{

return FALSE;

}

// If the job already exists

if (GetLastError() == ERROR_ALREADY_EXISTS)

{

CloseHandle(hJob);

hJob = NULL;

return FALSE;

}

JOBOBJECT_EXTENDED_LIMIT_INFORMATION jeli = { 0 };

jeli.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE;

if (!SetInformationJobObject(hJob, JobObjectExtendedLimitInformation,

&jeli, sizeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION)))

{

return FALSE;

}

#ifdef _M_ARM64

#if !MT_DEBUG_ONLY_X86

if (!CreateJobChild(MT_EXE_NAME))

{

return FALSE;

}

#endif

#if !MT_DEBUG_ONLY_X64

if (!CreateJobChild(MT_EXE_NAME64))

{

return FALSE;

}

#endif

#elif defined(_WIN64)

#if !MT_DEBUG_ONLY_X86

if (!CreateJobChild(MT_EXE_NAME))

{

return FALSE;

}

#endif

#endif

return TRUE;

}

BOOL Startup::CreateJobChild(LPCTSTR executableName)

{

TCHAR executablePath[MAX_PATH];

if (!GetSiblingExecutablePath(executableName, executablePath, ARRAYSIZE(executablePath)))

{

return FALSE;

}

TCHAR commandLine[MAX_PATH + 16];

_stprintf_s(commandLine, _T("\"%s\" -hook-host"), executablePath);

STARTUPINFO startupInfo = { 0 };

startupInfo.cb = sizeof(STARTUPINFO);

PROCESS_INFORMATION processInfo = { 0 };

// Create the companion process in a suspended state.

if (!CreateProcess(executablePath, commandLine, NULL, NULL, FALSE,

CREATE_SUSPENDED,

NULL, NULL, &startupInfo, &processInfo))

{

return FALSE;

}

// Assigns a process to an existing job object

if (!AssignProcessToJobObject(hJob, processInfo.hProcess))

{

TerminateProcess(processInfo.hProcess, 1);

CloseHandle(processInfo.hThread);

CloseHandle(processInfo.hProcess);

return FALSE;

}

// Execution of the thread is resumed

if (ResumeThread(processInfo.hThread) == (DWORD)-1)

{

TerminateProcess(processInfo.hProcess, 1);

CloseHandle(processInfo.hThread);

CloseHandle(processInfo.hProcess);

return FALSE;

}

CloseHandle(processInfo.hThread);

if (childProcessCount < ARRAYSIZE(hChildProcesses))

{

hChildProcesses[childProcessCount] = processInfo.hProcess;

childThreadIds[childProcessCount] = processInfo.dwThreadId;

childProcessCount++;

}

else

{

CloseHandle(processInfo.hProcess);

}

return TRUE;

}