Ventoy2Disk/Ventoy2Disk/DiskService.c

   1 /******************************************************************************
   2 * DiskService.c
   3 *
   4 * Copyright (c) 2021, longpanda <admin@ventoy.net>
   5 *
   6 * This program is free software; you can redistribute it and/or
   7 * modify it under the terms of the GNU General Public License as
   8 * published by the Free Software Foundation; either version 3 of the
   9 * License, or (at your option) any later version.
  10 *
  11 * This program is distributed in the hope that it will be useful, but
  12 * WITHOUT ANY WARRANTY; without even the implied warranty of
  13 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14 * General Public License for more details.
  15 *
  16 * You should have received a copy of the GNU General Public License
  17 * along with this program; if not, see <http://www.gnu.org/licenses/>.
  18 *
  19 */
  20
  21 #include <Windows.h>
  22 #include <winternl.h>
  23 #include <commctrl.h>
  24 #include <initguid.h>
  25 #include <vds.h>
  26 #include "Ventoy2Disk.h"
  27 #include "DiskService.h"
  28
  29 BOOL DISK_CleanDisk(int DriveIndex)
  30 {
  31         BOOL ret;
  32
  33         ret = VDS_CleanDisk(DriveIndex);
  34         if (!ret)
  35         {
  36                 ret = PSHELL_CleanDisk(DriveIndex);
  37         }
  38
  39         return ret;
  40 }
  41
  42
  43 BOOL DISK_DeleteVtoyEFIPartition(int DriveIndex, UINT64 EfiPartOffset)
  44 {
  45         BOOL ret;
  46
  47         ret = VDS_DeleteVtoyEFIPartition(DriveIndex, EfiPartOffset);
  48         if (!ret)
  49         {
  50                 ret = PSHELL_DeleteVtoyEFIPartition(DriveIndex, EfiPartOffset);
  51         }
  52
  53         return ret;
  54 }
  55
  56 BOOL DISK_ChangeVtoyEFI2ESP(int DriveIndex, UINT64 Offset)
  57 {
  58         BOOL ret;
  59
  60         ret = VDS_ChangeVtoyEFI2ESP(DriveIndex, Offset);
  61         if (!ret)
  62         {
  63                 ret = PSHELL_ChangeVtoyEFI2ESP(DriveIndex, Offset);
  64         }
  65
  66         return ret;
  67 }
  68
  69
  70 BOOL DISK_ChangeVtoyEFI2Basic(int DriveIndex, UINT64 Offset)
  71 {
  72         BOOL ret;
  73
  74         ret = VDS_ChangeVtoyEFI2Basic(DriveIndex, Offset);
  75         if (!ret)
  76         {
  77                 ret = PSHELL_ChangeVtoyEFI2Basic(DriveIndex, Offset);
  78         }
  79
  80         return ret;
  81 }
  82
  83 BOOL DISK_ChangeVtoyEFIAttr(int DriveIndex, UINT64 Offset, UINT64 Attr)
  84 {
  85         BOOL ret;
  86
  87         ret = VDS_ChangeVtoyEFIAttr(DriveIndex, Offset, Attr);
  88
  89         return ret;
  90 }
  91
  92 BOOL DISK_ShrinkVolume(int DriveIndex, const char* VolumeGuid, CHAR DriveLetter, UINT64 OldBytes, UINT64 ReduceBytes)
  93 {
  94         BOOL ret;
  95
  96         ret = VDS_ShrinkVolume(DriveIndex, VolumeGuid, DriveLetter, OldBytes, ReduceBytes);
  97         if (!ret)
  98         {
  99                 if (LASTERR == VDS_E_SHRINK_DIRTY_VOLUME)
 100                 {
 101                         Log("VDS shrink return dirty, no need to run powershell.");
 102                 }
 103                 else
 104                 {
 105                         ret = PSHELL_ShrinkVolume(DriveIndex, VolumeGuid, DriveLetter, OldBytes, ReduceBytes);
 106                 }
 107         }
 108
 109         return ret;
 110 }
 111
 112
 113
 114 // Output command
 115 typedef struct
 116 {
 117         DWORD Lines;
 118         PCHAR Output;
 119 } TEXTOUTPUT, * PTEXTOUTPUT;
 120
 121 // Callback command types
 122 typedef enum
 123 {
 124         PROGRESS,
 125         DONEWITHSTRUCTURE,
 126         UNKNOWN2,
 127         UNKNOWN3,
 128         UNKNOWN4,
 129         UNKNOWN5,
 130         INSUFFICIENTRIGHTS,
 131         UNKNOWN7,
 132         UNKNOWN8,
 133         UNKNOWN9,
 134         UNKNOWNA,
 135         DONE,
 136         UNKNOWNC,
 137         UNKNOWND,
 138         OUTPUT,
 139         STRUCTUREPROGRESS
 140 } CALLBACKCOMMAND;
 141
 142 // FMIFS callback definition
 143 typedef BOOLEAN(__stdcall* PFMIFSCALLBACK)(CALLBACKCOMMAND Command, DWORD SubAction, PVOID ActionInfo);
 144
 145
 146 // Chkdsk command in FMIFS
 147 typedef VOID(__stdcall* PCHKDSK)(PWCHAR DriveRoot,
 148         PWCHAR Format,
 149         BOOL CorrectErrors,
 150         BOOL Verbose,
 151         BOOL CheckOnlyIfDirty,
 152         BOOL ScanDrive,
 153         PVOID Unused2,
 154         PVOID Unused3,
 155         PFMIFSCALLBACK Callback);
 156
 157
 158 // media flags
 159 #define FMIFS_HARDDISK 0xC
 160 #define FMIFS_FLOPPY   0x8
 161 // Format command in FMIFS
 162 typedef VOID(__stdcall* PFORMATEX)(PWCHAR DriveRoot,
 163         DWORD MediaFlag,
 164         PWCHAR Format,
 165         PWCHAR Label,
 166         BOOL QuickFormat,
 167         DWORD ClusterSize,
 168         PFMIFSCALLBACK Callback);
 169
 170
 171 // FormatExCallback
 172 static int g_dll_format_error = 0;
 173 BOOLEAN __stdcall FormatExCallback(CALLBACKCOMMAND Command, DWORD Modifier, PVOID Argument)
 174 {
 175         PDWORD percent;
 176         PBOOLEAN status;
 177
 178         switch (Command)
 179         {
 180         case PROGRESS:
 181                 percent = (PDWORD)Argument;
 182                 Log("Format percent: %d \n", *percent);
 183                 break;
 184
 185         case OUTPUT:
 186                 break;
 187
 188         case DONE:
 189                 status = (PBOOLEAN)Argument;
 190                 if (*status == FALSE)
 191                 {
 192                         g_dll_format_error = 1;
 193                 }
 194                 else
 195                 {
 196                 }
 197
 198                 break;
 199         }
 200         return TRUE;
 201 }
 202
 203
 204 BOOL DLL_FormatVolume(char DriveLetter, int fs, DWORD ClusterSize)
 205 {
 206         PWCHAR  Label = L"Ventoy";
 207         PWCHAR  Format = NULL;
 208         WCHAR   RootDirectory[MAX_PATH] = { 0 };
 209         HMODULE ifsModule;
 210         PFORMATEX FormatEx;
 211
 212         ifsModule = LoadLibraryA("fmifs.dll");
 213         if (NULL == ifsModule)
 214         {
 215                 Log("LoadLibrary fmifs.dll failed %u", LASTERR);
 216                 return FALSE;
 217         }
 218
 219         Log("Find ifsModule");
 220
 221         FormatEx = (PFORMATEX)GetProcAddress(ifsModule, "FormatEx");
 222         if (FormatEx == NULL)
 223         {
 224                 Log("Failed to get FormatEx handler\n");
 225                 return FALSE;
 226         }
 227         Log("Find FormatEx=%p", FormatEx);
 228
 229         RootDirectory[0] = DriveLetter;
 230         RootDirectory[1] = L':';
 231         RootDirectory[2] = L'\\';
 232         RootDirectory[3] = (WCHAR)0;
 233
 234         DWORD media;
 235         DWORD driveType;
 236         driveType = GetDriveTypeW(RootDirectory);
 237         if (driveType != DRIVE_FIXED)
 238                 media = FMIFS_FLOPPY;
 239         if (driveType == DRIVE_FIXED)
 240                 media = FMIFS_HARDDISK;
 241
 242         Format = GetVentoyFsFmtNameByTypeW(fs);
 243
 244         g_dll_format_error = 0;
 245
 246         Log("Call FormatEx Function for %C: %s ClusterSize=%u(%uKB)", DriveLetter, GetVentoyFsFmtNameByTypeA(fs), ClusterSize, ClusterSize / 1024);
 247         FormatEx(RootDirectory, media, Format, Label, TRUE, ClusterSize, FormatExCallback);
 248         FreeLibrary(ifsModule);
 249
 250         if (g_dll_format_error)
 251         {
 252                 Log("Format failed by DLL");
 253                 return FALSE;
 254         }
 255
 256         Log("Format success by DLL");
 257         return TRUE;
 258 }
 259
 260
 261 BOOL DISK_FormatVolume(char DriveLetter, int fs, UINT64 VolumeSize)
 262 {
 263         DWORD ClusterSize = 0;
 264         BOOL ret = FALSE;
 265
 266         ClusterSize = (DWORD)GetClusterSize();
 267         Log("DISK_FormatVolume %C:\\ %s VolumeSize=%llu ClusterSize=%u(%uKB)",
 268                 DriveLetter, GetVentoyFsNameByType(fs), (ULONGLONG)VolumeSize, ClusterSize, ClusterSize/1024);
 269
 270         ret = DLL_FormatVolume(DriveLetter, fs, ClusterSize);
 271         if (!ret)
 272         {
 273                 ret = VDS_FormatVolume(DriveLetter, fs, ClusterSize);
 274                 if (!ret)
 275                 {
 276                         ret = DSPT_FormatVolume(DriveLetter, fs, ClusterSize);
 277                         if (!ret)
 278                         {
 279                                 ret = PSHELL_FormatVolume(DriveLetter, fs, ClusterSize);
 280                         }
 281                 }
 282         }
 283
 284         return ret;
 285 }