MdePkg/DxeRngLib: Refactor Rng algorithm selection

Add a library constructor which:
- locate the RNG prototocol and keep a reference to it in order to avoid
  locating it multiple times (for each random number generation)
- check which secure algorithm is available on the platform.
  This avoids to try each secure algorithm until finding one
  available for each random number generation call.

Signed-off-by: Pierre Gondois <pierre.gondois@arm.com>

MdePkg/Library/DxeRngLib/DxeRngLib.c

@@ -1,17 +1,129 @@
 /** @file
  Provides an implementation of the library class RngLib that uses the Rng protocol.
 
- Copyright (c) 2023, Arm Limited. All rights reserved.
+ Copyright (c) 2023 - 2024, Arm Limited. All rights reserved.
  Copyright (c) Microsoft Corporation. All rights reserved.
  SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
 #include <Uefi.h>
 #include <Library/UefiBootServicesTableLib.h>
+#include <Library/BaseMemoryLib.h>
 #include <Library/DebugLib.h>
+#include <Library/MemoryAllocationLib.h>
 #include <Library/RngLib.h>
 #include <Protocol/Rng.h>
 
+STATIC EFI_RNG_PROTOCOL  *mRngProtocol;
+STATIC UINTN             mFirstAlgo = MAX_UINTN;
+
+typedef struct {
+  /// Guid of the secure algorithm.
+  EFI_GUID       *Guid;
+
+  /// Algorithm name.
+  CONST CHAR8    *Name;
+
+  /// The algorithm is available for use.
+  BOOLEAN        Available;
+} SECURE_RNG_ALGO_ARRAY;
+
+//
+// These represent UEFI SPEC defined algorithms that should be supported by
+// the RNG protocol and are generally considered secure.
+//
+GLOBAL_REMOVE_IF_UNREFERENCED SECURE_RNG_ALGO_ARRAY  mSecureHashAlgorithms[] = {
+  {
+    &gEfiRngAlgorithmSp80090Ctr256Guid,  // SP800-90A DRBG CTR using AES-256
+    "DRBG-CTR",
+    FALSE,
+  },
+  {
+    &gEfiRngAlgorithmSp80090Hmac256Guid, // SP800-90A DRBG HMAC using SHA-256
+    "DRBG-HMAC",
+    FALSE,
+  },
+  {
+    &gEfiRngAlgorithmSp80090Hash256Guid, // SP800-90A DRBG Hash using SHA-256
+    "DRBG-Hash",
+    FALSE,
+  },
+  {
+    &gEfiRngAlgorithmRaw, // Raw data from NRBG (or TRNG)
+    "TRNG",
+    FALSE,
+  },
+};
+
+/**
+  Constructor routine to probe the available secure Rng algorithms.
+
+  @param  ImageHandle   The firmware allocated handle for the EFI image.
+  @param  SystemTable   A pointer to the EFI System Table.
+
+  @retval EFI_SUCCESS           Success.
+  @retval EFI_NOT_FOUND         Not found.
+  @retval EFI_INVALID_PARAMETER Invalid parameter.
+**/
+EFI_STATUS
+EFIAPI
+DxeRngLibConstructor (
+  IN EFI_HANDLE        ImageHandle,
+  IN EFI_SYSTEM_TABLE  *SystemTable
+  )
+{
+  EFI_STATUS         Status;
+  UINTN              RngArraySize;
+  UINTN              RngArrayCnt;
+  UINT32             Index;
+  UINT32             Index1;
+  EFI_RNG_ALGORITHM  *RngArray;
+
+  Status = gBS->LocateProtocol (&gEfiRngProtocolGuid, NULL, (VOID **)&mRngProtocol);
+  if (EFI_ERROR (Status) || (mRngProtocol == NULL)) {
+    DEBUG ((DEBUG_ERROR, "%a: Could not locate RNG protocol, Status = %r\n", __func__, Status));
+    return Status;
+  }
+
+  RngArraySize = 0;
+
+  Status = mRngProtocol->GetInfo (mRngProtocol, &RngArraySize, NULL);
+  if (EFI_ERROR (Status) && (Status != EFI_BUFFER_TOO_SMALL)) {
+    return Status;
+  } else if (RngArraySize == 0) {
+    return EFI_NOT_FOUND;
+  }
+
+  RngArrayCnt = RngArraySize / sizeof (*RngArray);
+
+  RngArray = AllocateZeroPool (RngArraySize);
+  if (RngArray == NULL) {
+    return EFI_OUT_OF_RESOURCES;
+  }
+
+  Status = mRngProtocol->GetInfo (mRngProtocol, &RngArraySize, RngArray);
+  if (EFI_ERROR (Status)) {
+    goto ExitHandler;
+  }
+
+  for (Index = 0; Index < RngArrayCnt; Index++) {
+    for (Index1 = 0; Index1 < ARRAY_SIZE (mSecureHashAlgorithms); Index1++) {
+      if (CompareGuid (&RngArray[Index], mSecureHashAlgorithms[Index1].Guid)) {
+        mSecureHashAlgorithms[Index1].Available = TRUE;
+        if (mFirstAlgo == MAX_UINTN) {
+          mFirstAlgo = Index1;
+        }
+
+        break;
+      }
+    }
+  }
+
+ExitHandler:
+  FreePool (RngArray);
+  return Status;
+}
+
 /**
   Routine Description:
 
@@ -32,55 +144,68 @@ GenerateRandomNumberViaNist800Algorithm (
   IN  UINTN  BufferSize
   )
 {
-  EFI_STATUS        Status;
+  EFI_STATUS             Status;
-  EFI_RNG_PROTOCOL  *RngProtocol;
+  UINTN                  Index;
-
+  SECURE_RNG_ALGO_ARRAY  *Algo;
-  RngProtocol = NULL;
 
   if (Buffer == NULL) {
     DEBUG ((DEBUG_ERROR, "%a: Buffer == NULL.\n", __func__));
     return EFI_INVALID_PARAMETER;
   }
 
-  Status = gBS->LocateProtocol (&gEfiRngProtocolGuid, NULL, (VOID **)&RngProtocol);
+  if (mRngProtocol == NULL) {
-  if (EFI_ERROR (Status) || (RngProtocol == NULL)) {
+    return EFI_NOT_FOUND;
-    DEBUG ((DEBUG_ERROR, "%a: Could not locate RNG prototocol, Status = %r\n", __func__, Status));
-    return Status;
   }
 
-  Status = RngProtocol->GetRNG (RngProtocol, &gEfiRngAlgorithmSp80090Ctr256Guid, BufferSize, Buffer);
+  // Try the first available algorithm.
-  DEBUG ((DEBUG_INFO, "%a: GetRNG algorithm CTR-256 - Status = %r\n", __func__, Status));
+  if (mFirstAlgo != MAX_UINTN) {
-  if (!EFI_ERROR (Status)) {
+    Algo   = &mSecureHashAlgorithms[mFirstAlgo];
-    return Status;
+    Status = mRngProtocol->GetRNG (mRngProtocol, Algo->Guid, BufferSize, Buffer);
+    DEBUG ((
+      DEBUG_INFO,
+      "%a: GetRNG algorithm %a - Status = %r\n",
+      __func__,
+      Algo->Name,
+      Status
+      ));
+    if (!EFI_ERROR (Status)) {
+      return Status;
+    }
+
+    Index = mFirstAlgo + 1;
+  } else {
+    Index = 0;
   }
 
-  Status = RngProtocol->GetRNG (RngProtocol, &gEfiRngAlgorithmSp80090Hmac256Guid, BufferSize, Buffer);
+  // Iterate over other available algorithms.
-  DEBUG ((DEBUG_INFO, "%a: GetRNG algorithm HMAC-256 - Status = %r\n", __func__, Status));
+  for ( ; Index < ARRAY_SIZE (mSecureHashAlgorithms); Index++) {
-  if (!EFI_ERROR (Status)) {
+    Algo = &mSecureHashAlgorithms[Index];
-    return Status;
+    if (!Algo->Available) {
-  }
+      continue;
+    }
 
-  Status = RngProtocol->GetRNG (RngProtocol, &gEfiRngAlgorithmSp80090Hash256Guid, BufferSize, Buffer);
+    Status = mRngProtocol->GetRNG (mRngProtocol, Algo->Guid, BufferSize, Buffer);
-  DEBUG ((DEBUG_INFO, "%a: GetRNG algorithm Hash-256 - Status = %r\n", __func__, Status));
+    DEBUG ((
-  if (!EFI_ERROR (Status)) {
+      DEBUG_INFO,
-    return Status;
+      "%a: GetRNG algorithm %a - Status = %r\n",
-  }
+      __func__,
-
+      Algo->Name,
-  Status = RngProtocol->GetRNG (RngProtocol, &gEfiRngAlgorithmRaw, BufferSize, Buffer);
+      Status
-  DEBUG ((DEBUG_INFO, "%a: GetRNG algorithm Raw - Status = %r\n", __func__, Status));
+      ));
-  if (!EFI_ERROR (Status)) {
+    if (!EFI_ERROR (Status)) {
-    return Status;
+      return Status;
+    }
   }
 
   // If all the other methods have failed, use the default method from the RngProtocol
-  Status = RngProtocol->GetRNG (RngProtocol, NULL, BufferSize, Buffer);
+  Status = mRngProtocol->GetRNG (mRngProtocol, NULL, BufferSize, Buffer);
   DEBUG ((DEBUG_INFO, "%a: GetRNG algorithm default - Status = %r\n", __func__, Status));
   if (!EFI_ERROR (Status)) {
     return Status;
   }
 
   // If we get to this point, we have failed
-  DEBUG ((DEBUG_ERROR, "%a: GetRNG() failed, staus = %r\n", __func__, Status));
+  DEBUG ((DEBUG_ERROR, "%a: GetRNG() failed, Status = %r\n", __func__, Status));
 
   return Status;
 }// GenerateRandomNumberViaNist800Algorithm()

MdePkg/Library/DxeRngLib/DxeRngLib.inf

@@ -15,6 +15,7 @@
   MODULE_TYPE     = DXE_DRIVER
   VERSION_STRING  = 1.0
   LIBRARY_CLASS   = RngLib|DXE_DRIVER DXE_RUNTIME_DRIVER UEFI_APPLICATION UEFI_DRIVER
+  CONSTRUCTOR     = DxeRngLibConstructor
 
 [Packages]
   MdePkg/MdePkg.dec
@@ -24,6 +25,7 @@
 
 [LibraryClasses]
   DebugLib
+  MemoryAllocationLib
   UefiBootServicesTableLib
 
 [Protocols]