fix a bug about injection plugin in GPT partition style.

marenan · Jul 24, 2020 · d063dbb · d063dbb
1 parent 89938e3
commit d063dbb
Show file tree

Hide file tree

Showing 5 changed files with 119 additions and 3 deletions.
diff --git a/INSTALL/ventoy/vtoyjump32.exe b/INSTALL/ventoy/vtoyjump32.exe
diff --git a/INSTALL/ventoy/vtoyjump64.exe b/INSTALL/ventoy/vtoyjump64.exe
diff --git a/vtoygpt/vtoygpt.c b/vtoygpt/vtoygpt.c
@@ -244,7 +244,7 @@ int DumpGptInfo(VTOY_GPT_INFO *pGptInfo)
     return 0;
 }
 
-#define VENTOY_EFI_PART_ATTR   0xC000000000000001ULL
+#define VENTOY_EFI_PART_ATTR   0x8000000000000001ULL
 
 int main(int argc, const char **argv)
 {

diff --git a/vtoyjump/vtoyjump/vtoyjump.c b/vtoyjump/vtoyjump/vtoyjump.c
@@ -513,6 +513,55 @@ static CHAR GetMountLogicalDrive(void)
 	return Letter;
 }
 
+UINT64 GetVentoyEfiPartStartSector(HANDLE hDrive)
+{
+	BOOL bRet;
+	DWORD dwSize; 
+	MBR_HEAD MBR;	
+	VTOY_GPT_INFO *pGpt = NULL;
+	UINT64 StartSector = 0;
+
+	SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
+
+	bRet = ReadFile(hDrive, &MBR, sizeof(MBR), &dwSize, NULL);
+	Log("Read MBR Ret:%u Size:%u code:%u", bRet, dwSize, LASTERR);
+
+	if ((!bRet) || (dwSize != sizeof(MBR)))
+	{
+		0;
+	}
+
+	if (MBR.PartTbl[0].FsFlag == 0xEE)
+	{
+		Log("GPT partition style");
+
+		pGpt = malloc(sizeof(VTOY_GPT_INFO));
+		if (!pGpt)
+		{
+			return 0;
+		}
+
+		SetFilePointer(hDrive, 0, NULL, FILE_BEGIN);
+		bRet = ReadFile(hDrive, pGpt, sizeof(VTOY_GPT_INFO), &dwSize, NULL);		
+		if ((!bRet) || (dwSize != sizeof(VTOY_GPT_INFO)))
+		{
+			Log("Failed to read gpt info %d %u %d", bRet, dwSize, LASTERR);
+			return 0;
+		}
+
+		StartSector = pGpt->PartTbl[1].StartLBA;
+		free(pGpt);
+	}
+	else
+	{
+		Log("MBR partition style");
+		StartSector = MBR.PartTbl[1].StartSectorId;
+	}
+
+	Log("GetVentoyEfiPart StartSector: %llu", StartSector);
+	return StartSector;
+}
+
 int VentoyMountISOByImdisk(const char *IsoPath, DWORD PhyDrive)
 {
 	int rc = 1;
@@ -543,7 +592,7 @@ int VentoyMountISOByImdisk(const char *IsoPath, DWORD PhyDrive)
 	}
 
 	g_FatPhyDrive = hDrive;
-	g_Part2StartSec = (LengthInfo.Length.QuadPart - VENTOY_EFI_PART_SIZE) / 512;
+	g_Part2StartSec = GetVentoyEfiPartStartSector(hDrive);
 
 	Log("Parse FAT fs...");
 
@@ -778,7 +827,7 @@ static int DecompressInjectionArchive(const char *archive, DWORD PhyDrive)
     }
 
     g_FatPhyDrive = hDrive;
-    g_Part2StartSec = (LengthInfo.Length.QuadPart - VENTOY_EFI_PART_SIZE) / 512;
+	g_Part2StartSec = GetVentoyEfiPartStartSector(hDrive);
 
     Log("Parse FAT fs...");
 

diff --git a/vtoyjump/vtoyjump/vtoyjump.h b/vtoyjump/vtoyjump/vtoyjump.h
@@ -72,6 +72,72 @@ typedef struct ventoy_windows_data
     UINT8 reserved[256];
 }ventoy_windows_data;
 
+
+
+typedef struct PART_TABLE
+{
+	UINT8  Active; // 0x00  0x80
+
+	UINT8  StartHead;
+	UINT16 StartSector : 6;
+	UINT16 StartCylinder : 10;
+
+	UINT8  FsFlag;
+
+	UINT8  EndHead;
+	UINT16 EndSector : 6;
+	UINT16 EndCylinder : 10;
+
+	UINT32 StartSectorId;
+	UINT32 SectorCount;
+}PART_TABLE;
+
+typedef struct MBR_HEAD
+{
+	UINT8 BootCode[446];
+	PART_TABLE PartTbl[4];
+	UINT8 Byte55;
+	UINT8 ByteAA;
+}MBR_HEAD;
+
+typedef struct VTOY_GPT_HDR
+{
+	CHAR   Signature[8]; /* EFI PART */
+	UINT8  Version[4];
+	UINT32 Length;
+	UINT32 Crc;
+	UINT8  Reserved1[4];
+	UINT64 EfiStartLBA;
+	UINT64 EfiBackupLBA;
+	UINT64 PartAreaStartLBA;
+	UINT64 PartAreaEndLBA;
+	GUID   DiskGuid;
+	UINT64 PartTblStartLBA;
+	UINT32 PartTblTotNum;
+	UINT32 PartTblEntryLen;
+	UINT32 PartTblCrc;
+	UINT8  Reserved2[420];
+}VTOY_GPT_HDR;
+
+typedef struct VTOY_GPT_PART_TBL
+{
+	GUID   PartType;
+	GUID   PartGuid;
+	UINT64 StartLBA;
+	UINT64 LastLBA;
+	UINT64 Attr;
+	UINT16 Name[36];
+}VTOY_GPT_PART_TBL;
+
+typedef struct VTOY_GPT_INFO
+{
+	MBR_HEAD MBR;
+	VTOY_GPT_HDR Head;
+	VTOY_GPT_PART_TBL PartTbl[128];
+}VTOY_GPT_INFO;
+
+
+
 #pragma pack()
 
 
@@ -84,5 +150,6 @@ typedef struct ventoy_windows_data
 	}\
 }
 
+#define LASTERR     GetLastError()
 
 #endif