From 94463e7bc121e983a133927b9d05ad508efdc21f Mon Sep 17 00:00:00 2001
From: Eeshan Londhe <eeshanlondhe@microsoft.com>
Date: Tue, 15 Sep 2026 12:11:03 -0700
Subject: [PATCH] UefiTestingPkg: Correctly handle cached DMA audit ACPI tables
 (#938)

## Description

Correctly handle cached DMA audit ACPI tables. Fixes
https://github.com/microsoft/mu_plus/issues/937

Tests call GetDmarAcpiTable(), but the Intel VT-d and IVRS
implementation caches the ACPI table and returns EFI_ALREADY_STARTED on
the second call without a reboot.

GetAcpiTable() will no longer return already started and will always
refresh the pointer with the data. On the caller side, the caching of
the global variable is handled so we don't call GetAcpiTable()
redundantly.

- [ ] Impacts functionality?
- [ ] Impacts security?
- [ ] Breaking change?
- [ ] Includes tests?
- [ ] Includes documentation?

## How This Was Tested

Tested on Intel platform and all tests now pass.

## Integration Instructions

N/A

(cherry picked from commit a8b711bf1ba37d14437b9779f4ddd24135a08d8f)
---
 .../AuditTests/DMAProtectionAudit/UEFI/Acpi.c         | 11 ++++-------
 .../AuditTests/DMAProtectionAudit/UEFI/Acpi.h         |  5 +++--
 .../DMAProtectionAudit/UEFI/IVRS/IVRSAcpiTable.c      |  4 ++++
 .../DMAProtectionAudit/UEFI/VTd/DmarAcpiTable.c       |  4 ++++
 4 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/Acpi.c b/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/Acpi.c
index 5b41d3a89f..c38376712d 100644
--- a/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/Acpi.c
+++ b/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/Acpi.c
@@ -142,8 +142,9 @@ FindAcpiPtr (
 /**
   Get the ACPI table.
 
-  @retval EFI_SUCCESS           The ACPI table is got.
-  @retval EFI_NOT_FOUND         The ACPI table is not found.
+  @retval EFI_SUCCESS            The ACPI table is found.
+  @retval EFI_INVALID_PARAMETER  The ACPI Table is NULL.
+  @retval EFI_NOT_FOUND          The ACPI table is not found.
 **/
 EFI_STATUS
 GetAcpiTable (
@@ -156,8 +157,6 @@ GetAcpiTable (
 
   if (AcpiTable == NULL) {
     return EFI_INVALID_PARAMETER;
-  } else if (*AcpiTable != NULL) {
-    return EFI_ALREADY_STARTED;
   }
 
   *AcpiTable = NULL;
@@ -172,12 +171,10 @@ GetAcpiTable (
                );
   }
 
-  if (EFI_ERROR (Status)) {
+  if (EFI_ERROR (Status) || (AcpiConfigurationTable == NULL)) {
     return EFI_NOT_FOUND;
   }
 
-  ASSERT (AcpiConfigurationTable != NULL);
-
   *AcpiTable = FindAcpiPtr (
                  (EFI_ACPI_2_0_ROOT_SYSTEM_DESCRIPTION_POINTER *)AcpiConfigurationTable,
                  AcpiSignature
diff --git a/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/Acpi.h b/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/Acpi.h
index 8bec31dcae..f7166bd04e 100644
--- a/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/Acpi.h
+++ b/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/Acpi.h
@@ -30,8 +30,9 @@ typedef struct {
 /**
   Get the ACPI table.
 
-  @retval EFI_SUCCESS           The ACPI table is got.
-  @retval EFI_NOT_FOUND         The ACPI table is not found.
+  @retval EFI_SUCCESS            The ACPI table is found.
+  @retval EFI_INVALID_PARAMETER  The ACPI Table is NULL.
+  @retval EFI_NOT_FOUND          The ACPI table is not found.
 **/
 EFI_STATUS
 GetAcpiTable (
diff --git a/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/IVRS/IVRSAcpiTable.c b/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/IVRS/IVRSAcpiTable.c
index 8495647c56..bd5ccd2b62 100644
--- a/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/IVRS/IVRSAcpiTable.c
+++ b/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/IVRS/IVRSAcpiTable.c
@@ -259,5 +259,9 @@ GetIvrsAcpiTable (
   VOID
   )
 {
+  if (mAcpiIVRSTable != NULL) {
+    return EFI_SUCCESS;
+  }
+
   return GetAcpiTable (IVRS_HEADER_SIGNATURE, (VOID **)&mAcpiIVRSTable);
 }
diff --git a/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/VTd/DmarAcpiTable.c b/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/VTd/DmarAcpiTable.c
index ea572b9a89..1d78ae6c51 100644
--- a/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/VTd/DmarAcpiTable.c
+++ b/UefiTestingPkg/AuditTests/DMAProtectionAudit/UEFI/VTd/DmarAcpiTable.c
@@ -325,5 +325,9 @@ GetDmarAcpiTable (
   VOID
   )
 {
+  if (mAcpiDmarTable != NULL) {
+    return EFI_SUCCESS;
+  }
+
   return GetAcpiTable (EFI_ACPI_4_0_DMA_REMAPPING_TABLE_SIGNATURE, (VOID **)&mAcpiDmarTable);
 }