Ventoy2Disk/Ventoy2Disk/AlertSuppress.c

   1 /******************************************************************************
   2 * AlertSuppress.c
   3 *
   4 * Copyright (c) 2022, longpanda <admin@ventoy.net>
   5 * Copyright (c) 2011-2022 Pete Batard <pete@akeo.ie>
   6 *
   7 * This program is free software; you can redistribute it and/or
   8 * modify it under the terms of the GNU General Public License as
   9 * published by the Free Software Foundation; either version 3 of the
  10 * License, or (at your option) any later version.
  11 *
  12 * This program is distributed in the hope that it will be useful, but
  13 * WITHOUT ANY WARRANTY; without even the implied warranty of
  14 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  15 * General Public License for more details.
  16 *
  17 * You should have received a copy of the GNU General Public License
  18 * along with this program; if not, see <http://www.gnu.org/licenses/>.
  19 *
  20 */
  21
  22 #include <Windows.h>
  23 #include <winternl.h>
  24 #include <commctrl.h>
  25 #include <initguid.h>
  26 #include <vds.h>
  27 #include "Ventoy2Disk.h"
  28
  29
  30 #define sfree(p) do {if (p != NULL) {free((void*)(p)); p = NULL;}} while(0)
  31 #define wconvert(p)     wchar_t* w ## p = utf8_to_wchar(p)
  32 #define walloc(p, size) wchar_t* w ## p = (p == NULL)?NULL:(wchar_t*)calloc(size, sizeof(wchar_t))
  33 #define wfree(p) sfree(w ## p)
  34
  35 #define static_strcpy(dst, src) strcpy_s(dst, sizeof(dst), src)
  36 #define static_strcat(dst, src) strcat_s(dst, sizeof(dst), src)
  37
  38
  39 #define wchar_to_utf8_no_alloc(wsrc, dest, dest_size) \
  40         WideCharToMultiByte(CP_UTF8, 0, wsrc, -1, dest, dest_size, NULL, NULL)
  41 #define utf8_to_wchar_no_alloc(src, wdest, wdest_size) \
  42         MultiByteToWideChar(CP_UTF8, 0, src, -1, wdest, wdest_size)
  43
  44
  45 /*
  46  * Converts an UTF8 string to UTF-16 (allocate returned string)
  47  * Returns NULL on error
  48  */
  49 static __inline wchar_t* utf8_to_wchar(const char* str)
  50 {
  51         int size = 0;
  52         wchar_t* wstr = NULL;
  53
  54         if (str == NULL)
  55                 return NULL;
  56
  57         // Convert the empty string too
  58         if (str[0] == 0)
  59                 return (wchar_t*)calloc(1, sizeof(wchar_t));
  60
  61         // Find out the size we need to allocate for our converted string
  62         size = MultiByteToWideChar(CP_UTF8, 0, str, -1, NULL, 0);
  63         if (size <= 1)  // An empty string would be size 1
  64                 return NULL;
  65
  66         if ((wstr = (wchar_t*)calloc(size, sizeof(wchar_t))) == NULL)
  67                 return NULL;
  68
  69         if (utf8_to_wchar_no_alloc(str, wstr, size) != size) {
  70                 sfree(wstr);
  71                 return NULL;
  72         }
  73         return wstr;
  74 }
  75
  76 static char g_FormatDiskTitle[256];
  77 static char g_FormatDiskButton[256];
  78
  79 static char g_LocNotAvaliableTitle[256]; //Location is not available
  80 static char g_InsertDiskTitle[256]; // Insert disk
  81
  82 static char system_dir[MAX_PATH], sysnative_dir[MAX_PATH];
  83
  84 static HWINEVENTHOOK ap_weh = NULL;
  85
  86
  87 static __inline UINT GetSystemDirectoryU(char* lpBuffer, UINT uSize)
  88 {
  89         UINT ret = 0, err = ERROR_INVALID_DATA;
  90         // coverity[returned_null]
  91         walloc(lpBuffer, uSize);
  92         ret = GetSystemDirectoryW(wlpBuffer, uSize);
  93         err = GetLastError();
  94         if ((ret != 0) && ((ret = wchar_to_utf8_no_alloc(wlpBuffer, lpBuffer, uSize)) == 0)) {
  95                 err = GetLastError();
  96         }
  97         wfree(lpBuffer);
  98         SetLastError(err);
  99         return ret;
 100 }
 101
 102
 103 static char* ToLocaleName(DWORD lang_id)
 104 {
 105         static char mui_str[LOCALE_NAME_MAX_LENGTH];
 106         wchar_t wmui_str[LOCALE_NAME_MAX_LENGTH];
 107
 108         if (LCIDToLocaleName(lang_id, wmui_str, LOCALE_NAME_MAX_LENGTH, 0) > 0) {
 109                 wchar_to_utf8_no_alloc(wmui_str, mui_str, LOCALE_NAME_MAX_LENGTH);
 110         }
 111         else {
 112                 static_strcpy(mui_str, "en-US");
 113         }
 114         return mui_str;
 115 }
 116
 117 static __inline int LoadStringU(HINSTANCE hInstance, UINT uID, LPSTR lpBuffer, int nBufferMax)
 118 {
 119         int ret;
 120         DWORD err = ERROR_INVALID_DATA;
 121         if (nBufferMax == 0) {
 122                 // read-only pointer to resource mode is not supported
 123                 SetLastError(ERROR_INVALID_PARAMETER);
 124                 return 0;
 125         }
 126         // coverity[returned_null]
 127         walloc(lpBuffer, nBufferMax);
 128         ret = LoadStringW(hInstance, uID, wlpBuffer, nBufferMax);
 129         err = GetLastError();
 130         if ((ret > 0) && ((ret = wchar_to_utf8_no_alloc(wlpBuffer, lpBuffer, nBufferMax)) == 0)) {
 131                 err = GetLastError();
 132         }
 133         wfree(lpBuffer);
 134         SetLastError(err);
 135         return ret;
 136 }
 137
 138 /*
 139  * typedefs for the function prototypes. Use the something like:
 140  *   PF_DECL(FormatEx);
 141  * which translates to:
 142  *   FormatEx_t pfFormatEx = NULL;
 143  * in your code, to declare the entrypoint and then use:
 144  *   PF_INIT(FormatEx, Fmifs);
 145  * which translates to:
 146  *   pfFormatEx = (FormatEx_t) GetProcAddress(GetDLLHandle("fmifs"), "FormatEx");
 147  * to make it accessible.
 148  */
 149 #define         MAX_LIBRARY_HANDLES 64
 150 extern HMODULE  OpenedLibrariesHandle[MAX_LIBRARY_HANDLES];
 151 extern UINT16 OpenedLibrariesHandleSize;
 152 #define         OPENED_LIBRARIES_VARS HMODULE OpenedLibrariesHandle[MAX_LIBRARY_HANDLES]; uint16_t OpenedLibrariesHandleSize = 0
 153 #define         CLOSE_OPENED_LIBRARIES while(OpenedLibrariesHandleSize > 0) FreeLibrary(OpenedLibrariesHandle[--OpenedLibrariesHandleSize])
 154
 155 static __inline HMODULE GetLibraryHandle(char* szLibraryName) {
 156         HMODULE h = NULL;
 157         wchar_t* wszLibraryName = NULL;
 158         int size;
 159         if (szLibraryName == NULL || szLibraryName[0] == 0)
 160                 goto out;
 161         size = MultiByteToWideChar(CP_UTF8, 0, szLibraryName, -1, NULL, 0);
 162         if ((size <= 1) || ((wszLibraryName = (wchar_t*)calloc(size, sizeof(wchar_t))) == NULL) ||
 163                 (MultiByteToWideChar(CP_UTF8, 0, szLibraryName, -1, wszLibraryName, size) != size))
 164                 goto out;
 165         // If the library is already opened, just return a handle (that doesn't need to be freed)
 166         if ((h = GetModuleHandleW(wszLibraryName)) != NULL)
 167                 goto out;
 168         // Sanity check
 169         if (OpenedLibrariesHandleSize >= MAX_LIBRARY_HANDLES) {
 170                 Log("Error: MAX_LIBRARY_HANDLES is too small");
 171                 goto out;
 172         }
 173         h = LoadLibraryExW(wszLibraryName, NULL, LOAD_LIBRARY_SEARCH_SYSTEM32);
 174         // Some Windows 7 platforms (most likely the ones missing KB2533623 per the
 175         // official LoadLibraryEx doc) can return ERROR_INVALID_PARAMETER when using
 176         // the Ex() version. If that's the case, fallback to using LoadLibraryW().
 177         if ((h == NULL) && (SCODE_CODE(GetLastError()) == ERROR_INVALID_PARAMETER))
 178                 h = LoadLibraryW(wszLibraryName);
 179         if (h != NULL)
 180                 OpenedLibrariesHandle[OpenedLibrariesHandleSize++] = h;
 181         else
 182                 Log("Unable to load '%S.dll': %u", wszLibraryName, LASTERR);
 183 out:
 184         free(wszLibraryName);
 185         return h;
 186 }
 187
 188 #define PF_TYPE(api, ret, proc, args)           typedef ret (api *proc##_t)args
 189 #define PF_DECL(proc)                                           static proc##_t pf##proc = NULL
 190 #define PF_TYPE_DECL(api, ret, proc, args)      PF_TYPE(api, ret, proc, args); PF_DECL(proc)
 191 #define PF_INIT(proc, name)                                     if (pf##proc == NULL) pf##proc = \
 192         (proc##_t) GetProcAddress(GetLibraryHandle(#name), #proc)
 193 #define PF_INIT_OR_OUT(proc, name)                      do {PF_INIT(proc, name);         \
 194         if (pf##proc == NULL) {Log("Unable to locate %s() in '%s.dll': %u",  \
 195         #proc, #name, LASTERR); goto out;} } while(0)
 196 #define PF_INIT_OR_SET_STATUS(proc, name)       do {PF_INIT(proc, name);         \
 197         if ((pf##proc == NULL) && (NT_SUCCESS(status))) status = STATUS_NOT_IMPLEMENTED; } while(0)
 198
 199 static BOOL is_x64(void)
 200 {
 201         BOOL ret = FALSE;
 202         PF_TYPE_DECL(WINAPI, BOOL, IsWow64Process, (HANDLE, PBOOL));
 203         // Detect if we're running a 32 or 64 bit system
 204         if (sizeof(uintptr_t) < 8) {
 205                 PF_INIT(IsWow64Process, Kernel32);
 206                 if (pfIsWow64Process != NULL) {
 207                         (*pfIsWow64Process)(GetCurrentProcess(), &ret);
 208                 }
 209         }
 210         else {
 211                 ret = TRUE;
 212         }
 213         return ret;
 214 }
 215
 216
 217 static __inline UINT GetSystemWindowsDirectoryU(char* lpBuffer, UINT uSize)
 218 {
 219         UINT ret = 0, err = ERROR_INVALID_DATA;
 220         // coverity[returned_null]
 221         walloc(lpBuffer, uSize);
 222         ret = GetSystemWindowsDirectoryW(wlpBuffer, uSize);
 223         err = GetLastError();
 224         if ((ret != 0) && ((ret = wchar_to_utf8_no_alloc(wlpBuffer, lpBuffer, uSize)) == 0)) {
 225                 err = GetLastError();
 226         }
 227         wfree(lpBuffer);
 228         SetLastError(err);
 229         return ret;
 230 }
 231
 232 static __inline HMODULE LoadLibraryU(LPCSTR lpFileName)
 233 {
 234         HMODULE ret;
 235         DWORD err = ERROR_INVALID_DATA;
 236         wconvert(lpFileName);
 237         ret = LoadLibraryW(wlpFileName);
 238         err = GetLastError();
 239         wfree(lpFileName);
 240         SetLastError(err);
 241         return ret;
 242 }
 243
 244 #pragma pack(1)
 245 typedef struct {
 246         WORD      dlgVer;
 247         WORD      signature;
 248         DWORD     helpID;
 249         DWORD     exStyle;
 250         DWORD     style;
 251         WORD      cDlgItems;
 252         short     x;
 253         short     y;
 254         short     cx;
 255         short     cy;
 256         //sz_Or_Ord menu;
 257         //sz_Or_Ord windowClass;
 258         //WCHAR     title[titleLen];
 259         //WORD      pointsize;
 260         //WORD      weight;
 261         //BYTE      italic;
 262         //BYTE      charset;
 263         //WCHAR     typeface[stringLen];
 264 } DLGTEMPLATEEX;
 265 #pragma pack()
 266
 267 static BOOL LoadDialogCaption(HMODULE hMui, DWORD ID, CHAR* title, DWORD len)
 268 {
 269         BOOL bRet = FALSE;
 270         int WordNum = 0;
 271         HRSRC hDlg = NULL;
 272         DLGTEMPLATEEX* pDlgTempEx = NULL;
 273         HGLOBAL hTemplate = NULL;
 274         WORD* pWordData = NULL;
 275
 276         hDlg = FindResource(hMui, MAKEINTRESOURCE(1024), RT_DIALOG);
 277         if (hDlg)
 278         {
 279                 hTemplate = LoadResource(hMui, hDlg);
 280                 if (hTemplate)
 281                 {
 282                         pDlgTempEx = (DLGTEMPLATEEX*)LockResource(hTemplate);
 283                         if (pDlgTempEx)
 284                         {
 285                                 if (pDlgTempEx->signature != 0xFFFF)
 286                                 {
 287                                         return FALSE;
 288                                 }
 289
 290                                 pWordData = (WORD *)(pDlgTempEx + 1);
 291
 292                                 //skip menu
 293                                 if (*pWordData == 0x0000)
 294                                 {
 295                                         pWordData += 1;
 296                                 }
 297                                 else if (*pWordData == 0xFFFF)
 298                                 {
 299                                         pWordData += 2;
 300                                 }
 301                                 else
 302                                 {
 303                                         while (*pWordData++)
 304                                         {
 305                                                 ;
 306                                         }
 307                                 }
 308
 309                                 //skip windowClass
 310                                 if (*pWordData == 0x0000)
 311                                 {
 312                                         pWordData += 1;
 313                                 }
 314                                 else if (*pWordData == 0xFFFF)
 315                                 {
 316                                         pWordData += 2;
 317                                 }
 318                                 else
 319                                 {
 320                                         while (*pWordData++)
 321                                         {
 322                                                 ;
 323                                         }
 324                                 }
 325
 326                                 wchar_to_utf8_no_alloc(pWordData, title, len);
 327                                 bRet = TRUE;
 328                         }
 329                 }
 330         }
 331
 332         return bRet;
 333 }
 334
 335 BOOL SetAlertPromptMessages(void)
 336 {
 337         HMODULE hMui;
 338         char mui_path[MAX_PATH];
 339
 340         if (GetSystemDirectoryU(system_dir, sizeof(system_dir)) == 0) {
 341                 Log("Could not get system directory: %u", LASTERR);
 342                 static_strcpy(system_dir, "C:\\Windows\\System32");
 343         }
 344
 345         // Construct Sysnative ourselves as there is no GetSysnativeDirectory() call
 346         // By default (64bit app running on 64 bit OS or 32 bit app running on 32 bit OS)
 347         // Sysnative and System32 are the same
 348         static_strcpy(sysnative_dir, system_dir);
 349         // But if the app is 32 bit and the OS is 64 bit, Sysnative must differ from System32
 350 #if (defined(VTARCH_X86) || defined(VTARCH_ARM))
 351         if (is_x64()) {
 352                 if (GetSystemWindowsDirectoryU(sysnative_dir, sizeof(sysnative_dir)) == 0) {
 353                         Log("Could not get Windows directory: %u", LASTERR);
 354                         static_strcpy(sysnative_dir, "C:\\Windows");
 355                 }
 356                 static_strcat(sysnative_dir, "\\Sysnative");
 357         }
 358 #endif
 359
 360         Log("system_dir=<%s>", system_dir);
 361         Log("sysnative_dir=<%s>", sysnative_dir);
 362
 363         sprintf_s(mui_path, MAX_PATH, "%s\\%s\\shell32.dll.mui", sysnative_dir, ToLocaleName(GetUserDefaultUILanguage()));
 364
 365         hMui = LoadLibraryU(mui_path);
 366         if (hMui)
 367         {
 368                 Log("LoadLibrary shell32.dll.mui SUCCESS");
 369         }
 370         else
 371         {
 372                 Log("LoadLibrary shell32.dll.mui FAILED");
 373                 return FALSE;
 374         }
 375
 376
 377         // String Table:
 378         // 4097 = "You need to format the disk in drive %c: before you can use it." (dialog text)
 379         // 4125 = "Microsoft Windows" (dialog title)
 380         // 4126 = "Format disk" (button)
 381         if (LoadStringU(hMui, 4125, g_FormatDiskTitle, sizeof(g_FormatDiskTitle)) <= 0) {
 382                 static_strcpy(g_FormatDiskTitle, "Microsoft Windows");
 383                 Log("Warning: Could not locate localized format prompt title string in '%s': %u", mui_path, LASTERR);
 384         }
 385
 386         if (LoadStringU(hMui, 4126, g_FormatDiskButton, sizeof(g_FormatDiskButton)) <= 0) {
 387                 static_strcpy(g_FormatDiskButton, "Format disk");
 388                 Log("Warning: Could not locate localized format prompt button string in '%s': %u", mui_path, LASTERR);
 389         }
 390
 391         // 32964 = "Location is not available"
 392         if (LoadStringU(hMui, 32964, g_LocNotAvaliableTitle, sizeof(g_LocNotAvaliableTitle)) <= 0) {
 393                 static_strcpy(g_LocNotAvaliableTitle, "Location is not available");
 394                 Log("Warning: Could not locate localized format prompt title string in '%s': %u", mui_path, LASTERR);
 395         }
 396
 397         if (!LoadDialogCaption(hMui, 1024, g_InsertDiskTitle, sizeof(g_InsertDiskTitle)))
 398         {
 399                 static_strcpy(g_InsertDiskTitle, "Insert disk");
 400                 Log("Warning: Could not locate insert disk title string in '%s': %u", mui_path, LASTERR);
 401         }
 402
 403         FreeLibrary(hMui);
 404         return TRUE;
 405 }
 406
 407 static __inline int GetWindowTextU(HWND hWnd, char* lpString, int nMaxCount)
 408 {
 409         int ret = 0;
 410         DWORD err = ERROR_INVALID_DATA;
 411         if (nMaxCount < 0)
 412                 return 0;
 413         // Handle the empty string as GetWindowTextW() returns 0 then
 414         if ((lpString != NULL) && (nMaxCount > 0))
 415                 lpString[0] = 0;
 416         // coverity[returned_null]
 417         walloc(lpString, nMaxCount);
 418         ret = GetWindowTextW(hWnd, wlpString, nMaxCount);
 419         err = GetLastError();
 420         // coverity[var_deref_model]
 421         if ((ret != 0) && ((ret = wchar_to_utf8_no_alloc(wlpString, lpString, nMaxCount)) == 0)) {
 422                 err = GetLastError();
 423         }
 424         wfree(lpString);
 425         lpString[nMaxCount - 1] = 0;
 426         SetLastError(err);
 427         return ret;
 428 }
 429
 430
 431 /*
 432  * The following function calls are used to automatically detect and close the native
 433  * Windows format prompt "You must format the disk in drive X:". To do that, we use
 434  * an event hook that gets triggered whenever a window is placed in the foreground.
 435  * In that hook, we look for a dialog that has style WS_POPUPWINDOW and has the relevant
 436  * title. However, because the title in itself is too generic (the expectation is that
 437  * it will be "Microsoft Windows") we also enumerate all the child controls from that
 438  * prompt, using another callback, until we find one that contains the text we expect
 439  * for the "Format disk" button.
 440  * Oh, and since all of these strings are localized, we must first pick them up from
 441  * the relevant mui's.
 442  */
 443 static BOOL CALLBACK AlertPromptCallback(HWND hWnd, LPARAM lParam)
 444 {
 445         char str[128];
 446         BOOL* found = (BOOL*)lParam;
 447
 448         if (GetWindowTextU(hWnd, str, sizeof(str)) == 0)
 449                 return TRUE;
 450         if (strcmp(str, g_FormatDiskButton) == 0)
 451                 *found = TRUE;
 452         return TRUE;
 453 }
 454
 455 static volatile BOOL g_AlertPromptHookEnable = FALSE;
 456
 457 void SetAlertPromptHookEnable(BOOL enable)
 458 {
 459         g_AlertPromptHookEnable = enable;
 460 }
 461
 462 static void CALLBACK AlertPromptHook(HWINEVENTHOOK hWinEventHook, DWORD Event, HWND hWnd, LONG idObject, LONG idChild, DWORD dwEventThread, DWORD dwmsEventTime)
 463 {
 464         char str[128];
 465         BOOL found;
 466
 467         if (Event != EVENT_SYSTEM_FOREGROUND)
 468         {
 469                 return;
 470         }
 471
 472         if (!g_AlertPromptHookEnable)
 473         {
 474                 return;
 475         }
 476
 477         //GetWindowTextU(hWnd, str, sizeof(str));
 478         //Log("###### EVENT_SYSTEM_FOREGROUND Windows prompt <%s> #######", str);
 479
 480         if (GetWindowLongPtr(hWnd, GWL_STYLE) & WS_POPUPWINDOW) {
 481                 str[0] = 0;
 482                 GetWindowTextU(hWnd, str, sizeof(str));
 483                 if (strcmp(str, g_FormatDiskTitle) == 0)
 484                 {
 485                         found = FALSE;
 486                         EnumChildWindows(hWnd, AlertPromptCallback, (LPARAM)&found);
 487                         if (found)
 488                         {
 489                                 SendMessage(hWnd, WM_COMMAND, (WPARAM)IDCANCEL, (LPARAM)0);
 490                                 Log("###### Detect 'Windows format' prompt, now close it. #######");
 491                         }
 492                 }
 493                 else if (strcmp(str, g_LocNotAvaliableTitle) == 0)
 494                 {
 495                         SendMessage(hWnd, WM_COMMAND, (WPARAM)IDCANCEL, (LPARAM)0);
 496                         Log("###### Detect 'Location is not available' prompt, now close it. #######");
 497                 }
 498                 else if (strcmp(str, g_InsertDiskTitle) == 0)
 499                 {
 500                         SendMessage(hWnd, WM_COMMAND, (WPARAM)IDCANCEL, (LPARAM)0);
 501                         Log("###### Detect 'Insert disk' prompt, now close it. #######");
 502                 }
 503         }
 504 }
 505
 506
 507 BOOL SetAlertPromptHook(void)
 508 {
 509         if (ap_weh != NULL)
 510                 return TRUE;    // No need to set again if active
 511         ap_weh = SetWinEventHook(EVENT_SYSTEM_FOREGROUND, EVENT_SYSTEM_FOREGROUND, NULL,
 512                 AlertPromptHook, 0, 0, WINEVENT_OUTOFCONTEXT | WINEVENT_SKIPOWNPROCESS);
 513         return (ap_weh != NULL);
 514 }
 515
 516 BOOL AlertSuppressInit(void)
 517 {
 518         BOOL bRet;
 519
 520         SetAlertPromptMessages();
 521
 522         bRet = SetAlertPromptHook();
 523         Log("SetAlertPromptHook %s", bRet ? "SUCCESS" : "FAILED");
 524
 525         return TRUE;
 526 }