Fixes and improvements

GuardDLL/GravitGuard.cpp

@@ -127,6 +127,7 @@ GravitGuard::CheckResultAll GravitGuard::check(DLLSource source)
 #endif
 	bool isOk = false;
 	bool found_unknown_module = false;
+	bool suspect = false;
 	for (int i = 0; i < n; i++) {
 		StackTraceElement el;
 		el.address = buf[i];
@@ -146,10 +147,19 @@ GravitGuard::CheckResultAll GravitGuard::check(DLLSource source)
 			el.fileName = getModuleFileName(module);
 			el.source = getSource(el.fileName);
 			if (!is_known_module(el.fileName)) {
-				if (el.source != DLLSource::SYSTEM) {
+				std::wstring_view view = el.fileName;
+				if (el.source == DLLSource::SYSTEM) {
+					if (view.ends_with(L"\\ucrtbase.dll")) { // ucrtbase.dll injected by system without LdrLoadDll
+						gg->add_known_module(el.fileName);
+					}
+					else {
+						//suspect = true;
+					}
+				}
+				else {
 					found_unknown_module = true;
+					this_element_in_unknown_module = true;
 				}
-				this_element_in_unknown_module = true;
 			}
 			if (!is_win_internal_module(module)) {
 				isOk = true;
@@ -169,6 +179,9 @@ GravitGuard::CheckResultAll GravitGuard::check(DLLSource source)
 		return CheckResultAll{ CheckResult::TERMINATE_APP };
 	}
 	if (isOk) {
+		if (suspect) {
+			return CheckResultAll{ CheckResult::ONLY_SIGNED };
+		}
 		return CheckResultAll{ CheckResult::PASS };
 	}
 	else {

GuardDLL/Logger.cpp

@@ -1,16 +1,58 @@
 #include "pch.h"
+#include "guard.h"
 #include <format>
+#ifdef ENABLE_LOGGING_CONSOLE
 #include <iostream>
+#endif
+#ifdef ENABLE_LOGGING_FILE
+#include <fstream>
+#endif
 #include "Logger.h"
 
 Logger logger;
+#ifdef ENABLE_LOGGING_FILE
+std::fstream* logfile;
+#endif
+
+Logger::Logger()
+{
+#ifdef ENABLE_LOGGING_FILE
+	logfile = new std::fstream(L"GravitGuard2.log", std::ios_base::app | std::ios_base::out);
+#endif
+}
+
+Logger::~Logger()
+{
+#ifdef ENABLE_LOGGING_FILE
+	delete logfile;
+#endif
+}
+
+std::string Logger::toUtf8(std::wstring_view wstr)
+{
+	size_t size = wstr.size() * 2;
+	std::string appdata(size, ' ');
+	size_t nCnt = WideCharToMultiByte(CP_UTF8, 0, wstr.data(), wstr.size(), appdata.data(), size, 0, 0);
+	appdata.resize(nCnt);
+	return appdata;
+}
 
 void Logger::log(std::string_view str)
 {
+#ifdef ENABLE_LOGGING_CONSOLE
 	std::cout << str << std::endl;
+#endif
+#ifdef ENABLE_LOGGING_FILE
+	*logfile << str << std::endl;
+#endif
 }
 
 void Logger::log(std::wstring_view str)
 {
+#ifdef ENABLE_LOGGING_CONSOLE
 	std::wcout << str << std::endl;
+#endif
+#ifdef ENABLE_LOGGING_FILE
+	*logfile << toUtf8(str) << std::endl;
+#endif
 }

GuardDLL/Logger.h

@@ -4,13 +4,16 @@
 class Logger
 {
 public:
+	Logger();
+	~Logger();
 	enum class Level {
 		TRACE, DEBUG, INFO, CRITICAL
 	};
 	template<class... T>
 	void log(Level level, const std::format_string<T...> _Fmt,T... _Args);
 	template<class... T>
 	void log(Level level, const std::wformat_string<T...> _Fmt, T... _Args);
+	std::string toUtf8(std::wstring_view wstr);
 private:
 	void log(std::string_view str);
 	void log(std::wstring_view str);

GuardDLL/hooks.cpp

@@ -41,18 +41,46 @@ inline void checkTerminateActions(GravitGuard::CheckResultAll result) {
 
 NTSTATUS NTAPI LdrLoadDllHook(UINT32 Flags, PUINT32 Reserved, PUNICODE_STRING ModuleFileName, PHANDLE ModuleHandle) {
 	std::wstring_view fileName(ModuleFileName->Buffer, ModuleFileName->Length/2);
+#ifdef ENABLE_DLL_THREAD_LOCAL_OPTIMIZATION
+	thread_local std::vector<std::wstring> checked;
+	for (auto l : checked) {
+		if (l == fileName) {
+			return ldrLoadDllHook->call_original(std::move(Flags), std::move(Reserved), std::move(ModuleFileName), std::move(ModuleHandle));
+		}
+	}
+#endif
 	logger.log(Logger::Level::DEBUG, L"LdrLoadDll {}", fileName);
 	auto source = gg->getSource(fileName);
 	auto result = gg->check(source);
 	checkTerminateActions(result);
 	if (result.result == GravitGuard::CheckResult::CANCEL) {
+		logger.log(Logger::Level::CRITICAL, L"LdrLoadDll {} cancelled", fileName);
 		return -1;
 	}
 	if (result.result == GravitGuard::CheckResult::ONLY_SIGNED) {
+		logger.log(Logger::Level::CRITICAL, L"LdrLoadDll {} only signed", fileName);
 		Flags |= LOAD_LIBRARY_REQUIRE_SIGNED_TARGET;
 	}
-	gg->add_known_module(fileName);
-	return ldrLoadDllHook->call_original(std::move(Flags), std::move(Reserved), std::move(ModuleFileName), std::move(ModuleHandle));
+	if (!gg->is_known_module(fileName)) {
+		gg->add_known_module(fileName);
+	}
+	auto retCode = ldrLoadDllHook->call_original(std::move(Flags), std::move(Reserved), std::move(ModuleFileName), std::move(ModuleHandle));
+	if (NT_SUCCESS(retCode)) {
+		auto realModuleFileName = gg->getModuleFileName(*((HMODULE*)ModuleHandle));
+		if (fileName != realModuleFileName) {
+			logger.log(Logger::Level::DEBUG, L"Found module {}: {}", fileName, realModuleFileName);
+			if (!gg->is_known_module(realModuleFileName)) {
+				gg->add_known_module(realModuleFileName);
+			}
+		}
+#ifdef ENABLE_DLL_THREAD_LOCAL_OPTIMIZATION
+		checked.push_back(std::wstring(fileName));
+#endif
+	}
+	else {
+		logger.log(Logger::Level::DEBUG, L"Fail to load module {}", fileName);
+	}
+	return retCode;
 }
 
 NTSTATUS NTAPI LdrpLoadDllHook(BOOLEAN Redirected, PWSTR DllPath, PULONG DllCharacteristics, PUNICODE_STRING DllName, PVOID* BaseAddress, BOOLEAN CallInit) {

GuardDLL/include/guard.h

@@ -1,4 +1,7 @@
 #pragma once
+#define ENABLE_LOGGING_FILE
+#define ENABLE_LOGGING_CONSOLE
+#define ENABLE_DLL_THREAD_LOCAL_OPTIMIZATION
 #include <string>
 #define GUARD_EXPORT    __declspec(dllexport)
 extern "C" {