--xdg improvement

[Ventoy.git] / vtoyjump / vtoyjump / vtoyjump.c

diff --git a/vtoyjump/vtoyjump/vtoyjump.c b/vtoyjump/vtoyjump/vtoyjump.c
index c2ab0a120e7856c0c402d4b1a83fc5b048f05b53..7c3d80377ce2012ad466ff0fce8aa160f1ea0c3f 100644 (file)
--- a/vtoyjump/vtoyjump/vtoyjump.c
+++ b/vtoyjump/vtoyjump/vtoyjump.c
@@ -1418,6 +1418,40 @@ End:
     return Ret; \r
 }\r
 \r
+static BOOL CheckVentoyDisk(DWORD DiskNum)\r
+{\r
+    DWORD dwSize = 0;\r
+    CHAR PhyPath[128];\r
+    UINT8 SectorBuf[512];\r
+    HANDLE Handle;\r
+    UINT8 check[8] = { 0x56, 0x54, 0x00, 0x47, 0x65, 0x00, 0x48, 0x44 };\r
+\r
+    sprintf_s(PhyPath, sizeof(PhyPath), "\\\\.\\PhysicalDrive%d", DiskNum);\r
+    Handle = CreateFileA(PhyPath, GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, 0, OPEN_EXISTING, 0, 0);\r
+    if (Handle == INVALID_HANDLE_VALUE)\r
+    {\r
+        Log("Could not open the disk<%s>, error:%u", PhyPath, GetLastError());\r
+        return FALSE;\r
+    }\r
+\r
+    if (!ReadFile(Handle, SectorBuf, sizeof(SectorBuf), &dwSize, NULL))\r
+    {\r
+        Log("ReadFile failed, dwSize:%u  error:%u", dwSize, GetLastError());\r
+        CloseHandle(Handle);\r
+        return FALSE;\r
+    }\r
+\r
+    CloseHandle(Handle);\r
+\r
+    if (memcmp(SectorBuf + 0x190, check, 8) == 0)\r
+    {\r
+        return TRUE;\r
+    }\r
+\r
+    return FALSE;\r
+}\r
+\r
+\r
 static int VentoyHook(ventoy_os_param *param)\r
 {\r
     int i;\r
@@ -1429,6 +1463,7 @@ static int VentoyHook(ventoy_os_param *param)
     CHAR VtoyLetter;\r
     DWORD Drives;\r
     DWORD NewDrives;\r
+    DWORD VtoyDiskNum;\r
     UINT32 DiskSig;\r
     UINT32 VtoySig;\r
        DISK_EXTENT DiskExtent;\r
@@ -1457,6 +1492,8 @@ static int VentoyHook(ventoy_os_param *param)
                 if (IsFileExist("%s", IsoPath))\r
                 {\r
                     Log("File exist under %C:", Letter);\r
+                    memset(UUID, 0, sizeof(UUID));\r
+                    memset(&DiskExtent, 0, sizeof(DiskExtent));\r
                     if (GetPhyDiskUUID(Letter, UUID, NULL, &DiskExtent) == 0)\r
                     {\r
                         if (memcmp(UUID, param->vtoy_disk_guid, 16) == 0)\r
@@ -1510,6 +1547,9 @@ static int VentoyHook(ventoy_os_param *param)
             {\r
                 if (Drives & 0x01)\r
                 {\r
+                    memset(UUID, 0, sizeof(UUID));\r
+                    memset(&VtoyDiskExtent, 0, sizeof(VtoyDiskExtent));\r
+                    DiskSig = 0;\r
                     if (GetPhyDiskUUID(VtoyLetter, UUID, &DiskSig, &VtoyDiskExtent) == 0)\r
                     {\r
                         Log("DiskSig=%08X PartStart=%lld", DiskSig, VtoyDiskExtent.StartingOffset.QuadPart);\r
@@ -1543,17 +1583,31 @@ static int VentoyHook(ventoy_os_param *param)
             Log("Failed to find ventoy disk");\r
             return 1;\r
         }\r
+\r
+        VtoyDiskNum = VtoyDiskExtent.DiskNumber;\r
     }\r
     else\r
     {\r
         VtoyLetter = Letter;\r
         Log("No vlnk mode %C", Letter);\r
+\r
+        VtoyDiskNum = DiskExtent.DiskNumber;\r
+    }\r
+\r
+    if (CheckVentoyDisk(VtoyDiskNum))\r
+    {\r
+        Log("Disk check OK %C: %u", VtoyLetter, VtoyDiskNum);\r
+    }\r
+    else\r
+    {\r
+        Log("Failed to check ventoy disk %u", VtoyDiskNum);\r
+        return 1;\r
     }\r
 \r
     Drives = GetLogicalDrives();\r
     Log("Drives before mount: 0x%x", Drives);\r
 \r
-    rc = MountIsoFile(IsoPath, DiskExtent.DiskNumber);\r
+    rc = MountIsoFile(IsoPath, VtoyDiskNum);\r
 \r
     NewDrives = GetLogicalDrives();\r
     Log("Drives after mount: 0x%x (0x%x)", NewDrives, (NewDrives ^ Drives));\r
@@ -1588,7 +1642,7 @@ static int VentoyHook(ventoy_os_param *param)
     }\r
 \r
     // for protect\r
-    rc = DeleteVentoyPart2MountPoint(VtoyDiskExtent.DiskNumber);\r
+    rc = DeleteVentoyPart2MountPoint(VtoyDiskNum);\r
     Log("Delete ventoy mountpoint: %s", rc == 0 ? "SUCCESS" : "NO NEED");\r
     \r
     if (g_windows_data.auto_install_script[0])\r
@@ -1615,7 +1669,7 @@ static int VentoyHook(ventoy_os_param *param)
         if (IsFileExist("%s", IsoPath))\r
         {\r
             Log("decompress injection archive %s...", IsoPath);\r
-            DecompressInjectionArchive(IsoPath, VtoyDiskExtent.DiskNumber);\r
+            DecompressInjectionArchive(IsoPath, VtoyDiskNum);\r
 \r
             if (IsFileExist("%s", AUTO_RUN_BAT))\r
             {\r