]> glassweightruler.freedombox.rocks Git - Ventoy.git/blob - EDK2/edk2_mod/edk2-edk2-stable201911/MdeModulePkg/Application/Ventoy/VentoyProtocol.c
projects / Ventoy.git / blob
? search:


a6e168e05a1e0d873adb8c84226cf67089acf03b
[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 + secNum > VirtSec)
479 {
480 secNum = VirtSec - Lba;
481 }
482
483 if (secNum > g_sector_flag_num)
484 {
485 cur_flag = AllocatePool(secNum * sizeof(ventoy_sector_flag));
486 if (NULL == cur_flag)
487 {
488 return EFI_OUT_OF_RESOURCES;
489 }
490
491 FreePool(g_sector_flag);
492 g_sector_flag = cur_flag;
493 g_sector_flag_num = secNum;
494 }
495
496 for (curlba = Lba, cur_flag = g_sector_flag, j = 0; j < secNum; j++, curlba++, cur_flag++)
497 {
498 cur_flag->flag = 0;
499 for (node = g_virt_chunk, i = 0; i < g_virt_chunk_num; i++, node++)
500 {
501 if (curlba >= node->mem_sector_start && curlba < node->mem_sector_end)
502 {
503 CopyMem((UINT8 *)Buffer + j * 2048,
504 (char *)g_virt_chunk + node->mem_sector_offset + (curlba - node->mem_sector_start) * 2048,
505 2048);
506 cur_flag->flag = 1;
507 break;
508 }
509 else if (curlba >= node->remap_sector_start && curlba < node->remap_sector_end)
510 {
511 cur_flag->remap_lba = node->org_sector_start + curlba - node->remap_sector_start;
512 cur_flag->flag = 2;
513 break;
514 }
515 }
516 }
517
518 for (curlba = Lba, cur_flag = g_sector_flag, j = 0; j < secNum; j++, curlba++, cur_flag++)
519 {
520 if (cur_flag->flag == 2)
521 {
522 if (lastlba == 0)
523 {
524 lastbuffer = (UINT8 *)Buffer + j * 2048;
525 lastlba = cur_flag->remap_lba;
526 lbacount = 1;
527 }
528 else if (lastlba + lbacount == cur_flag->remap_lba)
529 {
530 lbacount++;
531 }
532 else
533 {
534 ventoy_read_iso_sector(lastlba, lbacount, lastbuffer);
535 lastbuffer = (UINT8 *)Buffer + j * 2048;
536 lastlba = cur_flag->remap_lba;
537 lbacount = 1;
538 }
539 }
540 }
541
542 if (lbacount > 0)
543 {
544 ventoy_read_iso_sector(lastlba, lbacount, lastbuffer);
545 }
546
547 return EFI_SUCCESS;
548 }
549
550 EFI_STATUS EFIAPI ventoy_block_io_write
551 (
552 IN EFI_BLOCK_IO_PROTOCOL *This,
553 IN UINT32 MediaId,
554 IN EFI_LBA Lba,
555 IN UINTN BufferSize,
556 IN VOID *Buffer
557 )
558 {
559 UINT32 secNum = 0;
560 UINT64 offset = 0;
561
562 (VOID)This;
563 (VOID)MediaId;
564
565 if (!gSector512Mode)
566 {
567 return EFI_WRITE_PROTECTED;
568 }
569
570 secNum = BufferSize / 2048;
571 offset = Lba * 2048;
572
573 return ventoy_write_iso_sector(Lba, secNum, Buffer);
574 }
575
576 EFI_STATUS EFIAPI ventoy_block_io_flush(IN EFI_BLOCK_IO_PROTOCOL *This)
577 {
578 (VOID)This;
579 return EFI_SUCCESS;
580 }
581
582
583 EFI_STATUS EFIAPI ventoy_fill_device_path(VOID)
584 {
585 UINTN NameLen = 0;
586 UINT8 TmpBuf[128] = {0};
587 VENDOR_DEVICE_PATH *venPath = NULL;
588
589 venPath = (VENDOR_DEVICE_PATH *)TmpBuf;
590 NameLen = StrSize(VTOY_BLOCK_DEVICE_PATH_NAME);
591 venPath->Header.Type = HARDWARE_DEVICE_PATH;
592 venPath->Header.SubType = HW_VENDOR_DP;
593 venPath->Header.Length[0] = sizeof(VENDOR_DEVICE_PATH) + NameLen;
594 venPath->Header.Length[1] = 0;
595 CopyMem(&venPath->Guid, &gVtoyBlockDevicePathGuid, sizeof(EFI_GUID));
596 CopyMem(venPath + 1, VTOY_BLOCK_DEVICE_PATH_NAME, NameLen);
597
598 gBlockData.Path = AppendDevicePathNode(NULL, (EFI_DEVICE_PATH_PROTOCOL *)TmpBuf);
599 gBlockData.DevicePathCompareLen = sizeof(VENDOR_DEVICE_PATH) + NameLen;
600
601 debug("gBlockData.Path=<%s>\n", ConvertDevicePathToText(gBlockData.Path, FALSE, FALSE));
602
603 return EFI_SUCCESS;
604 }
605
606 EFI_STATUS EFIAPI ventoy_connect_driver(IN EFI_HANDLE ControllerHandle, IN CONST CHAR16 *DrvName)
607 {
608 UINTN i = 0;
609 UINTN Count = 0;
610 CHAR16 *DriverName = NULL;
611 EFI_HANDLE *Handles = NULL;
612 EFI_HANDLE DrvHandles[2] = { NULL };
613 EFI_STATUS Status = EFI_SUCCESS;
614 EFI_COMPONENT_NAME_PROTOCOL *NameProtocol = NULL;
615 EFI_COMPONENT_NAME2_PROTOCOL *Name2Protocol = NULL;
616
617 debug("ventoy_connect_driver <%s>...", DrvName);
618
619 Status = gBS->LocateHandleBuffer(ByProtocol, &gEfiComponentName2ProtocolGuid,
620 NULL, &Count, &Handles);
621 if (EFI_ERROR(Status))
622 {
623 return Status;
624 }
625
626 for (i = 0; i < Count; i++)
627 {
628 Status = gBS->HandleProtocol(Handles[i], &gEfiComponentName2ProtocolGuid, (VOID **)&Name2Protocol);
629 if (EFI_ERROR(Status))
630 {
631 continue;
632 }
633
634 Status = Name2Protocol->GetDriverName(Name2Protocol, "en", &DriverName);
635 if (EFI_ERROR(Status) || NULL == DriverName)
636 {
637 continue;
638 }
639
640 if (StrStr(DriverName, DrvName))
641 {
642 debug("Find driver name2:<%s>: <%s>", DriverName, DrvName);
643 DrvHandles[0] = Handles[i];
644 break;
645 }
646 }
647
648 if (i < Count)
649 {
650 Status = gBS->ConnectController(ControllerHandle, DrvHandles, NULL, TRUE);
651 debug("ventoy_connect_driver:<%s> <%r>", DrvName, Status);
652 goto end;
653 }
654
655 debug("%s NOT found, now try COMPONENT_NAME", DrvName);
656
657 Count = 0;
658 FreePool(Handles);
659 Handles = NULL;
660
661 Status = gBS->LocateHandleBuffer(ByProtocol, &gEfiComponentNameProtocolGuid,
662 NULL, &Count, &Handles);
663 if (EFI_ERROR(Status))
664 {
665 return Status;
666 }
667
668 for (i = 0; i < Count; i++)
669 {
670 Status = gBS->HandleProtocol(Handles[i], &gEfiComponentNameProtocolGuid, (VOID **)&NameProtocol);
671 if (EFI_ERROR(Status))
672 {
673 continue;
674 }
675
676 Status = NameProtocol->GetDriverName(NameProtocol, "en", &DriverName);
677 if (EFI_ERROR(Status))
678 {
679 continue;
680 }
681
682 if (StrStr(DriverName, DrvName))
683 {
684 debug("Find driver name:<%s>: <%s>", DriverName, DrvName);
685 DrvHandles[0] = Handles[i];
686 break;
687 }
688 }
689
690 if (i < Count)
691 {
692 Status = gBS->ConnectController(ControllerHandle, DrvHandles, NULL, TRUE);
693 debug("ventoy_connect_driver:<%s> <%r>", DrvName, Status);
694 goto end;
695 }
696
697 Status = EFI_NOT_FOUND;
698
699 end:
700 FreePool(Handles);
701
702 return Status;
703 }
704
705 EFI_STATUS EFIAPI ventoy_block_io_read_512
706 (
707 IN EFI_BLOCK_IO_PROTOCOL *This,
708 IN UINT32 MediaId,
709 IN EFI_LBA Lba,
710 IN UINTN BufferSize,
711 OUT VOID *Buffer
712 )
713 {
714 EFI_LBA Mod;
715 UINTN ReadSize;
716 UINT8 *CurBuf = NULL;
717 EFI_STATUS Status = EFI_SUCCESS;
718
719 debug("ventoy_block_io_read_512 %lu %lu\n", Lba, BufferSize / 512);
720
721 CurBuf = (UINT8 *)Buffer;
722
723 Mod = Lba % 4;
724 if (Mod > 0)
725 {
726 Status |= g_sector_2048_read(This, MediaId, Lba / 4, 2048, g_sector_buf);
727
728 if (BufferSize <= (4 - Mod) * 512)
729 {
730 CopyMem(CurBuf, g_sector_buf + Mod * 512, BufferSize);
731 return EFI_SUCCESS;
732 }
733 else
734 {
735 ReadSize = (4 - Mod) * 512;
736 CopyMem(CurBuf, g_sector_buf + Mod * 512, ReadSize);
737 CurBuf += ReadSize;
738 Lba += (4 - Mod);
739 BufferSize -= ReadSize;
740 }
741 }
742
743 if (BufferSize >= 2048)
744 {
745 ReadSize = BufferSize / 2048 * 2048;
746
747 Status |= g_sector_2048_read(This, MediaId, Lba / 4, ReadSize, CurBuf);
748 CurBuf += ReadSize;
749
750 Lba += ReadSize / 512;
751 BufferSize -= ReadSize;
752 }
753
754 if (BufferSize > 0)
755 {
756 Status |= g_sector_2048_read(This, MediaId, Lba / 4, 2048, g_sector_buf);
757 CopyMem(CurBuf, g_sector_buf, BufferSize);
758 }
759
760 return Status;
761 }
762
763 EFI_STATUS EFIAPI ventoy_block_io_write_512
764 (
765 IN EFI_BLOCK_IO_PROTOCOL *This,
766 IN UINT32 MediaId,
767 IN EFI_LBA Lba,
768 IN UINTN BufferSize,
769 IN VOID *Buffer
770 )
771 {
772 EFI_LBA Mod;
773 UINTN ReadSize;
774 UINT8 *CurBuf = NULL;
775 EFI_STATUS Status = EFI_SUCCESS;
776
777 debug("ventoy_block_io_write_512 %lu %lu\n", Lba, BufferSize / 512);
778
779 CurBuf = (UINT8 *)Buffer;
780
781 Mod = Lba % 4;
782 if (Mod > 0)
783 {
784 Status |= g_sector_2048_read(This, MediaId, Lba / 4, 2048, g_sector_buf);
785
786 if (BufferSize <= (4 - Mod) * 512)
787 {
788 CopyMem(g_sector_buf + Mod * 512, CurBuf, BufferSize);
789 return g_sector_2048_write(This, MediaId, Lba / 4, 2048, g_sector_buf);
790 }
791 else
792 {
793 ReadSize = (4 - Mod) * 512;
794 CopyMem(g_sector_buf + Mod * 512, CurBuf, ReadSize);
795 g_sector_2048_write(This, MediaId, Lba / 4, 2048, g_sector_buf);
796
797 CurBuf += ReadSize;
798 Lba += (4 - Mod);
799 BufferSize -= ReadSize;
800 }
801 }
802
803 if (BufferSize >= 2048)
804 {
805 ReadSize = BufferSize / 2048 * 2048;
806
807 Status |= g_sector_2048_write(This, MediaId, Lba / 4, ReadSize, CurBuf);
808 CurBuf += ReadSize;
809
810 Lba += ReadSize / 512;
811 BufferSize -= ReadSize;
812 }
813
814 if (BufferSize > 0)
815 {
816 Status |= g_sector_2048_read(This, MediaId, Lba / 4, 2048, g_sector_buf);
817
818 CopyMem(g_sector_buf, CurBuf, BufferSize);
819 g_sector_2048_write(This, MediaId, Lba / 4, 2048, g_sector_buf);
820 }
821
822 return Status;
823 }
824
825 EFI_STATUS EFIAPI ventoy_install_blockio(IN EFI_HANDLE ImageHandle, IN UINT64 ImgSize)
826 {
827 EFI_STATUS Status = EFI_SUCCESS;
828 EFI_BLOCK_IO_PROTOCOL *pBlockIo = &(gBlockData.BlockIo);
829
830 ventoy_fill_device_path();
831
832 debug("install block io protocol %p", ImageHandle);
833 ventoy_debug_pause();
834
835 if (gSector512Mode)
836 {
837 gBlockData.Media.BlockSize = 512;
838 gBlockData.Media.LastBlock = ImgSize / 512 - 1;
839 }
840 else
841 {
842 gBlockData.Media.BlockSize = 2048;
843 gBlockData.Media.LastBlock = ImgSize / 2048 - 1;
844 }
845
846 gBlockData.Media.ReadOnly = TRUE;
847 gBlockData.Media.MediaPresent = 1;
848 gBlockData.Media.LogicalBlocksPerPhysicalBlock = 1;
849
850 pBlockIo->Revision = EFI_BLOCK_IO_PROTOCOL_REVISION3;
851 pBlockIo->Media = &(gBlockData.Media);
852 pBlockIo->Reset = ventoy_block_io_reset;
853
854 if (gSector512Mode)
855 {
856 g_sector_2048_read = gMemdiskMode ? ventoy_block_io_ramdisk_read : ventoy_block_io_read;
857 g_sector_2048_write = gMemdiskMode ? ventoy_block_io_ramdisk_write : ventoy_block_io_write;
858 pBlockIo->ReadBlocks = ventoy_block_io_read_512;
859 pBlockIo->WriteBlocks = ventoy_block_io_write_512;
860 }
861 else
862 {
863 pBlockIo->ReadBlocks = gMemdiskMode ? ventoy_block_io_ramdisk_read : ventoy_block_io_read;
864 pBlockIo->WriteBlocks = ventoy_block_io_write;
865 }
866
867 pBlockIo->FlushBlocks = ventoy_block_io_flush;
868
869 Status = gBS->InstallMultipleProtocolInterfaces(&gBlockData.Handle,
870 &gEfiBlockIoProtocolGuid, &gBlockData.BlockIo,
871 &gEfiDevicePathProtocolGuid, gBlockData.Path,
872 NULL);
873 debug("Install protocol %r %p", Status, gBlockData.Handle);
874 if (EFI_ERROR(Status))
875 {
876 return Status;
877 }
878
879 Status = ventoy_connect_driver(gBlockData.Handle, L"Disk I/O Driver");
880 debug("Connect disk IO driver %r", Status);
881
882 Status = ventoy_connect_driver(gBlockData.Handle, L"Partition Driver");
883 debug("Connect partition driver %r", Status);
884 if (EFI_ERROR(Status))
885 {
886 Status = gBS->ConnectController(gBlockData.Handle, NULL, NULL, TRUE);
887 debug("Connect all controller %r", Status);
888 }
889
890 ventoy_debug_pause();
891
892 return EFI_SUCCESS;
893 }
894
895 #if 0
896 /* For file replace */
897 #endif
898
899 STATIC EFI_STATUS EFIAPI
900 ventoy_wrapper_fs_open(EFI_FILE_HANDLE This, EFI_FILE_HANDLE *New, CHAR16 *Name, UINT64 Mode, UINT64 Attributes)
901 {
902 (VOID)This;
903 (VOID)New;
904 (VOID)Name;
905 (VOID)Mode;
906 (VOID)Attributes;
907 return EFI_SUCCESS;
908 }
909
910 STATIC EFI_STATUS EFIAPI
911 ventoy_wrapper_file_open_ex(EFI_FILE_HANDLE This, EFI_FILE_HANDLE *New, CHAR16 *Name, UINT64 Mode, UINT64 Attributes, EFI_FILE_IO_TOKEN *Token)
912 {
913 return ventoy_wrapper_fs_open(This, New, Name, Mode, Attributes);
914 }
915
916 STATIC EFI_STATUS EFIAPI
917 ventoy_wrapper_file_delete(EFI_FILE_HANDLE This)
918 {
919 (VOID)This;
920 return EFI_SUCCESS;
921 }
922
923 STATIC EFI_STATUS EFIAPI
924 ventoy_wrapper_file_set_info(EFI_FILE_HANDLE This, EFI_GUID *Type, UINTN Len, VOID *Data)
925 {
926 return EFI_SUCCESS;
927 }
928
929 STATIC EFI_STATUS EFIAPI
930 ventoy_wrapper_file_flush(EFI_FILE_HANDLE This)
931 {
932 (VOID)This;
933 return EFI_SUCCESS;
934 }
935
936 /* Ex version */
937 STATIC EFI_STATUS EFIAPI
938 ventoy_wrapper_file_flush_ex(EFI_FILE_HANDLE This, EFI_FILE_IO_TOKEN *Token)
939 {
940 (VOID)This;
941 (VOID)Token;
942 return EFI_SUCCESS;
943 }
944
945
946 STATIC EFI_STATUS EFIAPI
947 ventoy_wrapper_file_write(EFI_FILE_HANDLE This, UINTN *Len, VOID *Data)
948 {
949 (VOID)This;
950 (VOID)Len;
951 (VOID)Data;
952
953 return EFI_WRITE_PROTECTED;
954 }
955
956 STATIC EFI_STATUS EFIAPI
957 ventoy_wrapper_file_write_ex(IN EFI_FILE_PROTOCOL *This, IN OUT EFI_FILE_IO_TOKEN *Token)
958 {
959 return ventoy_wrapper_file_write(This, &(Token->BufferSize), Token->Buffer);
960 }
961
962
963 STATIC EFI_STATUS EFIAPI
964 ventoy_wrapper_file_close(EFI_FILE_HANDLE This)
965 {
966 (VOID)This;
967 return EFI_SUCCESS;
968 }
969
970
971 STATIC EFI_STATUS EFIAPI
972 ventoy_wrapper_file_set_pos(EFI_FILE_HANDLE This, UINT64 Position)
973 {
974 (VOID)This;
975
976 if (Position <= g_efi_file_replace.FileSizeBytes)
977 {
978 g_efi_file_replace.CurPos = Position;
979 }
980 else
981 {
982 g_efi_file_replace.CurPos = g_efi_file_replace.FileSizeBytes;
983 }
984
985 return EFI_SUCCESS;
986 }
987
988 STATIC EFI_STATUS EFIAPI
989 ventoy_wrapper_file_get_pos(EFI_FILE_HANDLE This, UINT64 *Position)
990 {
991 (VOID)This;
992
993 *Position = g_efi_file_replace.CurPos;
994
995 return EFI_SUCCESS;
996 }
997
998
999 STATIC EFI_STATUS EFIAPI
1000 ventoy_wrapper_file_get_info(EFI_FILE_HANDLE This, EFI_GUID *Type, UINTN *Len, VOID *Data)
1001 {
1002 EFI_FILE_INFO *Info = (EFI_FILE_INFO *) Data;
1003
1004 debug("ventoy_wrapper_file_get_info ... %u", *Len);
1005
1006 if (!CompareGuid(Type, &gEfiFileInfoGuid))
1007 {
1008 return EFI_INVALID_PARAMETER;
1009 }
1010
1011 if (*Len == 0)
1012 {
1013 *Len = 384;
1014 return EFI_BUFFER_TOO_SMALL;
1015 }
1016
1017 ZeroMem(Data, sizeof(EFI_FILE_INFO));
1018
1019 Info->Size = sizeof(EFI_FILE_INFO);
1020 Info->FileSize = g_efi_file_replace.FileSizeBytes;
1021 Info->PhysicalSize = g_efi_file_replace.FileSizeBytes;
1022 Info->Attribute = EFI_FILE_READ_ONLY;
1023 //Info->FileName = EFI_FILE_READ_ONLY;
1024
1025 *Len = Info->Size;
1026
1027 return EFI_SUCCESS;
1028 }
1029
1030 STATIC EFI_STATUS EFIAPI
1031 ventoy_wrapper_file_read(EFI_FILE_HANDLE This, UINTN *Len, VOID *Data)
1032 {
1033 EFI_LBA Lba;
1034 UINTN ReadLen = *Len;
1035
1036 (VOID)This;
1037
1038 debug("ventoy_wrapper_file_read ... %u", *Len);
1039
1040 if (g_efi_file_replace.CurPos + ReadLen > g_efi_file_replace.FileSizeBytes)
1041 {
1042 ReadLen = g_efi_file_replace.FileSizeBytes - g_efi_file_replace.CurPos;
1043 }
1044
1045 Lba = g_efi_file_replace.CurPos / 2048 + g_efi_file_replace.BlockIoSectorStart;
1046
1047 ventoy_block_io_read(NULL, 0, Lba, ReadLen, Data);
1048
1049 *Len = ReadLen;
1050
1051 g_efi_file_replace.CurPos += ReadLen;
1052
1053 return EFI_SUCCESS;
1054 }
1055
1056 STATIC EFI_STATUS EFIAPI
1057 ventoy_wrapper_file_read_ex(IN EFI_FILE_PROTOCOL *This, IN OUT EFI_FILE_IO_TOKEN *Token)
1058 {
1059 return ventoy_wrapper_file_read(This, &(Token->BufferSize), Token->Buffer);
1060 }
1061
1062 STATIC EFI_STATUS EFIAPI ventoy_wrapper_file_procotol(EFI_FILE_PROTOCOL *File)
1063 {
1064 File->Revision = EFI_FILE_PROTOCOL_REVISION2;
1065 File->Open = ventoy_wrapper_fs_open;
1066 File->Close = ventoy_wrapper_file_close;
1067 File->Delete = ventoy_wrapper_file_delete;
1068 File->Read = ventoy_wrapper_file_read;
1069 File->Write = ventoy_wrapper_file_write;
1070 File->GetPosition = ventoy_wrapper_file_get_pos;
1071 File->SetPosition = ventoy_wrapper_file_set_pos;
1072 File->GetInfo = ventoy_wrapper_file_get_info;
1073 File->SetInfo = ventoy_wrapper_file_set_info;
1074 File->Flush = ventoy_wrapper_file_flush;
1075 File->OpenEx = ventoy_wrapper_file_open_ex;
1076 File->ReadEx = ventoy_wrapper_file_read_ex;
1077 File->WriteEx = ventoy_wrapper_file_write_ex;
1078 File->FlushEx = ventoy_wrapper_file_flush_ex;
1079
1080 return EFI_SUCCESS;
1081 }
1082
1083 STATIC EFI_STATUS EFIAPI ventoy_wrapper_file_open
1084 (
1085 EFI_FILE_HANDLE This,
1086 EFI_FILE_HANDLE *New,
1087 CHAR16 *Name,
1088 UINT64 Mode,
1089 UINT64 Attributes
1090 )
1091 {
1092 UINT32 i = 0;
1093 UINT32 j = 0;
1094 UINT64 Sectors = 0;
1095 EFI_STATUS Status = EFI_SUCCESS;
1096 CHAR8 TmpName[256];
1097 ventoy_virt_chunk *virt = NULL;
1098
1099 debug("## ventoy_wrapper_file_open <%s> ", Name);
1100
1101 Status = g_original_fopen(This, New, Name, Mode, Attributes);
1102 if (EFI_ERROR(Status))
1103 {
1104 return Status;
1105 }
1106
1107 if (g_file_replace_list && g_file_replace_list->magic == GRUB_FILE_REPLACE_MAGIC &&
1108 g_file_replace_list->new_file_virtual_id < g_virt_chunk_num)
1109 {
1110 AsciiSPrint(TmpName, sizeof(TmpName), "%s", Name);
1111 for (j = 0; j < 4; j++)
1112 {
1113 if (0 == AsciiStrCmp(g_file_replace_list[i].old_file_name[j], TmpName))
1114 {
1115 g_original_fclose(*New);
1116 *New = &g_efi_file_replace.WrapperHandle;
1117 ventoy_wrapper_file_procotol(*New);
1118
1119 virt = g_virt_chunk + g_file_replace_list->new_file_virtual_id;
1120
1121 Sectors = (virt->mem_sector_end - virt->mem_sector_start) + (virt->remap_sector_end - virt->remap_sector_start);
1122
1123 g_efi_file_replace.BlockIoSectorStart = virt->mem_sector_start;
1124 g_efi_file_replace.FileSizeBytes = Sectors * 2048;
1125
1126 if (gDebugPrint)
1127 {
1128 debug("## ventoy_wrapper_file_open <%s> BlockStart:%lu Sectors:%lu Bytes:%lu", Name,
1129 g_efi_file_replace.BlockIoSectorStart, Sectors, Sectors * 2048);
1130 sleep(3);
1131 }
1132
1133 return Status;
1134 }
1135 }
1136
1137 if (StrCmp(Name, L"\\EFI\\BOOT") == 0)
1138 {
1139 (*New)->Open = ventoy_wrapper_file_open;
1140 }
1141 }
1142
1143 return Status;
1144 }
1145
1146 EFI_STATUS EFIAPI ventoy_wrapper_open_volume
1147 (
1148 IN EFI_SIMPLE_FILE_SYSTEM_PROTOCOL *This,
1149 OUT EFI_FILE_PROTOCOL **Root
1150 )
1151 {
1152 EFI_STATUS Status = EFI_SUCCESS;
1153
1154 Status = g_original_open_volume(This, Root);
1155 if (!EFI_ERROR(Status))
1156 {
1157 g_original_fopen = (*Root)->Open;
1158 g_original_fclose = (*Root)->Close;
1159 (*Root)->Open = ventoy_wrapper_file_open;
1160 }
1161
1162 return Status;
1163 }
1164
1165 EFI_STATUS EFIAPI ventoy_wrapper_push_openvolume(IN EFI_SIMPLE_FILE_SYSTEM_PROTOCOL_OPEN_VOLUME OpenVolume)
1166 {
1167 g_original_open_volume = OpenVolume;
1168 return EFI_SUCCESS;
1169 }
1170
1171 #if 0
1172 /* For auto skip Windows 'Press any key to boot from CD or DVD ...' */
1173 #endif
1174
1175 STATIC EFI_STATUS EFIAPI ventoy_wrapper_read_key_ex
1176 (
1177 IN EFI_SIMPLE_TEXT_INPUT_EX_PROTOCOL *This,
1178 OUT EFI_KEY_DATA *KeyData
1179 )
1180 {
1181 /* only hook once before BCD file read */
1182 if (g_keyboard_hook_count == 0 && g_blockio_bcd_read_done == FALSE)
1183 {
1184 g_keyboard_hook_count++;
1185
1186 KeyData->Key.ScanCode = SCAN_DELETE;
1187 KeyData->Key.UnicodeChar = 0;
1188 KeyData->KeyState.KeyShiftState = 0;
1189 KeyData->KeyState.KeyToggleState = 0;
1190
1191 return EFI_SUCCESS;
1192 }
1193
1194 return g_org_read_key_ex(This, KeyData);
1195 }
1196
1197 EFI_STATUS EFIAPI ventoy_wrapper_read_key
1198 (
1199 IN EFI_SIMPLE_TEXT_INPUT_PROTOCOL *This,
1200 OUT EFI_INPUT_KEY *Key
1201 )
1202 {
1203 /* only hook once before BCD file read */
1204 if (g_keyboard_hook_count == 0 && g_blockio_bcd_read_done == FALSE)
1205 {
1206 g_keyboard_hook_count++;
1207
1208 Key->ScanCode = SCAN_DELETE;
1209 Key->UnicodeChar = 0;
1210 return EFI_SUCCESS;
1211 }
1212
1213 return g_org_read_key(This, Key);
1214 }
1215
1216 EFI_STATUS ventoy_hook_keyboard_start(VOID)
1217 {
1218 g_blockio_start_record_bcd = TRUE;
1219 g_blockio_bcd_read_done = FALSE;
1220 g_keyboard_hook_count = 0;
1221
1222 if (g_con_simple_input_ex)
1223 {
1224 g_org_read_key_ex = g_con_simple_input_ex->ReadKeyStrokeEx;
1225 g_con_simple_input_ex->ReadKeyStrokeEx = ventoy_wrapper_read_key_ex;
1226 }
1227
1228 g_org_read_key = gST->ConIn->ReadKeyStroke;
1229 gST->ConIn->ReadKeyStroke = ventoy_wrapper_read_key;
1230
1231 return EFI_SUCCESS;
1232 }
1233
1234 EFI_STATUS ventoy_hook_keyboard_stop(VOID)
1235 {
1236 g_blockio_start_record_bcd = FALSE;
1237 g_blockio_bcd_read_done = FALSE;
1238 g_keyboard_hook_count = 0;
1239
1240 if (g_con_simple_input_ex)
1241 {
1242 g_con_simple_input_ex->ReadKeyStrokeEx = g_org_read_key_ex;
1243 }
1244
1245 gST->ConIn->ReadKeyStroke = g_org_read_key;
1246
1247 return EFI_SUCCESS;
1248 }
1249
1250 #if 0
1251 /* Fixup the 1st cdrom influnce for Windows boot */
1252 #endif
1253
1254 STATIC EFI_STATUS EFIAPI ventoy_wrapper_locate_handle
1255 (
1256 IN EFI_LOCATE_SEARCH_TYPE SearchType,
1257 IN EFI_GUID *Protocol, OPTIONAL
1258 IN VOID *SearchKey, OPTIONAL
1259 IN OUT UINTN *BufferSize,
1260 OUT EFI_HANDLE *Buffer
1261 )
1262 {
1263 UINTN i;
1264 EFI_HANDLE Handle = NULL;
1265 EFI_STATUS Status = EFI_SUCCESS;
1266
1267 Status = g_org_locate_handle(SearchType, Protocol, SearchKey, BufferSize, Buffer);
1268
1269 if (EFI_SUCCESS == Status && Protocol && CompareGuid(&gEfiBlockIoProtocolGuid, Protocol))
1270 {
1271 for (i = 0; i < (*BufferSize) / sizeof(EFI_HANDLE); i++)
1272 {
1273 if (Buffer[i] == gBlockData.Handle)
1274 {
1275 Handle = Buffer[0];
1276 Buffer[0] = Buffer[i];
1277 Buffer[i] = Handle;
1278 break;
1279 }
1280 }
1281 }
1282
1283 return Status;
1284 }
1285
1286 EFI_STATUS ventoy_hook_1st_cdrom_start(VOID)
1287 {
1288 g_org_locate_handle = gBS->LocateHandle;
1289 gBS->LocateHandle = ventoy_wrapper_locate_handle;
1290
1291 return EFI_SUCCESS;
1292 }
1293
1294 EFI_STATUS ventoy_hook_1st_cdrom_stop(VOID)
1295 {
1296 gBS->LocateHandle = g_org_locate_handle;
1297 g_org_locate_handle = NULL;
1298
1299 return EFI_SUCCESS;
1300 }
1301
https://github.com/ventoy/Ventoy