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 064604058723a847847cb375fc90aaa32382d232..d2f2c3c2b1795fe8954461264d69a048eed8a773 100644 (file)
--- a/Vlnk/src/main_windows.c
+++ b/Vlnk/src/main_windows.c
@@ -14,6 +14,7 @@ static HWND g_create_button;
 static HWND g_parse_button;\r
 \r
 static BOOL g_ShowHelp = FALSE;\r
 static HWND g_parse_button;\r
 \r
 static BOOL g_ShowHelp = FALSE;\r

+static BOOL g_SaveAs = FALSE;\r

 static WCHAR g_CmdInFile[MAX_PATH];\r
 static WCHAR g_CmdOutFile[MAX_PATH];\r
 \r
 static WCHAR g_CmdInFile[MAX_PATH];\r
 static WCHAR g_CmdOutFile[MAX_PATH];\r
 \r


@@ -36,6 +37,7 @@ typedef enum MSGID
     MSGID_VLNK_POINT_TO,\r
     MSGID_VLNK_NO_DST,\r
     MSGID_FILE_NAME_TOO_LONG,\r
     MSGID_VLNK_POINT_TO,\r
     MSGID_VLNK_NO_DST,\r
     MSGID_FILE_NAME_TOO_LONG,\r

+    MSGID_INVALID_SUFFIX,\r

 \r
     MSGID_BUTT\r
 }MSGID;\r
 \r
     MSGID_BUTT\r
 }MSGID;\r


@@ -60,6 +62,7 @@ const WCHAR *g_msg_cn[MSGID_BUTT] =
     L"´Ë vlnk ÎļþÖ¸Ïò ",\r
     L"´Ë vlnk Ö¸ÏòµÄÎļþ²»´æÔÚ£¡",\r
     L"Îļþ·¾¶Ì«³¤£¡",\r
     L"´Ë vlnk ÎļþÖ¸Ïò ",\r
     L"´Ë vlnk Ö¸ÏòµÄÎļþ²»´æÔÚ£¡",\r
     L"Îļþ·¾¶Ì«³¤£¡",\r

+    L"·Ç·¨µÄvlnkÎļþºó׺Ãû!",\r

 };\r
 const WCHAR *g_msg_en[MSGID_BUTT] =\r
 {\r
 };\r
 const WCHAR *g_msg_en[MSGID_BUTT] =\r
 {\r


@@ -80,6 +83,7 @@ const WCHAR *g_msg_en[MSGID_BUTT] =
     L"The vlnk file point to ",\r
     L"The file pointed by the vlnk does NOT exist!",\r
     L"The file full path is too long!",\r
     L"The vlnk file point to ",\r
     L"The file pointed by the vlnk does NOT exist!",\r
     L"The file full path is too long!",\r

+    L"Invalid vlnk file suffix!",\r

 };\r
 \r
 const WCHAR **g_msg_lang = NULL;\r
 };\r
 \r
 const WCHAR **g_msg_lang = NULL;\r


@@ -363,11 +367,62 @@ End:
 }\r
 \r
 \r
 }\r
 \r
 \r

+static BOOL VentoyGetSaveFileName(HWND hWnd, WCHAR *szFile)\r
+{\r
+    OPENFILENAME ofn = { 0 };\r
+\r
+    ofn.lStructSize = sizeof(ofn);\r
+    ofn.hwndOwner = hWnd;\r
+    ofn.lpstrFilter = L"Vlnk File\0*.vlnk.iso;*.vlnk.img;*.vlnk.wim;*.vlnk.efi;*.vlnk.vhd;*.vlnk.vhdx;*.vlnk.vtoy;*.vlnk.dat\0";\r
+    ofn.nFilterIndex = 1;\r
+    ofn.lpstrFile = szFile;\r
+    ofn.nMaxFile = MAX_PATH;\r
+    ofn.Flags = OFN_EXPLORER | OFN_PATHMUSTEXIST | OFN_HIDEREADONLY | OFN_OVERWRITEPROMPT;\r
+    ofn.lpstrFileTitle = NULL;\r
+    ofn.nMaxFileTitle = 0;\r
+    ofn.lpstrInitialDir = NULL;\r
+\r
+    return GetSaveFileName(&ofn);\r
+}\r
+\r
+static BOOL IsSupportedVlnkSuffix(WCHAR *FileName)\r
+{\r
+    int len;\r
+\r
+    len = lstrlen(FileName);\r
+\r
+    if (len > 9)\r
+    {\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
+    }\r
+\r
+\r
+    if (len > 10)\r
+    {\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
+    }\r
+\r
+    return FALSE;\r
+}\r
+\r

 static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile)\r
 {\r
     int i;\r
     int end;\r
     int len;\r
 static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile)\r
 {\r
     int i;\r
     int end;\r
     int len;\r

+    BOOL SetOutFile = FALSE;\r

     UINT32 DiskSig;\r
     DISK_EXTENT DiskExtend;\r
     OPENFILENAME ofn = { 0 };\r
     UINT32 DiskSig;\r
     DISK_EXTENT DiskExtend;\r
     OPENFILENAME ofn = { 0 };\r


@@ -500,6 +555,7 @@ static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile)
         return 1;\r
     }\r
 \r
         return 1;\r
     }\r
 \r

