Ventoy2Disk/Ventoy2Disk/DiskService.c

   1 /******************************************************************************
   2 * DiskService.c
   3 *
   4 * Copyright (c) 2021, longpanda <admin@ventoy.net>
   5 * Copyright (c) 2011-2021 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 #include "DiskService.h"
  29
  30 static CHAR g_WindowsDir[MAX_PATH] = {0};
  31
  32 const CHAR* DISK_GetWindowsDir(void)
  33 {
  34         if (g_WindowsDir[0] == 0)
  35         {
  36                 GetEnvironmentVariableA("SystemRoot", g_WindowsDir, MAX_PATH);
  37                 if (g_WindowsDir[0] == 0)
  38                 {
  39                         sprintf_s(g_WindowsDir, MAX_PATH, "C:\\Windows");
  40                 }
  41         }
  42
  43         return g_WindowsDir;
  44 }
  45
  46 BOOL DISK_CleanDisk(int DriveIndex)
  47 {
  48         BOOL ret;
  49
  50         ret = VDS_CleanDisk(DriveIndex);
  51         if (!ret)
  52         {
  53                 ret = PSHELL_CleanDisk(DriveIndex);
  54         }
  55
  56         return ret;
  57 }
  58
  59
  60 BOOL DISK_DeleteVtoyEFIPartition(int DriveIndex, UINT64 EfiPartOffset)
  61 {
  62         BOOL ret;
  63
  64         ret = VDS_DeleteVtoyEFIPartition(DriveIndex, EfiPartOffset);
  65         if (!ret)
  66         {
  67                 ret = PSHELL_DeleteVtoyEFIPartition(DriveIndex, EfiPartOffset);
  68         }
  69
  70         return ret;
  71 }
  72
  73 BOOL DISK_ChangeVtoyEFI2ESP(int DriveIndex, UINT64 Offset)
  74 {
  75         BOOL ret;
  76
  77         ret = VDS_ChangeVtoyEFI2ESP(DriveIndex, Offset);
  78         if (!ret)
  79         {
  80                 ret = PSHELL_ChangeVtoyEFI2ESP(DriveIndex, Offset);
  81         }
  82
  83         return ret;
  84 }
  85
  86
  87 BOOL DISK_ChangeVtoyEFI2Basic(int DriveIndex, UINT64 Offset)
  88 {
  89         BOOL ret;
  90
  91         ret = VDS_ChangeVtoyEFI2Basic(DriveIndex, Offset);
  92         if (!ret)
  93         {
  94                 ret = PSHELL_ChangeVtoyEFI2Basic(DriveIndex, Offset);
  95         }
  96
  97         return ret;
  98 }
  99
 100 BOOL DISK_ChangeVtoyEFIAttr(int DriveIndex, UINT64 Offset, UINT64 Attr)
 101 {
 102         BOOL ret;
 103
 104         ret = VDS_ChangeVtoyEFIAttr(DriveIndex, Offset, Attr);
 105
 106         return ret;
 107 }
 108
 109 BOOL DISK_ShrinkVolume(int DriveIndex, const char* VolumeGuid, CHAR DriveLetter, UINT64 OldBytes, UINT64 ReduceBytes)
 110 {
 111         BOOL ret;
 112
 113         ret = VDS_ShrinkVolume(DriveIndex, VolumeGuid, DriveLetter, OldBytes, ReduceBytes);
 114         if (!ret)
 115         {
 116                 if (LASTERR == VDS_E_SHRINK_DIRTY_VOLUME)
 117                 {
 118                         Log("VDS shrink return dirty, no need to run powershell.");
 119                 }
 120                 else
 121                 {
 122                         ret = PSHELL_ShrinkVolume(DriveIndex, VolumeGuid, DriveLetter, OldBytes, ReduceBytes);
 123                 }
 124         }
 125
 126         return ret;
 127 }
 128
 129
 130
 131 // Output command
 132 typedef struct
 133 {
 134         DWORD Lines;
 135         PCHAR Output;
 136 } TEXTOUTPUT, * PTEXTOUTPUT;
 137
 138 /* Callback command types (some errorcode were filled from HPUSBFW V2.2.3 and their
 139    designation from docs.microsoft.com/windows/win32/api/vds/nf-vds-ivdsvolumemf2-formatex */
 140 typedef enum {
 141         FCC_PROGRESS,
 142         FCC_DONE_WITH_STRUCTURE,
 143         FCC_UNKNOWN2,
 144         FCC_INCOMPATIBLE_FILE_SYSTEM,
 145         FCC_UNKNOWN4,
 146         FCC_UNKNOWN5,
 147         FCC_ACCESS_DENIED,
 148         FCC_MEDIA_WRITE_PROTECTED,
 149         FCC_VOLUME_IN_USE,
 150         FCC_CANT_QUICK_FORMAT,
 151         FCC_UNKNOWNA,
 152         FCC_DONE,
 153         FCC_BAD_LABEL,
 154         FCC_UNKNOWND,
 155         FCC_OUTPUT,
 156         FCC_STRUCTURE_PROGRESS,
 157         FCC_CLUSTER_SIZE_TOO_SMALL,
 158         FCC_CLUSTER_SIZE_TOO_BIG,
 159         FCC_VOLUME_TOO_SMALL,
 160         FCC_VOLUME_TOO_BIG,
 161         FCC_NO_MEDIA_IN_DRIVE,
 162         FCC_UNKNOWN15,
 163         FCC_UNKNOWN16,
 164         FCC_UNKNOWN17,
 165         FCC_DEVICE_NOT_READY,
 166         FCC_CHECKDISK_PROGRESS,
 167         FCC_UNKNOWN1A,
 168         FCC_UNKNOWN1B,
 169         FCC_UNKNOWN1C,
 170         FCC_UNKNOWN1D,
 171         FCC_UNKNOWN1E,
 172         FCC_UNKNOWN1F,
 173         FCC_READ_ONLY_MODE,
 174         FCC_UNKNOWN21,
 175         FCC_UNKNOWN22,
 176         FCC_UNKNOWN23,
 177         FCC_UNKNOWN24,
 178         FCC_ALIGNMENT_VIOLATION,
 179 } FILE_SYSTEM_CALLBACK_COMMAND;
 180
 181 // FMIFS callback definition
 182 typedef BOOLEAN(__stdcall* PFMIFSCALLBACK)(FILE_SYSTEM_CALLBACK_COMMAND Command, DWORD SubAction, PVOID ActionInfo);
 183
 184
 185 // Chkdsk command in FMIFS
 186 typedef VOID(__stdcall* PCHKDSK)(PWCHAR DriveRoot,
 187         PWCHAR Format,
 188         BOOL CorrectErrors,
 189         BOOL Verbose,
 190         BOOL CheckOnlyIfDirty,
 191         BOOL ScanDrive,
 192         PVOID Unused2,
 193         PVOID Unused3,
 194         PFMIFSCALLBACK Callback);
 195
 196
 197 // media flags
 198 #define FMIFS_HARDDISK 0xC
 199 #define FMIFS_FLOPPY   0x8
 200 // Format command in FMIFS
 201 typedef VOID(__stdcall* PFORMATEX)(PWCHAR DriveRoot,
 202         DWORD MediaFlag,
 203         PWCHAR Format,
 204         PWCHAR Label,
 205         BOOL QuickFormat,
 206         DWORD ClusterSize,
 207         PFMIFSCALLBACK Callback);
 208
 209 #define FP_FORCE                            0x00000001
 210 #define FP_QUICK                            0x00000002
 211 #define FP_COMPRESSION                      0x00000004
 212 #define FP_DUPLICATE_METADATA               0x00000008
 213 #define FP_LARGE_FAT32                      0x00010000
 214 #define FP_NO_BOOT                          0x00020000
 215 #define FP_CREATE_PERSISTENCE_CONF          0x00040000
 216
 217 // FormatExCallback
 218 static int g_dll_format_error = 0;
 219 BOOLEAN __stdcall FormatExCallback(FILE_SYSTEM_CALLBACK_COMMAND Command, DWORD Modifier, PVOID Argument)
 220 {
 221         PDWORD percent;
 222         PBOOLEAN status;
 223
 224         switch (Command) {
 225         case FCC_PROGRESS:
 226                 percent = (PDWORD)Argument;
 227                 Log("Format percent: %u%%", *percent);
 228                 break;
 229         case FCC_STRUCTURE_PROGRESS:    // No progress on quick format
 230                 Log("Creating file system...");
 231                 break;
 232         case FCC_DONE:
 233                 status = (PBOOLEAN)Argument;
 234                 if (*status == FALSE)
 235                 {
 236                         Log("Format error: %u ERROR_NOT_SUPPORTED=%u", LASTERR, ERROR_NOT_SUPPORTED);
 237                         g_dll_format_error = 1;
 238                 }
 239                 else
 240                 {
 241                         Log("Format Done");
 242                 }
 243                 break;
 244         case FCC_DONE_WITH_STRUCTURE:
 245                 Log("Format FCC_DONE_WITH_STRUCTURE");
 246                 break;
 247         case FCC_INCOMPATIBLE_FILE_SYSTEM:
 248                 Log("Incompatible File System");
 249                 break;
 250         case FCC_ACCESS_DENIED:
 251                 Log("Access denied");
 252                 break;
 253         case FCC_MEDIA_WRITE_PROTECTED:
 254                 Log("Media is write protected");
 255                 break;
 256         case FCC_VOLUME_IN_USE:
 257                 Log("Volume is in use");
 258                 break;
 259         case FCC_DEVICE_NOT_READY:
 260                 Log("The device is not ready");
 261                 break;
 262         case FCC_CANT_QUICK_FORMAT:
 263                 Log("Cannot quick format this volume");
 264                 break;
 265         case FCC_BAD_LABEL:
 266                 Log("Bad label");
 267                 break;
 268         case FCC_OUTPUT:
 269                 Log("%s", ((PTEXTOUTPUT)Argument)->Output);
 270                 break;
 271         case FCC_CLUSTER_SIZE_TOO_BIG:
 272         case FCC_CLUSTER_SIZE_TOO_SMALL:
 273                 Log("Unsupported cluster size");
 274                 break;
 275         case FCC_VOLUME_TOO_BIG:
 276         case FCC_VOLUME_TOO_SMALL:
 277                 Log("Volume is too %s", (Command == FCC_VOLUME_TOO_BIG) ? "big" : "small");
 278                 break;
 279         case FCC_NO_MEDIA_IN_DRIVE:
 280                 Log("No media in drive");
 281                 break;
 282         case FCC_ALIGNMENT_VIOLATION:
 283                 Log("Partition start offset is not aligned to the cluster size");
 284                 break;
 285         default:
 286                 Log("FormatExCallback: Received unhandled command 0x%02X - aborting", Command);
 287                 break;
 288         }
 289
 290         return TRUE;
 291 }
 292
 293
 294 BOOL DLL_FormatVolume(char DriveLetter, int fs, DWORD ClusterSize)
 295 {
 296         PWCHAR  Label = L"Ventoy";
 297         PWCHAR  Format = NULL;
 298         WCHAR   RootDirectory[MAX_PATH] = { 0 };
 299         HMODULE ifsModule;
 300         PFORMATEX FormatEx;
 301
 302         ifsModule = LoadLibraryA("fmifs.dll");
 303         if (NULL == ifsModule)
 304         {
 305                 Log("LoadLibrary fmifs.dll failed %u", LASTERR);
 306                 return FALSE;
 307         }
 308
 309         Log("Find ifsModule");
 310
 311         FormatEx = (PFORMATEX)GetProcAddress(ifsModule, "FormatEx");
 312         if (FormatEx == NULL)
 313         {
 314                 Log("Failed to get FormatEx handler\n");
 315                 return FALSE;
 316         }
 317         Log("Find FormatEx=%p", FormatEx);
 318
 319         RootDirectory[0] = DriveLetter;
 320         RootDirectory[1] = L':';
 321         RootDirectory[2] = L'\\';
 322         RootDirectory[3] = (WCHAR)0;
 323
 324         DWORD media;
 325         DWORD driveType;
 326         driveType = GetDriveTypeW(RootDirectory);
 327         if (driveType != DRIVE_FIXED)
 328                 media = FMIFS_FLOPPY;
 329         if (driveType == DRIVE_FIXED)
 330                 media = FMIFS_HARDDISK;
 331
 332         Format = GetVentoyFsFmtNameByTypeW(fs);
 333
 334         g_dll_format_error = 0;
 335
 336         Log("Call FormatEx Function for %C: %s ClusterSize=%u(%uKB)", DriveLetter, GetVentoyFsFmtNameByTypeA(fs), ClusterSize, ClusterSize / 1024);
 337         FormatEx(RootDirectory, media, Format, Label, FP_FORCE | FP_QUICK, ClusterSize, FormatExCallback);
 338         FreeLibrary(ifsModule);
 339
 340         if (g_dll_format_error)
 341         {
 342                 Log("Format failed by DLL");
 343                 return FALSE;
 344         }
 345
 346         Log("Format success by DLL");
 347         return TRUE;
 348 }
 349
 350
 351 BOOL DISK_FormatVolume(char DriveLetter, int fs, UINT64 VolumeSize)
 352 {
 353         int i;
 354         DWORD ClusterSize = 0;
 355         BOOL ret = FALSE;
 356         FmtFunc astFmtFunc[] =
 357         {
 358                 FMT_DEF(VDS_FormatVolume),
 359                 FMT_DEF(DLL_FormatVolume),
 360                 FMT_DEF(PSHELL_FormatVolume),
 361                 FMT_DEF(DSPT_FormatVolume),
 362                 FMT_DEF(CMD_FormatVolume),
 363                 { NULL, NULL }
 364         };
 365
 366         ClusterSize = (DWORD)GetClusterSize();
 367         Log("DISK_FormatVolume %C:\\ %s VolumeSize=%llu ClusterSize=%u(%uKB)",
 368                 DriveLetter, GetVentoyFsNameByType(fs), (ULONGLONG)VolumeSize, ClusterSize, ClusterSize / 1024);
 369
 370         for (i = 0; astFmtFunc[i].formatFunc; i++)
 371         {
 372                 Log("%s ...", astFmtFunc[i].name);
 373                 ret = astFmtFunc[i].formatFunc(DriveLetter, fs, ClusterSize);
 374                 if (ret)
 375                 {
 376                         break;
 377                 }
 378         }
 379
 380         return ret;
 381 }