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