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 title_str[2][128], button_str[128];
  77 static char system_dir[MAX_PATH], sysnative_dir[MAX_PATH];
  78
  79 static HWINEVENTHOOK ap_weh = NULL;
  80
  81
  82 static __inline UINT GetSystemDirectoryU(char* lpBuffer, UINT uSize)
  83 {
  84         UINT ret = 0, err = ERROR_INVALID_DATA;
  85         // coverity[returned_null]
  86         walloc(lpBuffer, uSize);
  87         ret = GetSystemDirectoryW(wlpBuffer, uSize);
  88         err = GetLastError();
  89         if ((ret != 0) && ((ret = wchar_to_utf8_no_alloc(wlpBuffer, lpBuffer, uSize)) == 0)) {
  90                 err = GetLastError();
  91         }
  92         wfree(lpBuffer);
  93         SetLastError(err);
  94         return ret;
  95 }
  96
  97
  98 static char* ToLocaleName(DWORD lang_id)
  99 {
 100         static char mui_str[LOCALE_NAME_MAX_LENGTH];
 101         wchar_t wmui_str[LOCALE_NAME_MAX_LENGTH];
 102
 103         if (LCIDToLocaleName(lang_id, wmui_str, LOCALE_NAME_MAX_LENGTH, 0) > 0) {
 104                 wchar_to_utf8_no_alloc(wmui_str, mui_str, LOCALE_NAME_MAX_LENGTH);
 105         }
 106         else {
 107                 static_strcpy(mui_str, "en-US");
 108         }
 109         return mui_str;
 110 }
 111
 112 static __inline int LoadStringU(HINSTANCE hInstance, UINT uID, LPSTR lpBuffer, int nBufferMax)
 113 {
 114         int ret;
 115         DWORD err = ERROR_INVALID_DATA;
 116         if (nBufferMax == 0) {
 117                 // read-only pointer to resource mode is not supported
 118                 SetLastError(ERROR_INVALID_PARAMETER);
 119                 return 0;
 120         }
 121         // coverity[returned_null]
 122         walloc(lpBuffer, nBufferMax);
 123         ret = LoadStringW(hInstance, uID, wlpBuffer, nBufferMax);
 124         err = GetLastError();
 125         if ((ret > 0) && ((ret = wchar_to_utf8_no_alloc(wlpBuffer, lpBuffer, nBufferMax)) == 0)) {
 126                 err = GetLastError();
 127         }
 128         wfree(lpBuffer);
 129         SetLastError(err);
 130         return ret;
 131 }
 132
 133 /*
 134  * typedefs for the function prototypes. Use the something like:
 135  *   PF_DECL(FormatEx);
 136  * which translates to:
 137  *   FormatEx_t pfFormatEx = NULL;
 138  * in your code, to declare the entrypoint and then use:
 139  *   PF_INIT(FormatEx, Fmifs);
 140  * which translates to:
 141  *   pfFormatEx = (FormatEx_t) GetProcAddress(GetDLLHandle("fmifs"), "FormatEx");
 142  * to make it accessible.
 143  */
 144 #define         MAX_LIBRARY_HANDLES 64
 145 extern HMODULE  OpenedLibrariesHandle[MAX_LIBRARY_HANDLES];
 146 extern UINT16 OpenedLibrariesHandleSize;
 147 #define         OPENED_LIBRARIES_VARS HMODULE OpenedLibrariesHandle[MAX_LIBRARY_HANDLES]; uint16_t OpenedLibrariesHandleSize = 0
 148 #define         CLOSE_OPENED_LIBRARIES while(OpenedLibrariesHandleSize > 0) FreeLibrary(OpenedLibrariesHandle[--OpenedLibrariesHandleSize])
 149
 150 static __inline HMODULE GetLibraryHandle(char* szLibraryName) {
 151         HMODULE h = NULL;
 152         wchar_t* wszLibraryName = NULL;
 153         int size;
 154         if (szLibraryName == NULL || szLibraryName[0] == 0)
 155                 goto out;
 156         size = MultiByteToWideChar(CP_UTF8, 0, szLibraryName, -1, NULL, 0);
 157         if ((size <= 1) || ((wszLibraryName = (wchar_t*)calloc(size, sizeof(wchar_t))) == NULL) ||
 158                 (MultiByteToWideChar(CP_UTF8, 0, szLibraryName, -1, wszLibraryName, size) != size))
 159                 goto out;
 160         // If the library is already opened, just return a handle (that doesn't need to be freed)
 161         if ((h = GetModuleHandleW(wszLibraryName)) != NULL)
 162                 goto out;
 163         // Sanity check
 164         if (OpenedLibrariesHandleSize >= MAX_LIBRARY_HANDLES) {
 165                 Log("Error: MAX_LIBRARY_HANDLES is too small");
 166                 goto out;
 167         }
 168         h = LoadLibraryExW(wszLibraryName, NULL, LOAD_LIBRARY_SEARCH_SYSTEM32);
 169         // Some Windows 7 platforms (most likely the ones missing KB2533623 per the
 170         // official LoadLibraryEx doc) can return ERROR_INVALID_PARAMETER when using
 171         // the Ex() version. If that's the case, fallback to using LoadLibraryW().
 172         if ((h == NULL) && (SCODE_CODE(GetLastError()) == ERROR_INVALID_PARAMETER))
 173                 h = LoadLibraryW(wszLibraryName);
 174         if (h != NULL)
 175                 OpenedLibrariesHandle[OpenedLibrariesHandleSize++] = h;
 176         else
 177                 Log("Unable to load '%S.dll': %u", wszLibraryName, LASTERR);
 178 out:
 179         free(wszLibraryName);
 180         return h;
 181 }
 182
 183 #define PF_TYPE(api, ret, proc, args)           typedef ret (api *proc##_t)args
 184 #define PF_DECL(proc)                                           static proc##_t pf##proc = NULL
 185 #define PF_TYPE_DECL(api, ret, proc, args)      PF_TYPE(api, ret, proc, args); PF_DECL(proc)
 186 #define PF_INIT(proc, name)                                     if (pf##proc == NULL) pf##proc = \
 187         (proc##_t) GetProcAddress(GetLibraryHandle(#name), #proc)
 188 #define PF_INIT_OR_OUT(proc, name)                      do {PF_INIT(proc, name);         \
 189         if (pf##proc == NULL) {Log("Unable to locate %s() in '%s.dll': %u",  \
 190         #proc, #name, LASTERR); goto out;} } while(0)
 191 #define PF_INIT_OR_SET_STATUS(proc, name)       do {PF_INIT(proc, name);         \
 192         if ((pf##proc == NULL) && (NT_SUCCESS(status))) status = STATUS_NOT_IMPLEMENTED; } while(0)
 193
 194 static BOOL is_x64(void)
 195 {
 196         BOOL ret = FALSE;
 197         PF_TYPE_DECL(WINAPI, BOOL, IsWow64Process, (HANDLE, PBOOL));
 198         // Detect if we're running a 32 or 64 bit system
 199         if (sizeof(uintptr_t) < 8) {
 200                 PF_INIT(IsWow64Process, Kernel32);
 201                 if (pfIsWow64Process != NULL) {
 202                         (*pfIsWow64Process)(GetCurrentProcess(), &ret);
 203                 }
 204         }
 205         else {
 206                 ret = TRUE;
 207         }
 208         return ret;
 209 }
 210
 211
 212 static __inline UINT GetSystemWindowsDirectoryU(char* lpBuffer, UINT uSize)
 213 {
 214         UINT ret = 0, err = ERROR_INVALID_DATA;
 215         // coverity[returned_null]
 216         walloc(lpBuffer, uSize);
 217         ret = GetSystemWindowsDirectoryW(wlpBuffer, uSize);
 218         err = GetLastError();
 219         if ((ret != 0) && ((ret = wchar_to_utf8_no_alloc(wlpBuffer, lpBuffer, uSize)) == 0)) {
 220                 err = GetLastError();
 221         }
 222         wfree(lpBuffer);
 223         SetLastError(err);
 224         return ret;
 225 }
 226
 227 static __inline HMODULE LoadLibraryU(LPCSTR lpFileName)
 228 {
 229         HMODULE ret;
 230         DWORD err = ERROR_INVALID_DATA;
 231         wconvert(lpFileName);
 232         ret = LoadLibraryW(wlpFileName);
 233         err = GetLastError();
 234         wfree(lpFileName);
 235         SetLastError(err);
 236         return ret;
 237 }
 238
 239
 240 BOOL SetAlertPromptMessages(void)
 241 {
 242         HMODULE hMui;
 243         char mui_path[MAX_PATH];
 244
 245         if (GetSystemDirectoryU(system_dir, sizeof(system_dir)) == 0) {
 246                 Log("Could not get system directory: %u", LASTERR);
 247                 static_strcpy(system_dir, "C:\\Windows\\System32");
 248         }
 249
 250         // Construct Sysnative ourselves as there is no GetSysnativeDirectory() call
 251         // By default (64bit app running on 64 bit OS or 32 bit app running on 32 bit OS)
 252         // Sysnative and System32 are the same
 253         static_strcpy(sysnative_dir, system_dir);
 254         // But if the app is 32 bit and the OS is 64 bit, Sysnative must differ from System32
 255 #if (defined(VTARCH_X86) || defined(VTARCH_ARM))
 256         if (is_x64()) {
 257                 if (GetSystemWindowsDirectoryU(sysnative_dir, sizeof(sysnative_dir)) == 0) {
 258                         Log("Could not get Windows directory: %u", LASTERR);
 259                         static_strcpy(sysnative_dir, "C:\\Windows");
 260                 }
 261                 static_strcat(sysnative_dir, "\\Sysnative");
 262         }
 263 #endif
 264
 265         Log("system_dir=<%s>", system_dir);
 266         Log("sysnative_dir=<%s>", sysnative_dir);
 267
 268         sprintf_s(mui_path, MAX_PATH, "%s\\%s\\shell32.dll.mui", sysnative_dir, ToLocaleName(GetUserDefaultUILanguage()));
 269
 270         hMui = LoadLibraryU(mui_path);
 271         if (hMui)
 272         {
 273                 // 4097 = "You need to format the disk in drive %c: before you can use it." (dialog text)
 274                 // 4125 = "Microsoft Windows" (dialog title)
 275                 // 4126 = "Format disk" (button)
 276                 if (LoadStringU(hMui, 4125, title_str[0], sizeof(title_str[0])) <= 0) {
 277                         static_strcpy(title_str[0], "Microsoft Windows");
 278                         Log("Warning: Could not locate localized format prompt title string in '%s': %u", mui_path, LASTERR);
 279                 }
 280                 if (LoadStringU(hMui, 4126, button_str, sizeof(button_str)) <= 0) {
 281                         static_strcpy(button_str, "Format disk");
 282                         Log("Warning: Could not locate localized format prompt button string in '%s': %u", mui_path, LASTERR);
 283                 }
 284                 FreeLibrary(hMui);
 285
 286                 Log("LoadLibrary shell32.dll.mui SUCCESS");
 287         }
 288         else
 289         {
 290                 Log("LoadLibrary shell32.dll.mui FAILED");
 291         }
 292
 293         return TRUE;
 294 }
 295
 296 static __inline int GetWindowTextU(HWND hWnd, char* lpString, int nMaxCount)
 297 {
 298         int ret = 0;
 299         DWORD err = ERROR_INVALID_DATA;
 300         if (nMaxCount < 0)
 301                 return 0;
 302         // Handle the empty string as GetWindowTextW() returns 0 then
 303         if ((lpString != NULL) && (nMaxCount > 0))
 304                 lpString[0] = 0;
 305         // coverity[returned_null]
 306         walloc(lpString, nMaxCount);
 307         ret = GetWindowTextW(hWnd, wlpString, nMaxCount);
 308         err = GetLastError();
 309         // coverity[var_deref_model]
 310         if ((ret != 0) && ((ret = wchar_to_utf8_no_alloc(wlpString, lpString, nMaxCount)) == 0)) {
 311                 err = GetLastError();
 312         }
 313         wfree(lpString);
 314         lpString[nMaxCount - 1] = 0;
 315         SetLastError(err);
 316         return ret;
 317 }
 318
 319
 320 /*
 321  * The following function calls are used to automatically detect and close the native
 322  * Windows format prompt "You must format the disk in drive X:". To do that, we use
 323  * an event hook that gets triggered whenever a window is placed in the foreground.
 324  * In that hook, we look for a dialog that has style WS_POPUPWINDOW and has the relevant
 325  * title. However, because the title in itself is too generic (the expectation is that
 326  * it will be "Microsoft Windows") we also enumerate all the child controls from that
 327  * prompt, using another callback, until we find one that contains the text we expect
 328  * for the "Format disk" button.
 329  * Oh, and since all of these strings are localized, we must first pick them up from
 330  * the relevant mui's.
 331  */
 332 static BOOL CALLBACK AlertPromptCallback(HWND hWnd, LPARAM lParam)
 333 {
 334         char str[128];
 335         BOOL* found = (BOOL*)lParam;
 336
 337         if (GetWindowTextU(hWnd, str, sizeof(str)) == 0)
 338                 return TRUE;
 339         if (strcmp(str, button_str) == 0)
 340                 *found = TRUE;
 341         return TRUE;
 342 }
 343
 344 static void CALLBACK AlertPromptHook(HWINEVENTHOOK hWinEventHook, DWORD Event, HWND hWnd, LONG idObject, LONG idChild, DWORD dwEventThread, DWORD dwmsEventTime)
 345 {
 346         char str[128];
 347         BOOL found;
 348
 349         if (Event == EVENT_SYSTEM_FOREGROUND) {
 350                 if (GetWindowLongPtr(hWnd, GWL_STYLE) & WS_POPUPWINDOW) {
 351                         str[0] = 0;
 352                         GetWindowTextU(hWnd, str, sizeof(str));
 353                         if (strcmp(str, title_str[0]) == 0) {
 354                                 found = FALSE;
 355                                 EnumChildWindows(hWnd, AlertPromptCallback, (LPARAM)&found);
 356                                 if (found) {
 357                                         SendMessage(hWnd, WM_COMMAND, (WPARAM)IDCANCEL, (LPARAM)0);
 358                                         Log("###### Closed Windows format prompt #######");
 359                                 }
 360                         }
 361                 }
 362         }
 363 }
 364
 365
 366 BOOL SetAlertPromptHook(void)
 367 {
 368         if (ap_weh != NULL)
 369                 return TRUE;    // No need to set again if active
 370         ap_weh = SetWinEventHook(EVENT_SYSTEM_FOREGROUND, EVENT_SYSTEM_FOREGROUND, NULL,
 371                 AlertPromptHook, 0, 0, WINEVENT_OUTOFCONTEXT | WINEVENT_SKIPOWNPROCESS);
 372         return (ap_weh != NULL);
 373 }
 374
 375 BOOL AlertSuppressInit(void)
 376 {
 377         BOOL bRet;
 378
 379         SetAlertPromptMessages();
 380
 381         bRet = SetAlertPromptHook();
 382         Log("SetAlertPromptHook %s", bRet ? "SUCCESS" : "FAILED");
 383
 384         return TRUE;
 385 }