Prevent DLL search order hijacking for VentoyPlugson.exe and VentoyVlnk.exe

[Ventoy.git] / Vlnk / src / main_windows.c

diff --git a/Vlnk/src/main_windows.c b/Vlnk/src/main_windows.c
index 34f75483a99fedd59204184a6b0ddec97aba88f0..d2f2c3c2b1795fe8954461264d69a048eed8a773 100644 (file)
--- a/Vlnk/src/main_windows.c
+++ b/Vlnk/src/main_windows.c
@@ -393,12 +393,12 @@ static BOOL IsSupportedVlnkSuffix(WCHAR *FileName)
 \r
     if (len > 9)\r
     {\r
-        if (lstrcmp(FileName - 9, L".vlnk.iso") == 0 ||\r
-            lstrcmp(FileName - 9, L".vlnk.img") == 0 || \r
-            lstrcmp(FileName - 9, L".vlnk.wim") == 0 || \r
-            lstrcmp(FileName - 9, L".vlnk.vhd") == 0 || \r
-            lstrcmp(FileName - 9, L".vlnk.efi") == 0 || \r
-            lstrcmp(FileName - 9, L".vlnk.dat") == 0)\r
+        if (lstrcmp(FileName + len - 9, L".vlnk.iso") == 0 ||\r
+            lstrcmp(FileName + len - 9, L".vlnk.img") == 0 ||\r
+            lstrcmp(FileName + len - 9, L".vlnk.wim") == 0 ||\r
+            lstrcmp(FileName + len - 9, L".vlnk.vhd") == 0 ||\r
+            lstrcmp(FileName + len - 9, L".vlnk.efi") == 0 ||\r
+            lstrcmp(FileName + len - 9, L".vlnk.dat") == 0)\r
         {\r
             return TRUE;\r
         }\r
@@ -407,8 +407,8 @@ static BOOL IsSupportedVlnkSuffix(WCHAR *FileName)
 \r
     if (len > 10)\r
     {\r
-        if (lstrcmp(FileName - 10, L".vlnk.vhdx") == 0 ||\r
-            lstrcmp(FileName - 9, L".vlnk.vtoy") == 0)\r
+        if (lstrcmp(FileName + len - 10, L".vlnk.vhdx") == 0 ||\r
+            lstrcmp(FileName + len - 10, L".vlnk.vtoy") == 0)\r
         {\r
             return TRUE;\r
         }\r
@@ -831,6 +831,58 @@ static int ParseCmdLine(LPSTR lpCmdLine)
     return argc;\r
 }\r
 \r
+\r
+//\r
+//copy from Rufus\r
+//\r
+#include <delayimp.h>\r
+// For delay-loaded DLLs, use LOAD_LIBRARY_SEARCH_SYSTEM32 to avoid DLL search order hijacking.\r
+FARPROC WINAPI dllDelayLoadHook(unsigned dliNotify, PDelayLoadInfo pdli)\r
+{\r
+    if (dliNotify == dliNotePreLoadLibrary) {\r
+        // Windows 7 without KB2533623 does not support the LOAD_LIBRARY_SEARCH_SYSTEM32 flag.\r
+        // That is is OK, because the delay load handler will interrupt the NULL return value\r
+        // to mean that it should perform a normal LoadLibrary.\r
+        return (FARPROC)LoadLibraryExA(pdli->szDll, NULL, LOAD_LIBRARY_SEARCH_SYSTEM32);\r
+    }\r
+    return NULL;\r
+}\r
+\r
+#if defined(_MSC_VER)\r
+// By default the Windows SDK headers have a `const` while MinGW does not.\r
+const\r
+#endif\r
+PfnDliHook __pfnDliNotifyHook2 = dllDelayLoadHook;\r
+\r
+typedef BOOL(WINAPI *SetDefaultDllDirectories_t)(DWORD);\r
+static void DllProtect(void)\r
+{\r
+    SetDefaultDllDirectories_t pfSetDefaultDllDirectories = NULL;\r
+\r
+    // Disable loading system DLLs from the current directory (sideloading mitigation)\r
+    // PS: You know that official MSDN documentation for SetDllDirectory() that explicitly\r
+    // indicates that "If the parameter is an empty string (""), the call removes the current\r
+    // directory from the default DLL search order"? Yeah, that doesn't work. At all.\r
+    // Still, we invoke it, for platforms where the following call might actually work...\r
+    SetDllDirectoryA("");\r
+\r
+    // For libraries on the KnownDLLs list, the system will always load them from System32.\r
+    // For other DLLs we link directly to, we can delay load the DLL and use a delay load\r
+    // hook to load them from System32. Note that, for this to work, something like:\r
+    // 'somelib.dll;%(DelayLoadDLLs)' must be added to the 'Delay Loaded Dlls' option of\r
+    // the linker properties in Visual Studio (which means this won't work with MinGW).\r
+    // For all other DLLs, use SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_SYSTEM32).\r
+    // Finally, we need to perform the whole gymkhana below, where we can't call on\r
+    // SetDefaultDllDirectories() directly, because Windows 7 doesn't have the API exposed.\r
+    // Also, no, Coverity, we never need to care about freeing kernel32 as a library.\r
+    // coverity[leaked_storage]\r
+\r
+    pfSetDefaultDllDirectories = (SetDefaultDllDirectories_t)\r
+        GetProcAddress(LoadLibraryW(L"kernel32.dll"), "SetDefaultDllDirectories");\r
+    if (pfSetDefaultDllDirectories != NULL)\r
+        pfSetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_SYSTEM32);\r
+}\r
+\r
 int APIENTRY WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, INT nCmdShow)\r
 {\r
     DWORD dwAttrib;\r
@@ -838,6 +890,8 @@ int APIENTRY WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLi
 \r
     UNREFERENCED_PARAMETER(hPrevInstance);\r
 \r
+    DllProtect();\r
+\r
     if (GetUserDefaultUILanguage() == 0x0804)\r
     {\r
         g_msg_lang = g_msg_cn;\r
@@ -879,6 +933,13 @@ int APIENTRY WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLi
             return 1;\r
         }\r
 \r
+        if (!IsSupportedVlnkSuffix(g_CmdOutFile))\r
+        {\r
+            LogW(L"File <<%ls>> contains invalid vlnk suffix!\n", g_CmdOutFile);\r
+            VtoyMessageBox(NULL, g_msg_lang[MSGID_INVALID_SUFFIX], g_msg_lang[MSGID_ERROR], MB_OK | MB_ICONERROR);\r
+            return 1;\r
+        }\r
+\r
         return CreateVlnk(NULL, g_CurDirW, g_CmdInFile, g_CmdOutFile);\r
     }\r
     else\r