+\r

     Buf = malloc(VLNK_FILE_LEN);\r
     if (Buf)\r
     {\r
     Buf = malloc(VLNK_FILE_LEN);\r
     if (Buf)\r
     {\r


@@ -510,10 +566,34 @@ static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile)
         if (OutFile)\r
         {\r
             wcscpy_s(DstFullPath, MAX_PATH, OutFile);\r
         if (OutFile)\r
         {\r
             wcscpy_s(DstFullPath, MAX_PATH, OutFile);\r

+            SetOutFile = TRUE;\r

         }\r
         else\r
         {\r
             DefaultVlnkDstFullPath(Pos + 1, Dir, DstFullPath);\r
         }\r
         else\r
         {\r
             DefaultVlnkDstFullPath(Pos + 1, Dir, DstFullPath);\r

+\r
+            if (g_SaveAs)\r
+            {\r
+                wcscpy_s(szFile, MAX_PATH, DstFullPath);\r
+                if (VentoyGetSaveFileName(hWnd, szFile))\r
+                {\r
+                    if (IsSupportedVlnkSuffix(szFile))\r
+                    {\r
+                        wcscpy_s(DstFullPath, MAX_PATH, szFile);\r
+                        SetOutFile = TRUE;\r
+                    }\r
+                    else\r
+                    {\r
+                        VtoyMessageBox(hWnd, g_msg_lang[MSGID_INVALID_SUFFIX], g_msg_lang[MSGID_ERROR], MB_OK | MB_ICONERROR);\r
+                        LogA("Invalid vlnk suffix\n");\r
+                        goto end;\r
+                    }\r
+                }\r
+                else\r
+                {\r
+                    LogA("User cancel the save as diaglog, use default name\n");\r
+                }\r
+            }\r

         }\r
 \r
         LogW(L"vlnk output file path is <%ls>\n", DstFullPath);\r
         }\r
 \r
         LogW(L"vlnk output file path is <%ls>\n", DstFullPath);\r


@@ -524,7 +604,7 @@ static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile)
 \r
             LogW(L"Vlnk file create success <%ls>\n", DstFullPath);\r
 \r
 \r
             LogW(L"Vlnk file create success <%ls>\n", DstFullPath);\r
 \r

-            if (OutFile)\r
+            if (SetOutFile)\r

             {\r
                 swprintf_s(Msg, 1024, L"%ls\r\n\r\n%ls", g_msg_lang[MSGID_VLNK_SUCCESS], DstFullPath);\r
                 VtoyMessageBox(hWnd, Msg, g_msg_lang[MSGID_INFO], MB_OK | MB_ICONINFORMATION);\r
             {\r
                 swprintf_s(Msg, 1024, L"%ls\r\n\r\n%ls", g_msg_lang[MSGID_VLNK_SUCCESS], DstFullPath);\r
                 VtoyMessageBox(hWnd, Msg, g_msg_lang[MSGID_INFO], MB_OK | MB_ICONINFORMATION);\r


@@ -541,6 +621,7 @@ static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile)
             VtoyMessageBox(hWnd, g_msg_lang[MSGID_CREATE_FILE_ERR], g_msg_lang[MSGID_ERROR], MB_OK | MB_ICONERROR);\r
         }\r
 \r
             VtoyMessageBox(hWnd, g_msg_lang[MSGID_CREATE_FILE_ERR], g_msg_lang[MSGID_ERROR], MB_OK | MB_ICONERROR);\r
         }\r
 \r

+        end:\r

         free(Buf);\r
     }\r
 \r
         free(Buf);\r
     }\r
 \r


@@ -580,7 +661,6 @@ static CHAR GetDriveLetter(UINT32 disksig, UINT64 PartOffset)
     return 0;\r
 }\r
 \r
     return 0;\r
 }\r
 \r

-\r

 static int ParseVlnk(HWND hWnd)\r
 {\r
     int i;\r
 static int ParseVlnk(HWND hWnd)\r
 {\r
     int i;\r


@@ -728,6 +808,10 @@ static int ParseCmdLine(LPSTR lpCmdLine)
         {\r
             g_ShowHelp = TRUE;\r
         }\r
         {\r
             g_ShowHelp = TRUE;\r
         }\r

+        else if (lstrcmp(lpszArgv[i], L"-s") == 0 || lstrcmp(lpszArgv[i], L"-S") == 0)\r
+        {\r
+            g_SaveAs = TRUE;\r
+        }\r

         else if (lstrcmp(lpszArgv[i], L"-i") == 0 || lstrcmp(lpszArgv[i], L"-I") == 0)\r
         {\r
             if (i + 1 < argc)\r
         else if (lstrcmp(lpszArgv[i], L"-i") == 0 || lstrcmp(lpszArgv[i], L"-I") == 0)\r
         {\r
             if (i + 1 < argc)\r


@@ -747,6 +831,58 @@ static int ParseCmdLine(LPSTR lpCmdLine)
     return argc;\r
 }\r
 \r
     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
 int APIENTRY WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, INT nCmdShow)\r
 {\r
     DWORD dwAttrib;\r


@@ -754,6 +890,8 @@ int APIENTRY WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLi
 \r
     UNREFERENCED_PARAMETER(hPrevInstance);\r
 \r
 \r
     UNREFERENCED_PARAMETER(hPrevInstance);\r
 \r

+    DllProtect();\r
+\r

     if (GetUserDefaultUILanguage() == 0x0804)\r
     {\r
         g_msg_lang = g_msg_cn;\r
     if (GetUserDefaultUILanguage() == 0x0804)\r
     {\r
         g_msg_lang = g_msg_cn;\r


@@ -795,6 +933,13 @@ int APIENTRY WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLi
             return 1;\r
         }\r
 \r
             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
         return CreateVlnk(NULL, g_CurDirW, g_CmdInFile, g_CmdOutFile);\r
     }\r
     else\r