mirror
/
Ventoy
mirror of https://github.com/ventoy/Ventoy
2
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
Ventoy/vtoycli/partresize.c

684 lines
16 KiB
C
Raw Permalink Blame History

/******************************************************************************
* partresize.c ---- ventoy part resize util
*
* Copyright (c) 2021, longpanda <admin@ventoy.net>
*
* This program is free software; you can redistribute it and/or
* modify it under the terms of the GNU General Public License as
* published by the Free Software Foundation; either version 3 of the
* License, or (at your option) any later version.
*
* This program is distributed in the hope that it will be useful, but
* WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program; if not, see <http://www.gnu.org/licenses/>.
*
*/
#include <stdio.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <dirent.h>
#include <fat_filelib.h>
#include "vtoycli.h"
static int g_disk_fd = 0;
static UINT64 g_disk_offset = 0;
static GUID g_ZeroGuid = {0};
static GUID g_WindowsDataPartGuid = { 0xebd0a0a2, 0xb9e5, 0x4433, { 0x87, 0xc0, 0x68, 0xb6, 0xb7, 0x26, 0x99, 0xc7 } };
static int vtoy_disk_read(uint32 sector, uint8 *buffer, uint32 sector_count)
{
UINT64 offset = sector * 512ULL;
lseek(g_disk_fd, g_disk_offset + offset, SEEK_SET);
read(g_disk_fd, buffer, sector_count * 512);
return 1;
}
static int vtoy_disk_write(uint32 sector, uint8 *buffer, uint32 sector_count)
{
UINT64 offset = sector * 512ULL;
lseek(g_disk_fd, g_disk_offset + offset, SEEK_SET);
write(g_disk_fd, buffer, sector_count * 512);
return 1;
}
static int gpt_check(const char *disk)
{
int fd = -1;
int rc = 1;
VTOY_GPT_INFO *pGPT = NULL;
fd = open(disk, O_RDONLY);
if (fd < 0)
{
printf("Failed to open %s\n", disk);
goto out;
}
pGPT = malloc(sizeof(VTOY_GPT_INFO));
if (NULL == pGPT)
{
goto out;
}
memset(pGPT, 0, sizeof(VTOY_GPT_INFO));
read(fd, pGPT, sizeof(VTOY_GPT_INFO));
if (pGPT->MBR.PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
{
rc = 0;
}
out:
check_close(fd);
check_free(pGPT);
return rc;
}
static int part_check(const char *disk)
{
int i;
int fd = -1;
int rc = 0;
int Index = 0;
int Count = 0;
int PartStyle = 0;
UINT64 Part1Start;
UINT64 Part1End;
UINT64 NextPartStart;
UINT64 DiskSizeInBytes;
VTOY_GPT_INFO *pGPT = NULL;
DiskSizeInBytes = get_disk_size_in_byte(disk);
if (DiskSizeInBytes == 0)
{
printf("Failed to get disk size of %s\n", disk);
goto out;
}
fd = open(disk, O_RDONLY);
if (fd < 0)
{
printf("Failed to open %s\n", disk);
goto out;
}
pGPT = malloc(sizeof(VTOY_GPT_INFO));
if (NULL == pGPT)
{
goto out;
}
memset(pGPT, 0, sizeof(VTOY_GPT_INFO));
read(fd, pGPT, sizeof(VTOY_GPT_INFO));
if (pGPT->MBR.PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
{
PartStyle = 1;
}
else
{
PartStyle = 0;
}
if (PartStyle == 0)
{
PART_TABLE *PartTbl = pGPT->MBR.PartTbl;
for (Count = 0, i = 0; i < 4; i++)
{
if (PartTbl[i].SectorCount > 0)
{
printf("MBR Part%d SectorStart:%u SectorCount:%u\n", i + 1, PartTbl[i].StartSectorId, PartTbl[i].SectorCount);
Count++;
}
}
//We must have a free partition table for VTOYEFI partition
if (Count >= 4)
{
printf("###[FAIL] 4 MBR partition tables are all used.\n");
goto out;
}
if (PartTbl[0].SectorCount > 0)
{
Part1Start = PartTbl[0].StartSectorId;
Part1End = PartTbl[0].SectorCount + Part1Start;
}
else
{
printf("###[FAIL] MBR Partition 1 is invalid\n");
goto out;
}
Index = -1;
NextPartStart = DiskSizeInBytes / 512ULL;
for (i = 1; i < 4; i++)
{
if (PartTbl[i].SectorCount > 0 && NextPartStart > PartTbl[i].StartSectorId)
{
Index = i;
NextPartStart = PartTbl[i].StartSectorId;
}
}
NextPartStart *= 512ULL;
printf("DiskSize:%llu NextPartStart:%llu(LBA:%llu) Index:%d\n",
DiskSizeInBytes, NextPartStart, NextPartStart / 512ULL, Index);
}
else
{
VTOY_GPT_PART_TBL *PartTbl = pGPT->PartTbl;
for (Count = 0, i = 0; i < 128; i++)
{
if (memcmp(&(PartTbl[i].PartGuid), &g_ZeroGuid, sizeof(GUID)))
{
printf("GPT Part%d StartLBA:%llu LastLBA:%llu\n", i + 1, PartTbl[i].StartLBA, PartTbl[i].LastLBA);
Count++;
}
}
if (Count >= 128)
{
printf("###[FAIL] 128 GPT partition tables are all used.\n");
goto out;
}
if (memcmp(&(PartTbl[0].PartGuid), &g_ZeroGuid, sizeof(GUID)))
{
Part1Start = PartTbl[0].StartLBA;
Part1End = PartTbl[0].LastLBA + 1;
}
else
{
printf("###[FAIL] GPT Partition 1 is invalid\n");
goto out;
}
Index = -1;
NextPartStart = (pGPT->Head.PartAreaEndLBA + 1);
for (i = 1; i < 128; i++)
{
if (memcmp(&(PartTbl[i].PartGuid), &g_ZeroGuid, sizeof(GUID)) && NextPartStart > PartTbl[i].StartLBA)
{
Index = i;
NextPartStart = PartTbl[i].StartLBA;
}
}
NextPartStart *= 512ULL;
printf("DiskSize:%llu NextPartStart:%llu(LBA:%llu) Index:%d\n",
DiskSizeInBytes, NextPartStart, NextPartStart / 512ULL, Index);
}
printf("Valid partition table (%s): Valid partition count:%d\n", (PartStyle == 0) ? "MBR" : "GPT", Count);
//Partition 1 MUST start at 1MB
Part1Start *= 512ULL;
Part1End *= 512ULL;
printf("Partition 1 start at: %llu %lluKB, end:%llu, NextPartStart:%llu\n",
Part1Start, Part1Start / 1024, Part1End, NextPartStart);
if (Part1Start != SIZE_1MB)
{
printf("###[FAIL] Partition 1 is not start at 1MB\n");
goto out;
}
//If we have free space after partition 1
if (NextPartStart - Part1End >= VENTOY_EFI_PART_SIZE)
{
printf("Free space after partition 1 (%llu) is enough for VTOYEFI part\n", NextPartStart - Part1End);
rc = 1;
}
else if (NextPartStart == Part1End)
{
printf("There is no free space after partition 1\n");
rc = 2;
}
else
{
printf("The free space after partition 1 is not enough\n");
rc = 2;
}
out:
check_close(fd);
check_free(pGPT);
return rc;
}
static int secureboot_proc(char *disk, UINT64 part2start)
{
int rc = 0;
int size;
int fd = -1;
char *filebuf = NULL;
void *file = NULL;
fd = open(disk, O_RDWR);
if (fd < 0)
{
printf("Failed to open %s\n", disk);
return 1;
}
g_disk_fd = fd;
g_disk_offset = part2start * 512ULL;
fl_init();
if (0 == fl_attach_media(vtoy_disk_read, vtoy_disk_write))
{
file = fl_fopen("/EFI/BOOT/grubx64_real.efi", "rb");
printf("Open ventoy efi file %p\n", file);
if (file)
{
fl_fseek(file, 0, SEEK_END);
size = (int)fl_ftell(file);
fl_fseek(file, 0, SEEK_SET);
printf("ventoy x64 efi file size %d ...\n", size);
filebuf = (char *)malloc(size);
if (filebuf)
{
fl_fread(filebuf, 1, size, file);
}
fl_fclose(file);
fl_remove("/EFI/BOOT/BOOTX64.EFI");
fl_remove("/EFI/BOOT/grubx64.efi");
fl_remove("/EFI/BOOT/grubx64_real.efi");
fl_remove("/EFI/BOOT/MokManager.efi");
fl_remove("/EFI/BOOT/mmx64.efi");
fl_remove("/ENROLL_THIS_KEY_IN_MOKMANAGER.cer");
file = fl_fopen("/EFI/BOOT/BOOTX64.EFI", "wb");
printf("Open bootx64 efi file %p\n", file);
if (file)
{
if (filebuf)
{
fl_fwrite(filebuf, 1, size, file);
}
fl_fflush(file);
fl_fclose(file);
}
if (filebuf)
{
free(filebuf);
}
}
file = fl_fopen("/EFI/BOOT/grubia32_real.efi", "rb");
printf("Open ventoy ia32 efi file %p\n", file);
if (file)
{
fl_fseek(file, 0, SEEK_END);
size = (int)fl_ftell(file);
fl_fseek(file, 0, SEEK_SET);
printf("ventoy efi file size %d ...\n", size);
filebuf = (char *)malloc(size);
if (filebuf)
{
fl_fread(filebuf, 1, size, file);
}
fl_fclose(file);
fl_remove("/EFI/BOOT/BOOTIA32.EFI");
fl_remove("/EFI/BOOT/grubia32.efi");
fl_remove("/EFI/BOOT/grubia32_real.efi");
fl_remove("/EFI/BOOT/mmia32.efi");
file = fl_fopen("/EFI/BOOT/BOOTIA32.EFI", "wb");
printf("Open bootia32 efi file %p\n", file);
if (file)
{
if (filebuf)
{
fl_fwrite(filebuf, 1, size, file);
}
fl_fflush(file);
fl_fclose(file);
}
if (filebuf)
{
free(filebuf);
}
}
}
else
{
rc = 1;
}
fl_shutdown();
fsync(fd);
return rc;
}
static int VentoyFillMBRLocation(UINT64 DiskSizeInBytes, UINT32 StartSectorId, UINT32 SectorCount, PART_TABLE *Table)
{
UINT8 Head;
UINT8 Sector;
UINT8 nSector = 63;
UINT8 nHead = 8;
UINT32 Cylinder;
UINT32 EndSectorId;
while (nHead != 0 && (DiskSizeInBytes / 512 / nSector / nHead) > 1024)
{
nHead = (UINT8)nHead * 2;
}
if (nHead == 0)
{
nHead = 255;
}
Cylinder = StartSectorId / nSector / nHead;
Head = StartSectorId / nSector % nHead;
Sector = StartSectorId % nSector + 1;
Table->StartHead = Head;
Table->StartSector = Sector;
Table->StartCylinder = Cylinder;
EndSectorId = StartSectorId + SectorCount - 1;
Cylinder = EndSectorId / nSector / nHead;
Head = EndSectorId / nSector % nHead;
Sector = EndSectorId % nSector + 1;
Table->EndHead = Head;
Table->EndSector = Sector;
Table->EndCylinder = Cylinder;
Table->StartSectorId = StartSectorId;
Table->SectorCount = SectorCount;
return 0;
}
static int WriteDataToPhyDisk(int fd, UINT64 offset, void *buffer, int len)
{
ssize_t wrlen;
off_t newseek;
newseek = lseek(fd, offset, SEEK_SET);
if (newseek != offset)
{
printf("Failed to lseek %llu %lld %d\n", offset, (long long)newseek, errno);
return 0;
}
wrlen = write(fd, buffer, len);
if ((int)wrlen != len)
{
printf("Failed to write %d %d %d\n", len, (int)wrlen, errno);
return 0;
}
return 1;
}
static int VentoyFillBackupGptHead(VTOY_GPT_INFO *pInfo, VTOY_GPT_HDR *pHead)
{
UINT64 LBA;
UINT64 BackupLBA;
memcpy(pHead, &pInfo->Head, sizeof(VTOY_GPT_HDR));
LBA = pHead->EfiStartLBA;
BackupLBA = pHead->EfiBackupLBA;
pHead->EfiStartLBA = BackupLBA;
pHead->EfiBackupLBA = LBA;
pHead->PartTblStartLBA = BackupLBA + 1 - 33;
pHead->Crc = 0;
pHead->Crc = VtoyCrc32(pHead, pHead->Length);
return 0;
}
static int update_part_table(char *disk, UINT64 part2start)
{
int i;
int j;
int fd = -1;
int rc = 1;
int PartStyle = 0;
ssize_t len = 0;
UINT64 DiskSizeInBytes;
VTOY_GPT_INFO *pGPT = NULL;
VTOY_GPT_HDR *pBack = NULL;
DiskSizeInBytes = get_disk_size_in_byte(disk);
if (DiskSizeInBytes == 0)
{
printf("Failed to get disk size of %s\n", disk);
goto out;
}
fd = open(disk, O_RDWR);
if (fd < 0)
{
printf("Failed to open %s\n", disk);
goto out;
}
pGPT = malloc(sizeof(VTOY_GPT_INFO) + sizeof(VTOY_GPT_HDR));
if (NULL == pGPT)
{
goto out;
}
memset(pGPT, 0, sizeof(VTOY_GPT_INFO) + sizeof(VTOY_GPT_HDR));
pBack = (VTOY_GPT_HDR *)(pGPT + 1);
len = read(fd, pGPT, sizeof(VTOY_GPT_INFO));
if (len != (ssize_t)sizeof(VTOY_GPT_INFO))
{
printf("Failed to read partition table %d err:%d\n", (int)len, errno);
goto out;
}
if (pGPT->MBR.PartTbl[0].FsFlag == 0xEE && memcmp(pGPT->Head.Signature, "EFI PART", 8) == 0)
{
PartStyle = 1;
}
else
{
PartStyle = 0;
}
if (PartStyle == 0)
{
PART_TABLE *PartTbl = pGPT->MBR.PartTbl;
for (i = 1; i < 4; i++)
{
if (PartTbl[i].SectorCount == 0)
{
break;
}
}
if (i >= 4)
{
printf("###[FAIL] Can not find a free MBR partition table.\n");
goto out;
}
for (j = i - 1; j > 0; j--)
{
printf("Move MBR partition table %d --> %d\n", j + 1, j + 2);
memcpy(PartTbl + (j + 1), PartTbl + j, sizeof(PART_TABLE));
}
memset(PartTbl + 1, 0, sizeof(PART_TABLE));
VentoyFillMBRLocation(DiskSizeInBytes, (UINT32)part2start, VENTOY_EFI_PART_SIZE / 512, PartTbl + 1);
PartTbl[1].Active = 0x00;
PartTbl[1].FsFlag = 0xEF; // EFI System Partition
PartTbl[0].Active = 0x80; // bootable
PartTbl[0].SectorCount = (UINT32)part2start - 2048;
if (!WriteDataToPhyDisk(fd, 0, &(pGPT->MBR), 512))
{
printf("MBR write MBR failed\n");
goto out;
}
fsync(fd);
printf("MBR update partition table success.\n");
rc = 0;
}
else
{
VTOY_GPT_PART_TBL *PartTbl = pGPT->PartTbl;
for (i = 1; i < 128; i++)
{
if (memcmp(&(PartTbl[i].PartGuid), &g_ZeroGuid, sizeof(GUID)) == 0)
{
break;
}
}
if (i >= 128)
{
printf("###[FAIL] Can not find a free GPT partition table.\n");
goto out;
}
for (j = i - 1; j > 0; j--)
{
printf("Move GPT partition table %d --> %d\n", j + 1, j + 2);
memcpy(PartTbl + (j + 1), PartTbl + j, sizeof(VTOY_GPT_PART_TBL));
}
// to fix windows issue
memset(PartTbl + 1, 0, sizeof(VTOY_GPT_PART_TBL));
memcpy(&(PartTbl[1].PartType), &g_WindowsDataPartGuid, sizeof(GUID));
ventoy_gen_preudo_uuid(&(PartTbl[1].PartGuid));
PartTbl[0].LastLBA = part2start - 1;
PartTbl[1].StartLBA = PartTbl[0].LastLBA + 1;
PartTbl[1].LastLBA = PartTbl[1].StartLBA + VENTOY_EFI_PART_SIZE / 512 - 1;
PartTbl[1].Attr = 0xC000000000000001ULL;
PartTbl[1].Name[0] = 'V';
PartTbl[1].Name[1] = 'T';
PartTbl[1].Name[2] = 'O';
PartTbl[1].Name[3] = 'Y';
PartTbl[1].Name[4] = 'E';
PartTbl[1].Name[5] = 'F';
PartTbl[1].Name[6] = 'I';
PartTbl[1].Name[7] = 0;
//Update CRC
pGPT->Head.PartTblCrc = VtoyCrc32(pGPT->PartTbl, sizeof(pGPT->PartTbl));
pGPT->Head.Crc = 0;
pGPT->Head.Crc = VtoyCrc32(&(pGPT->Head), pGPT->Head.Length);
printf("pGPT->Head.EfiStartLBA=%llu\n", pGPT->Head.EfiStartLBA);
printf("pGPT->Head.EfiBackupLBA=%llu\n", pGPT->Head.EfiBackupLBA);
VentoyFillBackupGptHead(pGPT, pBack);
if (!WriteDataToPhyDisk(fd, pGPT->Head.EfiBackupLBA * 512, pBack, 512))
{
printf("GPT write backup head failed\n");
goto out;
}
if (!WriteDataToPhyDisk(fd, (pGPT->Head.EfiBackupLBA - 32) * 512, pGPT->PartTbl, 512 * 32))
{
printf("GPT write backup partition table failed\n");
goto out;
}
if (!WriteDataToPhyDisk(fd, 0, pGPT, 512 * 34))
{
printf("GPT write MBR & Main partition table failed\n");
goto out;
}
fsync(fd);
printf("GPT update partition table success.\n");
rc = 0;
}
out:
check_close(fd);
check_free(pGPT);
return rc;
}
int partresize_main(int argc, char **argv)
{
UINT64 sector;
if (argc != 3 && argc != 4)
{
printf("usage: partresize -c/-f /dev/sdb\n");
return 1;
}
if (strcmp(argv[1], "-c") == 0)
{
return part_check(argv[2]);
}
else if (strcmp(argv[1], "-s") == 0)
{
sector = strtoull(argv[3], NULL, 10);
return secureboot_proc(argv[2], sector);
}
else if (strcmp(argv[1], "-p") == 0)
{
sector = strtoull(argv[3], NULL, 10);
return update_part_table(argv[2], sector);
}
else if (strcmp(argv[1], "-t") == 0)
{
return gpt_check(argv[2]);
}
else
{
return 1;
}
}
Reference in New Issue View Git Blame Copy Permalink