1 files changed, 158 insertions, 0 deletions
diff --git a/libs/usvfs/test/tinjectlib_test/main.cpp b/libs/usvfs/test/tinjectlib_test/main.cpp
new file mode 100644
index 0000000..ce48f83
--- /dev/null
+++ b/libs/usvfs/test/tinjectlib_test/main.cpp
@@ -0,0 +1,158 @@
+#include <boost/filesystem.hpp>
+#include <boost/type_traits.hpp>
+#include <gtest/gtest.h>
+#include <injectlib.h>
+#include <spdlog/sinks/stdout_sinks.h>
+#include <spdlog/spdlog.h>
+#include <winapi.h>
+
+using namespace usvfs::shared;
+using namespace InjectLib;
+
+#if BOOST_ARCH_X86_64
+static const wchar_t INJECT_BIN[] = L"testinject_bin_x64.exe";
+#else
+static const wchar_t INJECT_BIN[] = L"testinject_bin_x86.exe";
+#endif
+
+#if BOOST_ARCH_X86_64
+static const wchar_t INJECT_LIB[] = L"testinject_dll_x64.dll";
+#else
+static const wchar_t INJECT_LIB[] = L"testinject_dll_x86.dll";
+#endif
+
+static std::shared_ptr<spdlog::logger> logger()
+{
+  std::shared_ptr<spdlog::logger> result = spdlog::get("test");
+  if (result.get() == nullptr) {
+    result = spdlog::stdout_logger_mt("test");
+  }
+  return result;
+}
+
+bool spawn(HANDLE& processHandle, HANDLE& threadHandle)
+{
+  STARTUPINFO si;
+  ::ZeroMemory(&si, sizeof(si));
+  si.cb = sizeof(si);
+
+  PROCESS_INFORMATION pi;
+  BOOL success = ::CreateProcess(INJECT_BIN, nullptr, nullptr, nullptr, FALSE,
+                                 CREATE_SUSPENDED, nullptr, nullptr, &si, &pi);
+
+  if (!success) {
+    throw windows_error("failed to start process");
+  }
+
+  processHandle = pi.hProcess;
+  threadHandle  = pi.hThread;
+
+  return true;
+}
+
+TEST(InjectingTest, InjectionNoInit)
+{
+  // Verify lib can inject without a init function
+
+  HANDLE process, thread;
+  spawn(process, thread);
+  EXPECT_NO_THROW(InjectLib::InjectDLL(process, thread, INJECT_LIB));
+  ResumeThread(thread);
+
+  DWORD res      = WaitForSingleObject(process, INFINITE);
+  DWORD exitCode = NO_ERROR;
+  res            = GetExitCodeProcess(process, &exitCode);
+  EXPECT_EQ(NOERROR, exitCode);
+
+  CloseHandle(process);
+  CloseHandle(thread);
+}
+
+TEST(InjectingTest, InjectionSimpleInit)
+{
+  // Verify lib can inject with a init function with null parameters
+
+  HANDLE process, thread;
+  spawn(process, thread);
+  EXPECT_NO_THROW(InjectLib::InjectDLL(process, thread, INJECT_LIB, "InitNoParam"));
+  ResumeThread(thread);
+
+  DWORD res      = WaitForSingleObject(process, INFINITE);
+  DWORD exitCode = NO_ERROR;
+  res            = GetExitCodeProcess(process, &exitCode);
+  EXPECT_EQ(10001, exitCode);  // used init function exits process with this exit code
+
+  CloseHandle(process);
+  CloseHandle(thread);
+}
+
+TEST(InjectingTest, InjectionComplexInit)
+{
+  // Verify lib can inject with a init function with null parameters
+
+  static const WCHAR param[] = L"magic_parameter";
+  HANDLE process, thread;
+  spawn(process, thread);
+  EXPECT_NO_THROW(InjectLib::InjectDLL(process, thread, INJECT_LIB, "InitComplexParam",
+                                       reinterpret_cast<LPCVOID>(param),
+                                       wcslen(param) * sizeof(WCHAR)));
+
+  ResumeThread(thread);
+
+  DWORD res      = WaitForSingleObject(process, INFINITE);
+  DWORD exitCode = NO_ERROR;
+  res            = GetExitCodeProcess(process, &exitCode);
+  EXPECT_EQ(10002, exitCode);  // used init function exits process with this exit code
+
+  CloseHandle(process);
+  CloseHandle(thread);
+}
+
+TEST(InjectingTest, InjectionNoQuitInit)
+{
+  // Verify lib can inject with a init function with null parameters
+
+  HANDLE process, thread;
+  spawn(process, thread);
+  EXPECT_NO_THROW(InjectLib::InjectDLL(process, thread, INJECT_LIB, "InitNoQuit"));
+  ResumeThread(thread);
+
+  DWORD res      = WaitForSingleObject(process, INFINITE);
+  DWORD exitCode = NO_ERROR;
+  res            = GetExitCodeProcess(process, &exitCode);
+  EXPECT_EQ(0, exitCode);  // expect regular exit from process
+
+  CloseHandle(process);
+  CloseHandle(thread);
+}
+
+TEST(InjectingTest, InjectionSkipInit)
+{
+  // verify the skip-on-missing mechanism for init function works
+
+  HANDLE process, thread;
+  spawn(process, thread);
+  EXPECT_NO_THROW(InjectLib::InjectDLL(process, thread, INJECT_LIB, "__InitInvalid",
+                                       nullptr, 0, true));
+  ResumeThread(thread);
+
+  DWORD res      = WaitForSingleObject(process, INFINITE);
+  DWORD exitCode = NO_ERROR;
+  res            = GetExitCodeProcess(process, &exitCode);
+  EXPECT_EQ(NOERROR, exitCode);
+
+  CloseHandle(process);
+  CloseHandle(thread);
+}
+
+int main(int argc, char** argv)
+{
+  auto logger = spdlog::stdout_logger_mt("usvfs");
+  logger->set_level(spdlog::level::warn);
+
+  boost::filesystem::path filePath(winapi::wide::getModuleFileName(nullptr));
+  SetCurrentDirectoryW(filePath.parent_path().wstring().c_str());
+
+  testing::InitGoogleTest(&argc, argv);
+  return RUN_ALL_TESTS();
+}