Files
UnrealEngine/Engine/Source/Programs/UnrealBuildAccelerator/TestApp/Private/Windows/UbaTestAppWin.cpp
2025-05-18 13:04:45 +08:00

323 lines
11 KiB
C++

// Copyright Epic Games, Inc. All Rights Reserved.
#define WINDOWS_LEAN_AND_MEAN
#include <Windows.h>
#include <wchar.h>
#include <stdarg.h>
#include <time.h>
int LogError(const wchar_t* format, ...)
{
va_list args;
va_start(args, format);
wchar_t buffer[1024];
_vsnwprintf_s(buffer, 1024, _TRUNCATE, format, args);
wprintf(L"%s\n", buffer);
va_end(args);
return -1;
}
int wmain(int argc, wchar_t* argv[])
{
HMODULE detoursHandle = GetModuleHandleW(L"UbaDetours.dll");
using UbaRequestNextProcessFunc = bool(unsigned int prevExitCode, wchar_t* outArguments, unsigned int outArgumentsCapacity);
static UbaRequestNextProcessFunc* requestNextProcess = (UbaRequestNextProcessFunc*)(void*)GetProcAddress(detoursHandle, "UbaRequestNextProcess");
if (argc == 1)
{
if (!detoursHandle)
return LogError(L"Did not find UbaDetours.dll in process!!!\n");
using UbaRunningRemoteFunc = bool();
UbaRunningRemoteFunc* runningRemoteFunc = (UbaRunningRemoteFunc*)GetProcAddress(detoursHandle, "UbaRunningRemote");
if (!runningRemoteFunc)
return LogError(L"Couldn't find UbaRunningRemote function in UbaDetours.dll");
bool runningRemote = (*runningRemoteFunc)();
HMODULE modules[] = { 0, detoursHandle, GetModuleHandleW(L"UbaTestApp.exe") };
for (HMODULE module : modules)
{
DWORD res1 = GetModuleFileNameW(module, NULL, 0);
if (GetLastError() != ERROR_INSUFFICIENT_BUFFER)
return LogError(L"Expected insufficient buffer");
if (res1 != 0)
return LogError(L"Expected zero");
wchar_t name[512];
DWORD realLen = GetModuleFileNameW(module, name, 512);
if (realLen == 0)
return LogError(L"Did not expect this function to fail");
if (GetLastError() == ERROR_INSUFFICIENT_BUFFER)
return LogError(L"Expected sufficient buffer");
name[realLen] = 254;
name[realLen+1] = 254;
DWORD res2 = GetModuleFileNameW(module, name, realLen);
if (GetLastError() != ERROR_INSUFFICIENT_BUFFER)
return LogError(L"Expected insufficient buffer");
if (res2 != realLen)
return LogError(L"Expected to return same as sent in");
if (name[realLen] != 254)
return LogError(L"Overwrite");
if (name[realLen-1] != 0)
return LogError(L"Not terminated");
name[realLen] = 254;
name[realLen + 1] = 254;
DWORD res3 = GetModuleFileNameW(module, name, realLen+1);
if (GetLastError() == ERROR_INSUFFICIENT_BUFFER)
return LogError(L"Expected sufficient buffer");
if (res3 != realLen)
return LogError(L"Expected to return same as sent in");
if (name[realLen+1] != 254)
return LogError(L"Overwrite");
if (name[realLen] != 0)
return LogError(L"Not terminated");
}
wchar_t currentDir[MAX_PATH];
DWORD currentDirLen = GetCurrentDirectoryW(MAX_PATH, currentDir);
if (!currentDirLen)
return LogError(L"GetCurrentDirectoryW failed");
currentDir[currentDirLen] = '\\';
currentDir[currentDirLen + 1] = 0;
wchar_t notepad[] = L"c:\\windows\\system32\\notepad.exe";
wchar_t localNotepad[MAX_PATH];
wcscpy_s(localNotepad, MAX_PATH, currentDir);
wcscat_s(localNotepad, MAX_PATH, L"notepad.exe");
if (!CopyFileW(notepad, localNotepad, false))
return LogError(L"CopyFileW failed");
{
HANDLE fh = CreateFileW(localNotepad, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
if (fh == INVALID_HANDLE_VALUE)
return LogError(L"Failed to open %s for read", localNotepad);
wchar_t path[MAX_PATH];
DWORD res = GetFinalPathNameByHandleW(fh, path, MAX_PATH, 0);
if (!res)
return LogError(L"GetFinalPathNameByHandleW failed");
if (res != wcslen(path))
return LogError(L"GetFinalPathNameByHandleW did not return length of string");
DWORD res2 = GetFinalPathNameByHandleW(fh, path, res, 0);
if (res2 != res + 1)
return LogError(L"GetFinalPathNameByHandleW should return full length plus terminating character");
DWORD res3 = GetFinalPathNameByHandleW(fh, path, res+1, 0);
if (res3 != res)
return LogError(L"GetFinalPathNameByHandleW should return full length plus terminating character");
// TODO: Test character after terminator char
if (!runningRemote)
GetFinalPathNameByHandleW(fh, path, MAX_PATH, VOLUME_NAME_NT); // Testing so it doesn't assert
CloseHandle(fh);
}
{
wchar_t testPath[] = L"R:.";
wchar_t fullPathName[MAX_PATH];
DWORD len = GetFullPathNameW(testPath, MAX_PATH, fullPathName, NULL);
if (len != 3)
return LogError(L"GetFullPathNameW failed");
testPath[0] = currentDir[0];
DWORD len2 = GetFullPathNameW(testPath, MAX_PATH, fullPathName, NULL);
if (len2 != currentDirLen)
return LogError(L"GetFullPathNameW returns length that does not match current dir");
if (memcmp(fullPathName, currentDir, len*sizeof(wchar_t)) != 0)
return LogError(L"GetFullPathNameW returned wrong path");
// TODO: Test character after terminator char
}
{
HANDLE fh = CreateFileW(L"FileW", GENERIC_WRITE, FILE_SHARE_WRITE, NULL, CREATE_ALWAYS, 0, NULL);
if (fh == INVALID_HANDLE_VALUE)
return LogError(L"Failed to create file File");
CloseHandle(fh);
if (!MoveFile(L"FileW", L"FileW2"))
return LogError(L"Failed to move file from FileW to FileW2");
if (!CopyFile(L"FileW2", L"FileWF", false))
return LogError(L"Failed to copy file from FileW2 to FileWF");
}
{
if (!CreateDirectoryW(L"DirA", NULL))
return LogError(L"Failed to create directory");
if (GetFileAttributesW(L"DirA") == 0)
return LogError(L"Failed to get attributes of directory");
if (!RemoveDirectoryW(L"DirA"))
return LogError(L"Failed to remove directory");
if (GetFileAttributesW(L"DirA") != INVALID_FILE_ATTRIBUTES)
return LogError(L"Found attributes of deleted directory");
if (CreateDirectoryW(L"Dir2\\Dir3", NULL))
return LogError(L"Should not succeed creation directory that exists");
if (GetLastError() != ERROR_ALREADY_EXISTS)
return LogError(L"Did not get correct error when failing to create existing directory");
if (GetFileAttributesW(L"Dir2\\Dir3\\Dir4\\Dir5") == INVALID_FILE_ATTRIBUTES)
return LogError(L"Failed to get attributes of directory");
}
{
STARTUPINFOW si;
memset(&si, 0, sizeof(si));
PROCESS_INFORMATION pi;
memset(&pi, 0, sizeof(pi));
wchar_t arg[1024];
wcscpy_s(arg, 1024, argv[0]);
wcscat_s(arg, 1024, L" -child");
if (!CreateProcessW(nullptr, arg, nullptr, nullptr, TRUE, 0, nullptr, nullptr, &si, &pi))
return LogError(L"Failed to create child process");
CloseHandle(pi.hThread);
if (WaitForSingleObject(pi.hProcess, 10000) != WAIT_OBJECT_0)
return LogError(L"Failed waiting for child process");
DWORD exitCode;
if (!GetExitCodeProcess(pi.hProcess, &exitCode) || exitCode)
return LogError(L"Child process failed");
CloseHandle(pi.hProcess);
}
}
else if (wcscmp(argv[1], L"-child") == 0)
{
if (GetFileAttributes(L"FileW2") == INVALID_FILE_ATTRIBUTES)
return LogError(L"Child process could not get attributes of FileW2");
if (GetFileAttributes(L"FileWF") == INVALID_FILE_ATTRIBUTES)
return LogError(L"Child process could not get attributes of FileWF");
if (GetFileAttributes(L"FileW") != INVALID_FILE_ATTRIBUTES)
return LogError(L"Child process found FileW which should not exist anymore");
}
else if (wcscmp(argv[1], L"-reuse") == 0)
{
wchar_t arguments[1024];
if (requestNextProcess(0, arguments, sizeof(arguments)))
return LogError(L"Didn't expect another process");
}
else if (wcsncmp(argv[1], L"-file=", 6) == 0)
{
wchar_t arguments[1024];
const wchar_t* file = argv[1] + 6;
while (true)
{
HANDLE rh = CreateFileW(file, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);
if (rh == INVALID_HANDLE_VALUE)
return LogError(L"Failed to open file %s", file);
char data[17] = {};
DWORD bytesRead;
if (!ReadFile(rh, data, 16, &bytesRead, NULL) || bytesRead != 16)
return LogError(L"Failed to read 16 bytes from file %s", file);
CloseHandle(rh);
srand(GetProcessId(GetCurrentProcess()));
Sleep(rand() % 2000);
wchar_t outFile[1024];
wcscpy_s(outFile, 1024, file);
outFile[wcslen(file)-3] = 0;
wcscat_s(outFile, 1024, L".out");
HANDLE wh = CreateFileW(outFile, GENERIC_WRITE, FILE_SHARE_WRITE, NULL, CREATE_ALWAYS, 0, NULL);
if (wh == INVALID_HANDLE_VALUE)
return LogError(L"Failed to create file File");
data[16] = 1;
DWORD bytesWritten;
if (!WriteFile(wh, data, 17, &bytesWritten, NULL) || bytesWritten != 17)
return LogError(L"Failed to read 16 bytes from file %s", file);
CloseHandle(wh);
// Request new process
if (!requestNextProcess(0, arguments, 1024))
break; // No process available, exit loop
file = arguments + 6;
}
return 0;
}
else if (wcsncmp(argv[1], L"-GetFileAttributes=", 6) == 0)
{
const wchar_t* str = argv[1] + 19;
DWORD attr = GetFileAttributesW(str);
return attr == INVALID_FILE_ATTRIBUTES ? 255 : attr;
}
else if (wcsncmp(argv[1], L"-stdout=", 8) == 0)
{
const wchar_t* str = argv[1] + 8;
if (wcscmp(str, L"rootprocess") == 0)
{
STARTUPINFOW si;
memset(&si, 0, sizeof(si));
PROCESS_INFORMATION pi;
memset(&pi, 0, sizeof(pi));
wchar_t arg[1024];
wcscpy_s(arg, 1024, argv[0]);
wcscat_s(arg, 1024, L" -stdout=childprocess");
//wcscpy_s(arg, 1024, L"\"c:\\sdk\\AutoSDK/HostWin64/Win64/MetalDeveloperTools/4.1/metal/macos/bin/metal.exe\" -v --target=air64-apple-darwin18.7.0 16384");
SECURITY_ATTRIBUTES saAttr;
saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
saAttr.bInheritHandle = TRUE;
saAttr.lpSecurityDescriptor = NULL;
HANDLE readPipe;
HANDLE writePipe;
if (!CreatePipe(&readPipe, &writePipe, &saAttr, 0))
return 1;
if (!SetHandleInformation(readPipe, HANDLE_FLAG_INHERIT, 0))
return 2;
si.dwFlags = STARTF_USESTDHANDLES;
si.hStdInput = GetStdHandle(STD_INPUT_HANDLE);
si.hStdOutput = writePipe;
si.hStdError = writePipe;
DWORD flags = 0;//CREATE_NO_WINDOW;
if (!CreateProcessW(nullptr, arg, nullptr, nullptr, TRUE, flags, nullptr, nullptr, &si, &pi))
return 3;
CloseHandle(pi.hThread);
CloseHandle(writePipe);
char buf[4096] = { 0 };
DWORD readCount = 0;
if (!::ReadFile(readPipe, buf, sizeof(buf) - 1, &readCount, NULL))
{
LogError(L"Failed to read pipe %u %u", GetLastError(), readCount);
return 4;
}
buf[readCount] = 0;
if (strncmp(buf, "childprocess", 12) != 0)
return 5;
if (WaitForSingleObject(pi.hProcess, INFINITE) != WAIT_OBJECT_0)
return 6;
CloseHandle(pi.hProcess);
}
wprintf(L"%s\n", str);
}
else
{
if (!detoursHandle)
return LogError(L"Did not find UbaDetours.dll in process!!!\n");
using u32 = unsigned int;
using UbaSendCustomMessageFunc = u32(const void* send, u32 sendSize, void* recv, u32 recvCapacity);
UbaSendCustomMessageFunc* sendMessage = (UbaSendCustomMessageFunc*)GetProcAddress(detoursHandle, "UbaSendCustomMessage");
if (!sendMessage)
return LogError(L"Couldn't find UbaSendCustomMessage function in UbaDetours.dll");
const wchar_t* helloMsg = L"Hello from client";
wchar_t response[256];
u32 responseSize = (*sendMessage)(helloMsg, u32(wcslen(helloMsg)) * 2, response, 256 * 2);
if (responseSize == 0)
return LogError(L"Didn't get proper response from session");
//wprintf(L"Recv: %.*s\n", responseSize / 2, response);
}
return 0;
}