Improvements to protocol support

sdk/xtdk/bltypes.h

@@ -94,8 +94,10 @@ typedef PLIST_ENTRY (*PBL_GET_MODULES_LIST)();
 typedef INT_PTR (*PBL_GET_SECURE_BOOT_STATUS)();
 typedef EFI_STATUS (*PBL_INSTALL_XT_PROTOCOL)(IN PVOID Interface, IN PEFI_GUID Guid);
 typedef EFI_STATUS (*PBL_INVOKE_BOOT_PROTOCOL)(IN PLIST_ENTRY OptionsList);
+typedef EFI_STATUS (*PBL_LOCATE_PROTOCOL_HANDLES)(OUT PEFI_HANDLE *Handles, OUT PUINT_PTR Count, IN PEFI_GUID ProtocolGuid);
 typedef EFI_STATUS (*PBL_OPEN_VOLUME)(IN PEFI_DEVICE_PATH_PROTOCOL DevicePath, OUT PEFI_HANDLE DiskHandle, OUT PEFI_FILE_HANDLE *FsHandle);
-typedef EFI_STATUS (*PBL_OPEN_XT_PROTOCOL)(OUT PEFI_HANDLE Handle, OUT PVOID *ProtocolHandler, IN PEFI_GUID ProtocolGuid);
+typedef EFI_STATUS (*PBL_OPEN_PROTOCOL)(OUT PEFI_HANDLE Handle, OUT PVOID *ProtocolHandler, IN PEFI_GUID ProtocolGuid);
+typedef EFI_STATUS (*PBL_OPEN_PROTOCOL_HANDLE)(IN EFI_HANDLE Handle, OUT PVOID *ProtocolHandler, IN PEFI_GUID ProtocolGuid);
 typedef EFI_STATUS (*PBL_READ_FILE)(IN PEFI_FILE_HANDLE DirHandle, IN CONST PWCHAR FileName, OUT PVOID *FileData, OUT PSIZE_T FileSize);
 typedef EFI_STATUS (*PBL_REGISTER_BOOT_PROTOCOL)(IN PWCHAR SystemType, IN PEFI_GUID BootProtocolGuid);
 typedef VOID (*PBL_REGISTER_XT_BOOT_MENU)(PVOID BootMenuRoutine);
@@ -273,7 +275,9 @@ typedef struct _XTBL_LOADER_PROTOCOL
         PBL_CLOSE_XT_PROTOCOL Close;
         PBL_GET_MODULES_LIST GetModulesList;
         PBL_INSTALL_XT_PROTOCOL Install;
