EDK2/edk2_mod/edk2-edk2-stable201911/MdeModulePkg/Application/Ventoy/VentoyProtocol.c

   1 /******************************************************************************
   2  * Ventoy.c
   3  *
   4  * Copyright (c) 2020, 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 <Uefi.h>
  22 #include <Library/DebugLib.h>
  23 #include <Library/PrintLib.h>
  24 #include <Library/UefiLib.h>
  25 #include <Library/BaseMemoryLib.h>
  26 #include <Library/DevicePathLib.h>
  27 #include <Library/MemoryAllocationLib.h>
  28 #include <Library/UefiBootServicesTableLib.h>
  29 #include <Library/UefiRuntimeServicesTableLib.h>
  30 #include <Library/UefiApplicationEntryPoint.h>
  31 #include <Protocol/LoadedImage.h>
  32 #include <Guid/FileInfo.h>
  33 #include <Guid/FileSystemInfo.h>
  34 #include <Protocol/BlockIo.h>
  35 #include <Protocol/RamDisk.h>
  36 #include <Protocol/SimpleFileSystem.h>
  37 #include <Ventoy.h>
  38
  39 UINTN g_iso_buf_size = 0;
  40 BOOLEAN gMemdiskMode = FALSE;
  41
  42 ventoy_sector_flag *g_sector_flag = NULL;
  43 UINT32 g_sector_flag_num = 0;
  44
  45 EFI_FILE_OPEN g_original_fopen = NULL;
  46 EFI_FILE_CLOSE g_original_fclose = NULL;
  47 EFI_SIMPLE_FILE_SYSTEM_PROTOCOL_OPEN_VOLUME g_original_open_volume = NULL;
  48
  49 /* EFI block device vendor device path GUID */
  50 EFI_GUID gVtoyBlockDevicePathGuid = VTOY_BLOCK_DEVICE_PATH_GUID;
  51
  52 #define VENTOY_ISO9660_SECTOR_OVERFLOW  2097152
  53
  54 BOOLEAN g_fixup_iso9660_secover_enable = FALSE;
  55 BOOLEAN g_fixup_iso9660_secover_start  = FALSE;
  56 UINT64  g_fixup_iso9660_secover_1st_secs = 0;
  57 UINT64  g_fixup_iso9660_secover_cur_secs = 0;
  58 UINT64  g_fixup_iso9660_secover_tot_secs = 0;
  59
  60 STATIC UINTN g_keyboard_hook_count = 0;
  61 STATIC BOOLEAN g_blockio_start_record_bcd = FALSE;
  62 STATIC BOOLEAN g_blockio_bcd_read_done = FALSE;
  63
  64 EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *g_con_simple_input_ex = NULL;
  65 STATIC EFI_INPUT_READ_KEY_EX g_org_read_key_ex = NULL;
  66 STATIC EFI_INPUT_READ_KEY g_org_read_key = NULL;
  67
  68 STATIC EFI_LOCATE_HANDLE g_org_locate_handle = NULL;
  69
  70 BOOLEAN ventoy_is_cdrom_dp_exist(VOID)
  71 {
  72     UINTN i = 0;
  73     UINTN Count = 0;
  74     EFI_HANDLE *Handles = NULL;
  75     EFI_STATUS Status = EFI_SUCCESS;
  76     EFI_DEVICE_PATH_PROTOCOL *DevicePath = NULL;
  77
  78     Status = gBS->LocateHandleBuffer(ByProtocol, &gEfiDevicePathProtocolGuid,
  79                                      NULL, &Count, &Handles);
  80     if (EFI_ERROR(Status))
  81     {
  82         return FALSE;
  83     }
  84
  85     for (i = 0; i < Count; i++)
  86     {
  87         Status = gBS->HandleProtocol(Handles[i], &gEfiDevicePathProtocolGuid, (VOID **)&DevicePath);
  88         if (EFI_ERROR(Status))
  89         {
  90             continue;
  91         }
  92
  93         while (!IsDevicePathEnd(DevicePath))
  94         {
  95             if (MEDIA_DEVICE_PATH == DevicePath->Type && MEDIA_CDROM_DP == DevicePath->SubType)
  96             {
  97                 FreePool(Handles);
  98                 return TRUE;
  99             }
 100
 101             DevicePath = NextDevicePathNode(DevicePath);
 102         }
 103     }
 104
 105     FreePool(Handles);
 106     return FALSE;
 107 }
 108
 109 #if 0
 110 /* Block IO procotol */
 111 #endif
 112
 113 EFI_STATUS EFIAPI ventoy_block_io_reset
 114 (
 115     IN EFI_BLOCK_IO_PROTOCOL          *This,
 116     IN BOOLEAN                        ExtendedVerification
 117 )
 118 {
 119     (VOID)This;
 120     (VOID)ExtendedVerification;
 121         return EFI_SUCCESS;
 122 }
 123
 124 STATIC EFI_STATUS EFIAPI ventoy_read_iso_sector
 125 (
 126     IN UINT64                 Sector,
 127     IN UINTN                  Count,
 128     OUT VOID                 *Buffer
 129 )
 130 {
 131     EFI_STATUS Status = EFI_SUCCESS;
 132     EFI_LBA MapLba = 0;
 133     UINT32 i = 0;
 134     UINTN secLeft = 0;
 135     UINTN secRead = 0;
 136     UINT64 ReadStart = 0;
 137     UINT64 ReadEnd = 0;
 138     UINT64 OverrideStart = 0;
 139     UINT64 OverrideEnd= 0;
 140     UINT8 *pCurBuf = (UINT8 *)Buffer;
 141     ventoy_img_chunk *pchunk = g_chunk;
 142     ventoy_override_chunk *pOverride = g_override_chunk;
 143     EFI_BLOCK_IO_PROTOCOL *pRawBlockIo = gBlockData.pRawBlockIo;
 144
 145     debug("read iso sector %lu  count %u", Sector, Count);
 146
 147     ReadStart = Sector * 2048;
 148     ReadEnd = (Sector + Count) * 2048;
 149
 150     for (i = 0; Count > 0 && i < g_img_chunk_num; i++, pchunk++)
 151     {
 152         if (Sector >= pchunk->img_start_sector && Sector <= pchunk->img_end_sector)
 153         {
 154             if (g_chain->disk_sector_size == 512)
 155             {
 156                 MapLba = (Sector - pchunk->img_start_sector) * 4 + pchunk->disk_start_sector;
 157             }
 158             else
 159             {
 160                 MapLba = (Sector - pchunk->img_start_sector) * 2048 / g_chain->disk_sector_size + pchunk->disk_start_sector;
 161             }
 162
 163             secLeft = pchunk->img_end_sector + 1 - Sector;
 164             secRead = (Count < secLeft) ? Count : secLeft;
 165
 166             Status = pRawBlockIo->ReadBlocks(pRawBlockIo, pRawBlockIo->Media->MediaId,
 167                                      MapLba, secRead * 2048, pCurBuf);
 168             if (EFI_ERROR(Status))
 169             {
 170                 debug("Raw disk read block failed %r LBA:%lu Count:%u", Status, MapLba, secRead);
 171                 return Status;
 172             }
 173
 174             Count -= secRead;
 175             Sector += secRead;
 176             pCurBuf += secRead * 2048;
 177         }
 178     }
 179
 180     if (ReadStart > g_chain->real_img_size_in_bytes)
 181     {
 182         return EFI_SUCCESS;
 183     }
 184
 185     /* override data */
 186     pCurBuf = (UINT8 *)Buffer;
 187     for (i = 0; i < g_override_chunk_num; i++, pOverride++)
 188     {
 189         OverrideStart = pOverride->img_offset;
 190         OverrideEnd = pOverride->img_offset + pOverride->override_size;
 191
 192         if (OverrideStart >= ReadEnd || ReadStart >= OverrideEnd)
 193         {
 194             continue;
 195         }
 196
 197         if (ReadStart <= OverrideStart)
 198         {
 199             if (ReadEnd <= OverrideEnd)
 200             {
 201                 CopyMem(pCurBuf + OverrideStart - ReadStart, pOverride->override_data, ReadEnd - OverrideStart);
 202             }
 203             else
 204             {
 205                 CopyMem(pCurBuf + OverrideStart - ReadStart, pOverride->override_data, pOverride->override_size);
 206             }
 207         }
 208         else
 209         {
 210             if (ReadEnd <= OverrideEnd)
 211             {
 212                 CopyMem(pCurBuf, pOverride->override_data + ReadStart - OverrideStart, ReadEnd - ReadStart);
 213             }
 214             else
 215             {
 216                 CopyMem(pCurBuf, pOverride->override_data + ReadStart - OverrideStart, OverrideEnd - ReadStart);
 217             }
 218         }
 219
 220         if (g_fixup_iso9660_secover_enable && (!g_fixup_iso9660_secover_start) &&
 221             pOverride->override_size == sizeof(ventoy_iso9660_override))
 222         {
 223             ventoy_iso9660_override *dirent = (ventoy_iso9660_override *)pOverride->override_data;
 224             if (dirent->first_sector >= VENTOY_ISO9660_SECTOR_OVERFLOW)
 225             {
 226                 g_fixup_iso9660_secover_start = TRUE;
 227                 g_fixup_iso9660_secover_cur_secs = 0;
 228             }
 229         }
 230     }
 231
 232     if (g_blockio_start_record_bcd && FALSE == g_blockio_bcd_read_done)
 233     {
 234         if (*(UINT32 *)Buffer == 0x66676572)
 235         {
 236             g_blockio_bcd_read_done = TRUE;
 237         }
 238     }
 239
 240     return EFI_SUCCESS;
 241 }
 242
 243 EFI_STATUS EFIAPI ventoy_block_io_ramdisk_read
 244 (
 245     IN EFI_BLOCK_IO_PROTOCOL          *This,
 246     IN UINT32                          MediaId,
 247     IN EFI_LBA                         Lba,
 248     IN UINTN                           BufferSize,
 249     OUT VOID                          *Buffer
 250 )
 251 {
 252     //debug("### ventoy_block_io_ramdisk_read sector:%u count:%u", (UINT32)Lba, (UINT32)BufferSize / 2048);
 253
 254     (VOID)This;
 255     (VOID)MediaId;
 256
 257     CopyMem(Buffer, (char *)g_chain + (Lba * 2048), BufferSize);
 258
 259     if (g_blockio_start_record_bcd && FALSE == g_blockio_bcd_read_done)
 260     {
 261         if (*(UINT32 *)Buffer == 0x66676572)
 262         {
 263             g_blockio_bcd_read_done = TRUE;
 264         }
 265     }
 266
 267         return EFI_SUCCESS;
 268 }
 269
 270 EFI_LBA EFIAPI ventoy_fixup_iso9660_sector(IN EFI_LBA Lba, UINT32 secNum)
 271 {
 272     UINT32 i = 0;
 273
 274     if (g_fixup_iso9660_secover_cur_secs > 0)
 275     {
 276         Lba += VENTOY_ISO9660_SECTOR_OVERFLOW;
 277         g_fixup_iso9660_secover_cur_secs += secNum;
 278         if (g_fixup_iso9660_secover_cur_secs >= g_fixup_iso9660_secover_tot_secs)
 279         {
 280             g_fixup_iso9660_secover_start = FALSE;
 281             goto end;
 282         }
 283     }
 284     else
 285     {
 286         ventoy_iso9660_override *dirent;
 287         ventoy_override_chunk *pOverride;
 288
 289         for (i = 0, pOverride = g_override_chunk; i < g_override_chunk_num; i++, pOverride++)
 290         {
 291             dirent = (ventoy_iso9660_override *)pOverride->override_data;
 292             if (Lba == dirent->first_sector)
 293             {
 294                 g_fixup_iso9660_secover_start = FALSE;
 295                 goto end;
 296             }
 297         }
 298
 299         if (g_fixup_iso9660_secover_start)
 300         {
 301             for (i = 0, pOverride = g_override_chunk; i < g_override_chunk_num; i++, pOverride++)
 302             {
 303                 dirent = (ventoy_iso9660_override *)pOverride->override_data;
 304                 if (Lba + VENTOY_ISO9660_SECTOR_OVERFLOW == dirent->first_sector)
 305                 {
 306                     g_fixup_iso9660_secover_tot_secs = (dirent->size + 2047) / 2048;
 307                     g_fixup_iso9660_secover_cur_secs = secNum;
 308                     if (g_fixup_iso9660_secover_cur_secs >= g_fixup_iso9660_secover_tot_secs)
 309                     {
 310                         g_fixup_iso9660_secover_start = FALSE;
 311                     }
 312                     Lba += VENTOY_ISO9660_SECTOR_OVERFLOW;
 313                     goto end;
 314                 }
 315             }
 316         }
 317     }
 318
 319 end:
 320     return Lba;
 321 }
 322
 323 EFI_STATUS EFIAPI ventoy_block_io_read
 324 (
 325     IN EFI_BLOCK_IO_PROTOCOL          *This,
 326     IN UINT32                          MediaId,
 327     IN EFI_LBA                         Lba,
 328     IN UINTN                           BufferSize,
 329     OUT VOID                          *Buffer
 330 )
 331 {
 332     UINT32 i = 0;
 333     UINT32 j = 0;
 334     UINT32 lbacount = 0;
 335     UINT32 secNum = 0;
 336     UINT64 offset = 0;
 337     EFI_LBA curlba = 0;
 338     EFI_LBA lastlba = 0;
 339     UINT8 *lastbuffer;
 340     ventoy_sector_flag *cur_flag;
 341     ventoy_virt_chunk *node;
 342
 343     //debug("### ventoy_block_io_read sector:%u count:%u", (UINT32)Lba, (UINT32)BufferSize / 2048);
 344
 345     secNum = BufferSize / 2048;
 346
 347     /* Workaround for SSTR PE loader error */
 348     if (g_fixup_iso9660_secover_start)
 349     {
 350         Lba = ventoy_fixup_iso9660_sector(Lba, secNum);
 351     }
 352
 353     offset = Lba * 2048;
 354
 355     if (offset + BufferSize <= g_chain->real_img_size_in_bytes)
 356     {
 357         return ventoy_read_iso_sector(Lba, secNum, Buffer);
 358     }
 359
 360     if (secNum > g_sector_flag_num)
 361     {
 362         cur_flag = AllocatePool(secNum * sizeof(ventoy_sector_flag));
 363         if (NULL == cur_flag)
 364         {
 365             return EFI_OUT_OF_RESOURCES;
 366         }
 367
 368         FreePool(g_sector_flag);
 369         g_sector_flag = cur_flag;
 370         g_sector_flag_num = secNum;
 371     }
 372
 373     for (curlba = Lba, cur_flag = g_sector_flag, j = 0; j < secNum; j++, curlba++, cur_flag++)
 374     {
 375         cur_flag->flag = 0;
 376         for (node = g_virt_chunk, i = 0; i < g_virt_chunk_num; i++, node++)
 377         {
 378             if (curlba >= node->mem_sector_start && curlba < node->mem_sector_end)
 379             {
 380                 CopyMem((UINT8 *)Buffer + j * 2048,
 381                        (char *)g_virt_chunk + node->mem_sector_offset + (curlba - node->mem_sector_start) * 2048,
 382                        2048);
 383                 cur_flag->flag = 1;
 384                 break;
 385             }
 386             else if (curlba >= node->remap_sector_start && curlba < node->remap_sector_end)
 387             {
 388                 cur_flag->remap_lba = node->org_sector_start + curlba - node->remap_sector_start;
 389                 cur_flag->flag = 2;
 390                 break;
 391             }
 392         }
 393     }
 394
 395     for (curlba = Lba, cur_flag = g_sector_flag, j = 0; j < secNum; j++, curlba++, cur_flag++)
 396     {
 397         if (cur_flag->flag == 2)
 398         {
 399             if (lastlba == 0)
 400             {
 401                 lastbuffer = (UINT8 *)Buffer + j * 2048;
 402                 lastlba = cur_flag->remap_lba;
 403                 lbacount = 1;
 404             }
 405             else if (lastlba + lbacount == cur_flag->remap_lba)
 406             {
 407                 lbacount++;
 408             }
 409             else
 410             {
 411                 ventoy_read_iso_sector(lastlba, lbacount, lastbuffer);
 412                 lastbuffer = (UINT8 *)Buffer + j * 2048;
 413                 lastlba = cur_flag->remap_lba;
 414                 lbacount = 1;
 415             }
 416         }
 417     }
 418
 419     if (lbacount > 0)
 420     {
 421         ventoy_read_iso_sector(lastlba, lbacount, lastbuffer);
 422     }
 423
 424         return EFI_SUCCESS;
 425 }
 426
 427 EFI_STATUS EFIAPI ventoy_block_io_write
 428 (
 429     IN EFI_BLOCK_IO_PROTOCOL          *This,
 430     IN UINT32                          MediaId,
 431     IN EFI_LBA                         Lba,
 432     IN UINTN                           BufferSize,
 433     IN VOID                           *Buffer
 434 )
 435 {
 436     (VOID)This;
 437     (VOID)MediaId;
 438     (VOID)Lba;
 439     (VOID)BufferSize;
 440     (VOID)Buffer;
 441         return EFI_WRITE_PROTECTED;
 442 }
 443
 444 EFI_STATUS EFIAPI ventoy_block_io_flush(IN EFI_BLOCK_IO_PROTOCOL *This)
 445 {
 446         (VOID)This;
 447         return EFI_SUCCESS;
 448 }
 449
 450
 451 EFI_STATUS EFIAPI ventoy_fill_device_path(VOID)
 452 {
 453     UINTN NameLen = 0;
 454     UINT8 TmpBuf[128] = {0};
 455     VENDOR_DEVICE_PATH *venPath = NULL;
 456
 457     venPath = (VENDOR_DEVICE_PATH *)TmpBuf;
 458     NameLen = StrSize(VTOY_BLOCK_DEVICE_PATH_NAME);
 459     venPath->Header.Type = HARDWARE_DEVICE_PATH;
 460     venPath->Header.SubType = HW_VENDOR_DP;
 461     venPath->Header.Length[0] = sizeof(VENDOR_DEVICE_PATH) + NameLen;
 462     venPath->Header.Length[1] = 0;
 463     CopyMem(&venPath->Guid, &gVtoyBlockDevicePathGuid, sizeof(EFI_GUID));
 464     CopyMem(venPath + 1, VTOY_BLOCK_DEVICE_PATH_NAME, NameLen);
 465
 466     gBlockData.Path = AppendDevicePathNode(NULL, (EFI_DEVICE_PATH_PROTOCOL *)TmpBuf);
 467     gBlockData.DevicePathCompareLen = sizeof(VENDOR_DEVICE_PATH) + NameLen;
 468
 469     debug("gBlockData.Path=<%s>\n", ConvertDevicePathToText(gBlockData.Path, FALSE, FALSE));
 470
 471     return EFI_SUCCESS;
 472 }
 473
 474 EFI_STATUS EFIAPI ventoy_connect_driver(IN EFI_HANDLE ControllerHandle, IN CONST CHAR16 *DrvName)
 475 {
 476     UINTN i = 0;
 477     UINTN Count = 0;
 478     CHAR16 *DriverName = NULL;
 479     EFI_HANDLE *Handles = NULL;
 480     EFI_HANDLE DrvHandles[2] = { NULL };
 481     EFI_STATUS Status = EFI_SUCCESS;
 482     EFI_COMPONENT_NAME_PROTOCOL *NameProtocol = NULL;
 483     EFI_COMPONENT_NAME2_PROTOCOL *Name2Protocol = NULL;
 484
 485     debug("ventoy_connect_driver <%s>...", DrvName);
 486
 487     Status = gBS->LocateHandleBuffer(ByProtocol, &gEfiComponentName2ProtocolGuid,
 488                                      NULL, &Count, &Handles);
 489     if (EFI_ERROR(Status))
 490     {
 491         return Status;
 492     }
 493
 494     for (i = 0; i < Count; i++)
 495     {
 496         Status = gBS->HandleProtocol(Handles[i], &gEfiComponentName2ProtocolGuid, (VOID **)&Name2Protocol);
 497         if (EFI_ERROR(Status))
 498         {
 499             continue;
 500         }
 501
 502         Status = Name2Protocol->GetDriverName(Name2Protocol, "en", &DriverName);
 503         if (EFI_ERROR(Status) || NULL == DriverName)
 504         {
 505             continue;
 506         }
 507
 508         if (StrStr(DriverName, DrvName))
 509         {
 510             debug("Find driver name2:<%s>: <%s>", DriverName, DrvName);
 511             DrvHandles[0] = Handles[i];
 512             break;
 513         }
 514     }
 515
 516     if (i < Count)
 517     {
 518         Status = gBS->ConnectController(ControllerHandle, DrvHandles, NULL, TRUE);
 519         debug("ventoy_connect_driver:<%s> <%r>", DrvName, Status);
 520         goto end;
 521     }
 522
 523     debug("%s NOT found, now try COMPONENT_NAME", DrvName);
 524
 525     Count = 0;
 526     FreePool(Handles);
 527     Handles = NULL;
 528
 529     Status = gBS->LocateHandleBuffer(ByProtocol, &gEfiComponentNameProtocolGuid,
 530                                      NULL, &Count, &Handles);
 531     if (EFI_ERROR(Status))
 532     {
 533         return Status;
 534     }
 535
 536     for (i = 0; i < Count; i++)
 537     {
 538         Status = gBS->HandleProtocol(Handles[i], &gEfiComponentNameProtocolGuid, (VOID **)&NameProtocol);
 539         if (EFI_ERROR(Status))
 540         {
 541             continue;
 542         }
 543
 544         Status = NameProtocol->GetDriverName(NameProtocol, "en", &DriverName);
 545         if (EFI_ERROR(Status))
 546         {
 547             continue;
 548         }
 549
 550         if (StrStr(DriverName, DrvName))
 551         {
 552             debug("Find driver name:<%s>: <%s>", DriverName, DrvName);
 553             DrvHandles[0] = Handles[i];
 554             break;
 555         }
 556     }
 557
 558     if (i < Count)
 559     {
 560         Status = gBS->ConnectController(ControllerHandle, DrvHandles, NULL, TRUE);
 561         debug("ventoy_connect_driver:<%s> <%r>", DrvName, Status);
 562         goto end;
 563     }
 564
 565     Status = EFI_NOT_FOUND;
 566
 567 end:
 568     FreePool(Handles);
 569
 570     return Status;
 571 }
 572
 573 EFI_STATUS EFIAPI ventoy_install_blockio(IN EFI_HANDLE ImageHandle, IN UINT64 ImgSize)
 574 {
 575     EFI_STATUS Status = EFI_SUCCESS;
 576     EFI_BLOCK_IO_PROTOCOL *pBlockIo = &(gBlockData.BlockIo);
 577
 578     ventoy_fill_device_path();
 579
 580     debug("install block io protocol %p", ImageHandle);
 581     ventoy_debug_pause();
 582
 583     gBlockData.Media.BlockSize = 2048;
 584     gBlockData.Media.LastBlock = ImgSize / 2048 - 1;
 585     gBlockData.Media.ReadOnly = TRUE;
 586     gBlockData.Media.MediaPresent = 1;
 587     gBlockData.Media.LogicalBlocksPerPhysicalBlock = 1;
 588
 589         pBlockIo->Revision = EFI_BLOCK_IO_PROTOCOL_REVISION3;
 590         pBlockIo->Media = &(gBlockData.Media);
 591         pBlockIo->Reset = ventoy_block_io_reset;
 592     pBlockIo->ReadBlocks = gMemdiskMode ? ventoy_block_io_ramdisk_read : ventoy_block_io_read;
 593         pBlockIo->WriteBlocks = ventoy_block_io_write;
 594         pBlockIo->FlushBlocks = ventoy_block_io_flush;
 595
 596     Status = gBS->InstallMultipleProtocolInterfaces(&gBlockData.Handle,
 597             &gEfiBlockIoProtocolGuid, &gBlockData.BlockIo,
 598             &gEfiDevicePathProtocolGuid, gBlockData.Path,
 599             NULL);
 600     debug("Install protocol %r %p", Status, gBlockData.Handle);
 601     if (EFI_ERROR(Status))
 602     {
 603         return Status;
 604     }
 605
 606     Status = ventoy_connect_driver(gBlockData.Handle, L"Disk I/O Driver");
 607     debug("Connect disk IO driver %r", Status);
 608
 609     Status = ventoy_connect_driver(gBlockData.Handle, L"Partition Driver");
 610     debug("Connect partition driver %r", Status);
 611     if (EFI_ERROR(Status))
 612     {
 613         Status = gBS->ConnectController(gBlockData.Handle, NULL, NULL, TRUE);
 614         debug("Connect all controller %r", Status);
 615     }
 616
 617     ventoy_debug_pause();
 618
 619     return EFI_SUCCESS;
 620 }
 621
 622 #if 0
 623 /* For file replace */
 624 #endif
 625
 626 STATIC EFI_STATUS EFIAPI
 627 ventoy_wrapper_fs_open(EFI_FILE_HANDLE This, EFI_FILE_HANDLE *New, CHAR16 *Name, UINT64 Mode, UINT64 Attributes)
 628 {
 629     (VOID)This;
 630     (VOID)New;
 631     (VOID)Name;
 632     (VOID)Mode;
 633     (VOID)Attributes;
 634     return EFI_SUCCESS;
 635 }
 636
 637 STATIC EFI_STATUS EFIAPI
 638 ventoy_wrapper_file_open_ex(EFI_FILE_HANDLE This, EFI_FILE_HANDLE *New, CHAR16 *Name, UINT64 Mode, UINT64 Attributes, EFI_FILE_IO_TOKEN *Token)
 639 {
 640         return ventoy_wrapper_fs_open(This, New, Name, Mode, Attributes);
 641 }
 642
 643 STATIC EFI_STATUS EFIAPI
 644 ventoy_wrapper_file_delete(EFI_FILE_HANDLE This)
 645 {
 646     (VOID)This;
 647         return EFI_SUCCESS;
 648 }
 649
 650 STATIC EFI_STATUS EFIAPI
 651 ventoy_wrapper_file_set_info(EFI_FILE_HANDLE This, EFI_GUID *Type, UINTN Len, VOID *Data)
 652 {
 653         return EFI_SUCCESS;
 654 }
 655
 656 STATIC EFI_STATUS EFIAPI
 657 ventoy_wrapper_file_flush(EFI_FILE_HANDLE This)
 658 {
 659     (VOID)This;
 660         return EFI_SUCCESS;
 661 }
 662
 663 /* Ex version */
 664 STATIC EFI_STATUS EFIAPI
 665 ventoy_wrapper_file_flush_ex(EFI_FILE_HANDLE This, EFI_FILE_IO_TOKEN *Token)
 666 {
 667     (VOID)This;
 668     (VOID)Token;
 669         return EFI_SUCCESS;
 670 }
 671
 672
 673 STATIC EFI_STATUS EFIAPI
 674 ventoy_wrapper_file_write(EFI_FILE_HANDLE This, UINTN *Len, VOID *Data)
 675 {
 676     (VOID)This;
 677     (VOID)Len;
 678     (VOID)Data;
 679
 680         return EFI_WRITE_PROTECTED;
 681 }
 682
 683 STATIC EFI_STATUS EFIAPI
 684 ventoy_wrapper_file_write_ex(IN EFI_FILE_PROTOCOL *This, IN OUT EFI_FILE_IO_TOKEN *Token)
 685 {
 686         return ventoy_wrapper_file_write(This, &(Token->BufferSize), Token->Buffer);
 687 }
 688
 689
 690 STATIC EFI_STATUS EFIAPI
 691 ventoy_wrapper_file_close(EFI_FILE_HANDLE This)
 692 {
 693     (VOID)This;
 694     return EFI_SUCCESS;
 695 }
 696
 697
 698 STATIC EFI_STATUS EFIAPI
 699 ventoy_wrapper_file_set_pos(EFI_FILE_HANDLE This, UINT64 Position)
 700 {
 701     (VOID)This;
 702
 703     if (Position <= g_efi_file_replace.FileSizeBytes)
 704     {
 705         g_efi_file_replace.CurPos = Position;
 706     }
 707     else
 708     {
 709         g_efi_file_replace.CurPos = g_efi_file_replace.FileSizeBytes;
 710     }
 711
 712     return EFI_SUCCESS;
 713 }
 714
 715 STATIC EFI_STATUS EFIAPI
 716 ventoy_wrapper_file_get_pos(EFI_FILE_HANDLE This, UINT64 *Position)
 717 {
 718     (VOID)This;
 719
 720     *Position = g_efi_file_replace.CurPos;
 721
 722     return EFI_SUCCESS;
 723 }
 724
 725
 726 STATIC EFI_STATUS EFIAPI
 727 ventoy_wrapper_file_get_info(EFI_FILE_HANDLE This, EFI_GUID *Type, UINTN *Len, VOID *Data)
 728 {
 729     EFI_FILE_INFO *Info = (EFI_FILE_INFO *) Data;
 730
 731     debug("ventoy_wrapper_file_get_info ... %u", *Len);
 732
 733     if (!CompareGuid(Type, &gEfiFileInfoGuid))
 734     {
 735         return EFI_INVALID_PARAMETER;
 736     }
 737
 738     if (*Len == 0)
 739     {
 740         *Len = 384;
 741         return EFI_BUFFER_TOO_SMALL;
 742     }
 743
 744     ZeroMem(Data, sizeof(EFI_FILE_INFO));
 745
 746     Info->Size = sizeof(EFI_FILE_INFO);
 747     Info->FileSize = g_efi_file_replace.FileSizeBytes;
 748     Info->PhysicalSize = g_efi_file_replace.FileSizeBytes;
 749     Info->Attribute = EFI_FILE_READ_ONLY;
 750     //Info->FileName = EFI_FILE_READ_ONLY;
 751
 752     *Len = Info->Size;
 753
 754     return EFI_SUCCESS;
 755 }
 756
 757 STATIC EFI_STATUS EFIAPI
 758 ventoy_wrapper_file_read(EFI_FILE_HANDLE This, UINTN *Len, VOID *Data)
 759 {
 760     EFI_LBA Lba;
 761     UINTN ReadLen = *Len;
 762
 763     (VOID)This;
 764
 765     debug("ventoy_wrapper_file_read ... %u", *Len);
 766
 767     if (g_efi_file_replace.CurPos + ReadLen > g_efi_file_replace.FileSizeBytes)
 768     {
 769         ReadLen = g_efi_file_replace.FileSizeBytes - g_efi_file_replace.CurPos;
 770     }
 771
 772     Lba = g_efi_file_replace.CurPos / 2048 + g_efi_file_replace.BlockIoSectorStart;
 773
 774     ventoy_block_io_read(NULL, 0, Lba, ReadLen, Data);
 775
 776     *Len = ReadLen;
 777
 778     g_efi_file_replace.CurPos += ReadLen;
 779
 780     return EFI_SUCCESS;
 781 }
 782
 783 STATIC EFI_STATUS EFIAPI
 784 ventoy_wrapper_file_read_ex(IN EFI_FILE_PROTOCOL *This, IN OUT EFI_FILE_IO_TOKEN *Token)
 785 {
 786         return ventoy_wrapper_file_read(This, &(Token->BufferSize), Token->Buffer);
 787 }
 788
 789 STATIC EFI_STATUS EFIAPI ventoy_wrapper_file_procotol(EFI_FILE_PROTOCOL *File)
 790 {
 791     File->Revision    = EFI_FILE_PROTOCOL_REVISION2;
 792     File->Open        = ventoy_wrapper_fs_open;
 793     File->Close       = ventoy_wrapper_file_close;
 794     File->Delete      = ventoy_wrapper_file_delete;
 795     File->Read        = ventoy_wrapper_file_read;
 796     File->Write       = ventoy_wrapper_file_write;
 797     File->GetPosition = ventoy_wrapper_file_get_pos;
 798     File->SetPosition = ventoy_wrapper_file_set_pos;
 799     File->GetInfo     = ventoy_wrapper_file_get_info;
 800     File->SetInfo     = ventoy_wrapper_file_set_info;
 801     File->Flush       = ventoy_wrapper_file_flush;
 802     File->OpenEx      = ventoy_wrapper_file_open_ex;
 803     File->ReadEx      = ventoy_wrapper_file_read_ex;
 804     File->WriteEx     = ventoy_wrapper_file_write_ex;
 805     File->FlushEx     = ventoy_wrapper_file_flush_ex;
 806
 807     return EFI_SUCCESS;
 808 }
 809
 810 STATIC EFI_STATUS EFIAPI ventoy_wrapper_file_open
 811 (
 812     EFI_FILE_HANDLE This,
 813     EFI_FILE_HANDLE *New,
 814     CHAR16 *Name,
 815     UINT64 Mode,
 816     UINT64 Attributes
 817 )
 818 {
 819     UINT32 i = 0;
 820     UINT32 j = 0;
 821     UINT64 Sectors = 0;
 822     EFI_STATUS Status = EFI_SUCCESS;
 823     CHAR8 TmpName[256];
 824     ventoy_virt_chunk *virt = NULL;
 825
 826     debug("## ventoy_wrapper_file_open <%s> ", Name);
 827
 828     Status = g_original_fopen(This, New, Name, Mode, Attributes);
 829     if (EFI_ERROR(Status))
 830     {
 831         return Status;
 832     }
 833
 834     if (g_file_replace_list && g_file_replace_list->magic == GRUB_FILE_REPLACE_MAGIC &&
 835         g_file_replace_list->new_file_virtual_id < g_virt_chunk_num)
 836     {
 837         AsciiSPrint(TmpName, sizeof(TmpName), "%s", Name);
 838         for (j = 0; j < 4; j++)
 839         {
 840             if (0 == AsciiStrCmp(g_file_replace_list[i].old_file_name[j], TmpName))
 841             {
 842                 g_original_fclose(*New);
 843                 *New = &g_efi_file_replace.WrapperHandle;
 844                 ventoy_wrapper_file_procotol(*New);
 845
 846                 virt = g_virt_chunk + g_file_replace_list->new_file_virtual_id;
 847
 848                 Sectors = (virt->mem_sector_end - virt->mem_sector_start) + (virt->remap_sector_end - virt->remap_sector_start);
 849
 850                 g_efi_file_replace.BlockIoSectorStart = virt->mem_sector_start;
 851                 g_efi_file_replace.FileSizeBytes = Sectors * 2048;
 852
 853                 if (gDebugPrint)
 854                 {
 855                     debug("## ventoy_wrapper_file_open <%s> BlockStart:%lu Sectors:%lu Bytes:%lu", Name,
 856                         g_efi_file_replace.BlockIoSectorStart, Sectors, Sectors * 2048);
 857                     sleep(3);
 858                 }
 859
 860                 return Status;
 861             }
 862         }
 863
 864         if (StrCmp(Name, L"\\EFI\\BOOT") == 0)
 865         {
 866             (*New)->Open = ventoy_wrapper_file_open;
 867         }
 868     }
 869
 870     return Status;
 871 }
 872
 873 EFI_STATUS EFIAPI ventoy_wrapper_open_volume
 874 (
 875     IN EFI_SIMPLE_FILE_SYSTEM_PROTOCOL     *This,
 876     OUT EFI_FILE_PROTOCOL                 **Root
 877 )
 878 {
 879     EFI_STATUS Status = EFI_SUCCESS;
 880
 881     Status = g_original_open_volume(This, Root);
 882     if (!EFI_ERROR(Status))
 883     {
 884         g_original_fopen = (*Root)->Open;
 885         g_original_fclose = (*Root)->Close;
 886         (*Root)->Open = ventoy_wrapper_file_open;
 887     }
 888
 889     return Status;
 890 }
 891
 892 EFI_STATUS EFIAPI ventoy_wrapper_push_openvolume(IN EFI_SIMPLE_FILE_SYSTEM_PROTOCOL_OPEN_VOLUME OpenVolume)
 893 {
 894     g_original_open_volume = OpenVolume;
 895     return EFI_SUCCESS;
 896 }
 897
 898 #if 0
 899 /* For auto skip Windows 'Press any key to boot from CD or DVD ...' */
 900 #endif
 901
 902 STATIC EFI_STATUS EFIAPI ventoy_wrapper_read_key_ex
 903 (
 904     IN  EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *This,
 905     OUT EFI_KEY_DATA                      *KeyData
 906 )
 907 {
 908     /* only hook once before BCD file read */
 909     if (g_keyboard_hook_count == 0 && g_blockio_bcd_read_done == FALSE)
 910     {
 911         g_keyboard_hook_count++;
 912
 913         KeyData->Key.ScanCode = SCAN_DELETE;
 914         KeyData->Key.UnicodeChar = 0;
 915         KeyData->KeyState.KeyShiftState = 0;
 916         KeyData->KeyState.KeyToggleState = 0;
 917
 918         return EFI_SUCCESS;
 919     }
 920
 921     return g_org_read_key_ex(This, KeyData);
 922 }
 923
 924 EFI_STATUS EFIAPI ventoy_wrapper_read_key
 925 (
 926     IN EFI_SIMPLE_TEXT_INPUT_PROTOCOL       *This,
 927     OUT EFI_INPUT_KEY                       *Key
 928 )
 929 {
 930     /* only hook once before BCD file read */
 931     if (g_keyboard_hook_count == 0 && g_blockio_bcd_read_done == FALSE)
 932     {
 933         g_keyboard_hook_count++;
 934
 935         Key->ScanCode = SCAN_DELETE;
 936         Key->UnicodeChar = 0;
 937         return EFI_SUCCESS;
 938     }
 939
 940     return g_org_read_key(This, Key);
 941 }
 942
 943 EFI_STATUS ventoy_hook_keyboard_start(VOID)
 944 {
 945     g_blockio_start_record_bcd = TRUE;
 946     g_blockio_bcd_read_done = FALSE;
 947     g_keyboard_hook_count = 0;
 948
 949     if (g_con_simple_input_ex)
 950     {
 951         g_org_read_key_ex = g_con_simple_input_ex->ReadKeyStrokeEx;
 952         g_con_simple_input_ex->ReadKeyStrokeEx = ventoy_wrapper_read_key_ex;
 953     }
 954
 955     g_org_read_key = gST->ConIn->ReadKeyStroke;
 956     gST->ConIn->ReadKeyStroke = ventoy_wrapper_read_key;
 957
 958     return EFI_SUCCESS;
 959 }
 960
 961 EFI_STATUS ventoy_hook_keyboard_stop(VOID)
 962 {
 963     g_blockio_start_record_bcd = FALSE;
 964     g_blockio_bcd_read_done = FALSE;
 965     g_keyboard_hook_count = 0;
 966
 967     if (g_con_simple_input_ex)
 968     {
 969         g_con_simple_input_ex->ReadKeyStrokeEx = g_org_read_key_ex;
 970     }
 971
 972     gST->ConIn->ReadKeyStroke = g_org_read_key;
 973
 974     return EFI_SUCCESS;
 975 }
 976
 977 #if 0
 978 /* Fixup the 1st cdrom influnce for Windows boot */
 979 #endif
 980
 981 STATIC EFI_STATUS EFIAPI ventoy_wrapper_locate_handle
 982 (
 983     IN     EFI_LOCATE_SEARCH_TYPE    SearchType,
 984     IN     EFI_GUID                 *Protocol,    OPTIONAL
 985     IN     VOID                     *SearchKey,   OPTIONAL
 986     IN OUT UINTN                    *BufferSize,
 987     OUT    EFI_HANDLE               *Buffer
 988 )
 989 {
 990     UINTN i;
 991     EFI_HANDLE Handle = NULL;
 992     EFI_STATUS Status = EFI_SUCCESS;
 993
 994     Status = g_org_locate_handle(SearchType, Protocol, SearchKey, BufferSize, Buffer);
 995
 996     if (EFI_SUCCESS == Status && Protocol && CompareGuid(&gEfiBlockIoProtocolGuid, Protocol))
 997     {
 998         for (i = 0; i < (*BufferSize) / sizeof(EFI_HANDLE); i++)
 999         {
1000             if (Buffer[i] == gBlockData.Handle)
1001             {
1002                 Handle = Buffer[0];
1003                 Buffer[0] = Buffer[i];
1004                 Buffer[i] = Handle;
1005                 break;
1006             }
1007         }
1008     }
1009
1010     return Status;
1011 }
1012
1013 EFI_STATUS ventoy_hook_1st_cdrom_start(VOID)
1014 {
1015     g_org_locate_handle = gBS->LocateHandle;
1016     gBS->LocateHandle = ventoy_wrapper_locate_handle;
1017
1018     return EFI_SUCCESS;
1019 }
1020
1021 EFI_STATUS ventoy_hook_1st_cdrom_stop(VOID)
1022 {
1023     gBS->LocateHandle = g_org_locate_handle;
1024     g_org_locate_handle = NULL;
1025
1026     return EFI_SUCCESS;
1027 }
1028