X-Git-Url: https://glassweightruler.freedombox.rocks/gitweb/Ventoy.git/blobdiff_plain/1ca48923da8cb7e351b646429cafbdd03eab5f9e..dcc588967716395dfdb51be9dcd24e64145edd1b:/Vlnk/src/main_windows.c diff --git a/Vlnk/src/main_windows.c b/Vlnk/src/main_windows.c index 0646040..d2f2c3c 100644 --- 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; static BOOL g_ShowHelp = FALSE; +static BOOL g_SaveAs = FALSE; static WCHAR g_CmdInFile[MAX_PATH]; static WCHAR g_CmdOutFile[MAX_PATH]; @@ -36,6 +37,7 @@ typedef enum MSGID MSGID_VLNK_POINT_TO, MSGID_VLNK_NO_DST, MSGID_FILE_NAME_TOO_LONG, + MSGID_INVALID_SUFFIX, MSGID_BUTT }MSGID; @@ -60,6 +62,7 @@ const WCHAR *g_msg_cn[MSGID_BUTT] = L"´Ë vlnk ÎļþÖ¸Ïò ", L"´Ë vlnk Ö¸ÏòµÄÎļþ²»´æÔÚ£¡", L"Îļþ·¾¶Ì«³¤£¡", + L"·Ç·¨µÄvlnkÎļþºó׺Ãû!", }; const WCHAR *g_msg_en[MSGID_BUTT] = { @@ -80,6 +83,7 @@ const WCHAR *g_msg_en[MSGID_BUTT] = L"The vlnk file point to ", L"The file pointed by the vlnk does NOT exist!", L"The file full path is too long!", + L"Invalid vlnk file suffix!", }; const WCHAR **g_msg_lang = NULL; @@ -363,11 +367,62 @@ End: } +static BOOL VentoyGetSaveFileName(HWND hWnd, WCHAR *szFile) +{ + OPENFILENAME ofn = { 0 }; + + ofn.lStructSize = sizeof(ofn); + ofn.hwndOwner = hWnd; + ofn.lpstrFilter = L"Vlnk File\0*.vlnk.iso;*.vlnk.img;*.vlnk.wim;*.vlnk.efi;*.vlnk.vhd;*.vlnk.vhdx;*.vlnk.vtoy;*.vlnk.dat\0"; + ofn.nFilterIndex = 1; + ofn.lpstrFile = szFile; + ofn.nMaxFile = MAX_PATH; + ofn.Flags = OFN_EXPLORER | OFN_PATHMUSTEXIST | OFN_HIDEREADONLY | OFN_OVERWRITEPROMPT; + ofn.lpstrFileTitle = NULL; + ofn.nMaxFileTitle = 0; + ofn.lpstrInitialDir = NULL; + + return GetSaveFileName(&ofn); +} + +static BOOL IsSupportedVlnkSuffix(WCHAR *FileName) +{ + int len; + + len = lstrlen(FileName); + + if (len > 9) + { + if (lstrcmp(FileName + len - 9, L".vlnk.iso") == 0 || + lstrcmp(FileName + len - 9, L".vlnk.img") == 0 || + lstrcmp(FileName + len - 9, L".vlnk.wim") == 0 || + lstrcmp(FileName + len - 9, L".vlnk.vhd") == 0 || + lstrcmp(FileName + len - 9, L".vlnk.efi") == 0 || + lstrcmp(FileName + len - 9, L".vlnk.dat") == 0) + { + return TRUE; + } + } + + + if (len > 10) + { + if (lstrcmp(FileName + len - 10, L".vlnk.vhdx") == 0 || + lstrcmp(FileName + len - 10, L".vlnk.vtoy") == 0) + { + return TRUE; + } + } + + return FALSE; +} + static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile) { int i; int end; int len; + BOOL SetOutFile = FALSE; UINT32 DiskSig; DISK_EXTENT DiskExtend; OPENFILENAME ofn = { 0 }; @@ -500,6 +555,7 @@ static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile) return 1; } + Buf = malloc(VLNK_FILE_LEN); if (Buf) { @@ -510,10 +566,34 @@ static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile) if (OutFile) { wcscpy_s(DstFullPath, MAX_PATH, OutFile); + SetOutFile = TRUE; } else { DefaultVlnkDstFullPath(Pos + 1, Dir, DstFullPath); + + if (g_SaveAs) + { + wcscpy_s(szFile, MAX_PATH, DstFullPath); + if (VentoyGetSaveFileName(hWnd, szFile)) + { + if (IsSupportedVlnkSuffix(szFile)) + { + wcscpy_s(DstFullPath, MAX_PATH, szFile); + SetOutFile = TRUE; + } + else + { + VtoyMessageBox(hWnd, g_msg_lang[MSGID_INVALID_SUFFIX], g_msg_lang[MSGID_ERROR], MB_OK | MB_ICONERROR); + LogA("Invalid vlnk suffix\n"); + goto end; + } + } + else + { + LogA("User cancel the save as diaglog, use default name\n"); + } + } } LogW(L"vlnk output file path is <%ls>\n", DstFullPath); @@ -524,7 +604,7 @@ static int CreateVlnk(HWND hWnd, WCHAR *Dir, WCHAR *InFile, WCHAR *OutFile) LogW(L"Vlnk file create success <%ls>\n", DstFullPath); - if (OutFile) + if (SetOutFile) { swprintf_s(Msg, 1024, L"%ls\r\n\r\n%ls", g_msg_lang[MSGID_VLNK_SUCCESS], DstFullPath); VtoyMessageBox(hWnd, Msg, g_msg_lang[MSGID_INFO], MB_OK | MB_ICONINFORMATION); @@ -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); } + end: free(Buf); } @@ -580,7 +661,6 @@ static CHAR GetDriveLetter(UINT32 disksig, UINT64 PartOffset) return 0; } - static int ParseVlnk(HWND hWnd) { int i; @@ -728,6 +808,10 @@ static int ParseCmdLine(LPSTR lpCmdLine) { g_ShowHelp = TRUE; } + else if (lstrcmp(lpszArgv[i], L"-s") == 0 || lstrcmp(lpszArgv[i], L"-S") == 0) + { + g_SaveAs = TRUE; + } else if (lstrcmp(lpszArgv[i], L"-i") == 0 || lstrcmp(lpszArgv[i], L"-I") == 0) { if (i + 1 < argc) @@ -747,6 +831,58 @@ static int ParseCmdLine(LPSTR lpCmdLine) return argc; } + +// +//copy from Rufus +// +#include +// For delay-loaded DLLs, use LOAD_LIBRARY_SEARCH_SYSTEM32 to avoid DLL search order hijacking. +FARPROC WINAPI dllDelayLoadHook(unsigned dliNotify, PDelayLoadInfo pdli) +{ + if (dliNotify == dliNotePreLoadLibrary) { + // Windows 7 without KB2533623 does not support the LOAD_LIBRARY_SEARCH_SYSTEM32 flag. + // That is is OK, because the delay load handler will interrupt the NULL return value + // to mean that it should perform a normal LoadLibrary. + return (FARPROC)LoadLibraryExA(pdli->szDll, NULL, LOAD_LIBRARY_SEARCH_SYSTEM32); + } + return NULL; +} + +#if defined(_MSC_VER) +// By default the Windows SDK headers have a `const` while MinGW does not. +const +#endif +PfnDliHook __pfnDliNotifyHook2 = dllDelayLoadHook; + +typedef BOOL(WINAPI *SetDefaultDllDirectories_t)(DWORD); +static void DllProtect(void) +{ + SetDefaultDllDirectories_t pfSetDefaultDllDirectories = NULL; + + // Disable loading system DLLs from the current directory (sideloading mitigation) + // PS: You know that official MSDN documentation for SetDllDirectory() that explicitly + // indicates that "If the parameter is an empty string (""), the call removes the current + // directory from the default DLL search order"? Yeah, that doesn't work. At all. + // Still, we invoke it, for platforms where the following call might actually work... + SetDllDirectoryA(""); + + // For libraries on the KnownDLLs list, the system will always load them from System32. + // For other DLLs we link directly to, we can delay load the DLL and use a delay load + // hook to load them from System32. Note that, for this to work, something like: + // 'somelib.dll;%(DelayLoadDLLs)' must be added to the 'Delay Loaded Dlls' option of + // the linker properties in Visual Studio (which means this won't work with MinGW). + // For all other DLLs, use SetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_SYSTEM32). + // Finally, we need to perform the whole gymkhana below, where we can't call on + // SetDefaultDllDirectories() directly, because Windows 7 doesn't have the API exposed. + // Also, no, Coverity, we never need to care about freeing kernel32 as a library. + // coverity[leaked_storage] + + pfSetDefaultDllDirectories = (SetDefaultDllDirectories_t) + GetProcAddress(LoadLibraryW(L"kernel32.dll"), "SetDefaultDllDirectories"); + if (pfSetDefaultDllDirectories != NULL) + pfSetDefaultDllDirectories(LOAD_LIBRARY_SEARCH_SYSTEM32); +} + int APIENTRY WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLine, INT nCmdShow) { DWORD dwAttrib; @@ -754,6 +890,8 @@ int APIENTRY WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLi UNREFERENCED_PARAMETER(hPrevInstance); + DllProtect(); + if (GetUserDefaultUILanguage() == 0x0804) { g_msg_lang = g_msg_cn; @@ -795,6 +933,13 @@ int APIENTRY WinMain(HINSTANCE hInstance, HINSTANCE hPrevInstance, LPSTR lpCmdLi return 1; } + if (!IsSupportedVlnkSuffix(g_CmdOutFile)) + { + LogW(L"File <<%ls>> contains invalid vlnk suffix!\n", g_CmdOutFile); + VtoyMessageBox(NULL, g_msg_lang[MSGID_INVALID_SUFFIX], g_msg_lang[MSGID_ERROR], MB_OK | MB_ICONERROR); + return 1; + } + return CreateVlnk(NULL, g_CurDirW, g_CmdInFile, g_CmdOutFile); } else