-        PBL_OPEN_XT_PROTOCOL Open;
+        PBL_LOCATE_PROTOCOL_HANDLES LocateHandles;
+        PBL_OPEN_PROTOCOL Open;
+        PBL_OPEN_PROTOCOL_HANDLE OpenHandle;
     } Protocol;
     struct
     {

xtldr/includes/xtldr.h

@@ -147,6 +147,12 @@ XTCDECL
 EFI_STATUS
 BlLoadModules(IN PWCHAR ModulesList);
 
+XTCDECL
+EFI_STATUS
+BlLocateProtocolHandles(OUT PEFI_HANDLE *Handles,
+                        OUT PUINT_PTR Count,
+                        IN PEFI_GUID ProtocolGuid);
+
 XTCDECL
 EFI_STATUS
 BlMemoryAllocatePages(IN UINT64 Pages,
@@ -178,6 +184,12 @@ BlOpenProtocol(OUT PEFI_HANDLE Handle,
                OUT PVOID *ProtocolHandler,
                IN PEFI_GUID ProtocolGuid);
 
+XTCDECL
+EFI_STATUS
+BlOpenProtocolHandle(IN EFI_HANDLE Handle,
+                     OUT PVOID *ProtocolHandler,
+                     IN PEFI_GUID ProtocolGuid);
+
 XTCDECL
 VOID
 BlQueryConsoleMode(OUT PUINT_PTR ResX,

xtldr/protocol.c

@@ -398,7 +398,32 @@ BlLoadModules(IN PWCHAR ModulesList)
 }
 
 /**
- * This routine locates and opens the requested XT Boot Loader or EFI protocol.
+ * Returns an array of handles that support the requested protocol.
+ *
+ * @param Handles
+ *        Supplies the address where a pointer to all handles found for the protocol interface.
+ *
+ * @param Count
+ *        Provides a number of the returned handles.
+ *
+ * @param ProtocolGuid
+ *        Supplies a pointer to the unique protocol GUID.
+ *
+ * @return This routine returns a status code.
+ *
+ * @since XT 1.0
+ */
+XTCDECL
+EFI_STATUS
+BlLocateProtocolHandles(OUT PEFI_HANDLE *Handles,
+                        OUT PUINT_PTR Count,
+                        IN PEFI_GUID ProtocolGuid)
+{
+    return EfiSystemTable->BootServices->LocateHandleBuffer(ByProtocol, ProtocolGuid, NULL, Count, Handles);
+}
+
+/**
+ * Locates and opens the requested XT Boot Loader or EFI protocol.
  *
  * @param Handle
  *        Supplies the address where a pointer to the handle for the protocol interface.
@@ -425,7 +450,7 @@ BlOpenProtocol(OUT PEFI_HANDLE Handle,
     UINT Index;
 
     /* Try to locate the handles */
-    Status = EfiSystemTable->BootServices->LocateHandleBuffer(ByProtocol, ProtocolGuid, NULL, &Count, &Handles);
+    Status = BlLocateProtocolHandles(&Handles, &Count, ProtocolGuid);
     if(Status != STATUS_EFI_SUCCESS)
     {
         /* Unable to get handles */
@@ -439,9 +464,7 @@ BlOpenProtocol(OUT PEFI_HANDLE Handle,
         for(Index = 0; Index < Count; Index++)
         {
             /* Try to open protocol */
-            Status = EfiSystemTable->BootServices->OpenProtocol(Handles[Index], ProtocolGuid,
+            Status = BlOpenProtocolHandle(Handles[Index], ProtocolHandler, ProtocolGuid);
-                                                                ProtocolHandler, EfiImageHandle, NULL,
-                                                                EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL);
 
             /* Check if successfully opened the loader protocol */
             if(Status == STATUS_EFI_SUCCESS)
@@ -467,6 +490,32 @@ BlOpenProtocol(OUT PEFI_HANDLE Handle,
 	return STATUS_EFI_SUCCESS;
 }
 
+/**
+ * Opens the requested XT Boot Loader or EFI protocol, if it is supported by the handle.
+ *
+ * @param Handle
+ *        Supplies a handle for the protocol interface that is being opened.
+ *
+ * @param ProtocolHandler
+ *        Supplies the address where a pointer to the opened protocol is returned.
+ *
+ * @param ProtocolGuid
+ *        Supplies a pointer to the unique protocol GUID.
+ *
+ * @return This routine returns a status code.
+ *
+ * @since XT 1.0
+ */
+XTCDECL
+EFI_STATUS
+BlOpenProtocolHandle(IN EFI_HANDLE Handle,
+                     OUT PVOID *ProtocolHandler,
+                     IN PEFI_GUID ProtocolGuid)
+{
+    return EfiSystemTable->BootServices->OpenProtocol(Handle, ProtocolGuid, ProtocolHandler, EfiImageHandle,
+                                                      NULL, EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL);
+}
+
 /**
  * Registers a boot menu callback routine, that will be used to display alternative boot menu.
  *
@@ -586,8 +635,10 @@ BlpInstallXtLoaderProtocol()
     BlpLdrProtocol.Memory.ZeroMemory = RtlZeroMemory;
     BlpLdrProtocol.Protocol.Close = BlCloseProtocol;
     BlpLdrProtocol.Protocol.GetModulesList = BlGetModulesList;
-    BlpLdrProtocol.Protocol.Open = BlOpenProtocol;
     BlpLdrProtocol.Protocol.Install = BlInstallProtocol;
+    BlpLdrProtocol.Protocol.LocateHandles = BlLocateProtocolHandles;
+    BlpLdrProtocol.Protocol.Open = BlOpenProtocol;
+    BlpLdrProtocol.Protocol.OpenHandle = BlOpenProtocolHandle;
     BlpLdrProtocol.Tui.DisplayErrorDialog = BlDisplayErrorDialog;
     BlpLdrProtocol.Tui.DisplayInfoDialog = BlDisplayInfoDialog;
     BlpLdrProtocol.Tui.DisplayInputDialog = BlDisplayInputDialog;

xtldr/volume.c

@@ -613,7 +613,7 @@ BlpDiscoverEfiBlockDevices(OUT PLIST_ENTRY BlockDevices)
     EFI_STATUS Status;
 
     /* Locate handles which support the disk I/O interface */
-    Status = EfiSystemTable->BootServices->LocateHandleBuffer(ByProtocol, &IoGuid, NULL, &HandlesCount, &Handles);
+    Status = BlLocateProtocolHandles(&Handles, &HandlesCount, &IoGuid);
     if(Status != STATUS_EFI_SUCCESS)
     {
         /* Failed to locate handles */
@@ -629,8 +629,7 @@ BlpDiscoverEfiBlockDevices(OUT PLIST_ENTRY BlockDevices)
 
         /* Open I/O protocol for given handle */
         Io = NULL;
-        Status = EfiSystemTable->BootServices->OpenProtocol(Handles[Index], &IoGuid, (PVOID *)&Io, EfiImageHandle,
+        Status = BlOpenProtocolHandle(Handles[Index], (PVOID *)&Io, &IoGuid);
-                                                            NULL, EFI_OPEN_PROTOCOL_BY_HANDLE_PROTOCOL);
         if(Status != STATUS_EFI_SUCCESS || Io == NULL)
         {
             /* Failed to open I/O protocol, skip it */