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