From 7005df708610b0fea9d0a06c4f3e654d5bd90ab7 Mon Sep 17 00:00:00 2001
From: Taylor Beebe <31827475+TaylorBeebe@users.noreply.github.com>
Date: Mon, 18 Sep 2023 15:07:53 -0700
Subject: [PATCH] MdeModulePkg: Don't Lock Variable Policy if Device is in Unit
 Test Mode

To enable testing variable policy functionality in the shell, it needs
to be unlocked when the test starts. This PR adds a check to the system
device state and avoids locking if it is in unit test mode.

- [x] Impacts functionality?
- **Functionality** - Does the change ultimately impact how firmware
functions?
- Examples: Add a new library, publish a new PPI, update an algorithm,
...
- [x] Impacts security?
- **Security** - Does the change have a direct security impact on an
application,
    flow, or firmware?
  - Examples: Crypto algorithm change, buffer overflow fix, parameter
    validation improvement, ...
- [ ] Breaking change?
- **Breaking change** - Will anyone consuming this change experience a
break
    in build or boot behavior?
- Examples: Add a new library class, move a module to a different repo,
call
    a function in a new library class in a pre-existing module, ...
- [ ] Includes tests?
  - **Tests** - Does the change include any explicit test code?
  - Examples: Unit tests, integration tests, robot tests, ...
- [ ] Includes documentation?
- **Documentation** - Does the change contain explicit documentation
additions
    outside direct code modifications (and comments)?
- Examples: Update readme file, add feature readme file, link to
documentation
    on an a separate Web page, ...

Tested by running the variable policy test app in Q35.

N/A
---
 MdeModulePkg/MdeModulePkg.dsc                            | 7 +++++--
 .../Variable/RuntimeDxe/VariablePolicyLockingCommon.c    | 9 +++++++++
 .../Universal/Variable/RuntimeDxe/VariableRuntimeDxe.inf | 1 +
 .../Variable/RuntimeDxe/VariableSmmRuntimeDxe.inf        | 1 +
 4 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/MdeModulePkg/MdeModulePkg.dsc b/MdeModulePkg/MdeModulePkg.dsc
index 3fce2695c31..e94e08ae6bc 100644
--- a/MdeModulePkg/MdeModulePkg.dsc
+++ b/MdeModulePkg/MdeModulePkg.dsc
@@ -112,13 +112,16 @@
   MmUnblockMemoryLib|MdePkg/Library/MmUnblockMemoryLib/MmUnblockMemoryLibNull.inf
   VariableFlashInfoLib|MdeModulePkg/Library/BaseVariableFlashInfoLib/BaseVariableFlashInfoLib.inf
   IpmiCommandLib|MdeModulePkg/Library/BaseIpmiCommandLibNull/BaseIpmiCommandLibNull.inf
-  SpiHcPlatformLib|MdeModulePkg/Library/BaseSpiHcPlatformLibNull/BaseSpiHcPlatformLibNull.inf
-   MemoryTypeInfoSecVarCheckLib|MdeModulePkg/Library/MemoryTypeInfoSecVarCheckLib/MemoryTypeInfoSecVarCheckLib.inf     # MU_CHANGE TCBZ1086
 
   AdvLoggerAccessLib|MdeModulePkg/Library/AdvLoggerAccessLibNull/AdvLoggerAccessLib.inf       ## MU_CHANGE
+  AdvLoggerAccessLib|MdeModulePkg/Library/AdvLoggerAccessLibNull/AdvLoggerAccessLib.inf                               ## MS_CHANGE
+  MemoryTypeInfoSecVarCheckLib|MdeModulePkg/Library/MemoryTypeInfoSecVarCheckLib/MemoryTypeInfoSecVarCheckLib.inf     # MU_CHANGE TCBZ1086
+  DeviceStateLib|MdeModulePkg/Library/DeviceStateLib/DeviceStateLib.inf                                            # MU_CHANGE
+
   PanicLib|MdePkg/Library/BasePanicLibNull/BasePanicLibNull.inf  # MU_CHANGE
 
   NULL|MdePkg/Library/StackCheckLibNull/StackCheckLibNull.inf # MU_CHANGE: /GS and -fstack-protector support
+
 [LibraryClasses.EBC.PEIM]
   IoLib|MdePkg/Library/PeiIoLibCpuIo/PeiIoLibCpuIo.inf
 
diff --git a/MdeModulePkg/Universal/Variable/RuntimeDxe/VariablePolicyLockingCommon.c b/MdeModulePkg/Universal/Variable/RuntimeDxe/VariablePolicyLockingCommon.c
index 1043b03c076..ec742333c20 100644
--- a/MdeModulePkg/Universal/Variable/RuntimeDxe/VariablePolicyLockingCommon.c
+++ b/MdeModulePkg/Universal/Variable/RuntimeDxe/VariablePolicyLockingCommon.c
@@ -10,6 +10,7 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 #include <Uefi.h>
 #include <Library/UefiLib.h>
 #include <Library/DebugLib.h>
+#include <Library/DeviceStateLib.h>
 #include <Library/UefiBootServicesTableLib.h>
 
 #include <Protocol/VariablePolicy.h>
@@ -61,6 +62,14 @@ LockPolicyInterfaceAtReadyToBoot (
 {
   EFI_STATUS  Status;
 
+  DEBUG_CODE_BEGIN ();
+  if ((GetDeviceState () & DEVICE_STATE_UNIT_TEST_MODE) != 0) {
+    DEBUG ((DEBUG_INFO, "[%a] Unit test mode is enabled -- skipping variable policy lock.\n", __FUNCTION__));
+    return;
+  }
+
+  DEBUG_CODE_END ();
+
   if (mCallbackInterface != NULL) {
     DEBUG ((DEBUG_INFO, "[%a] Invoking pre-lock callback.\n", __FUNCTION__));
     Status = mCallbackInterface->PreLock (mVariablePolicy);
diff --git a/MdeModulePkg/Universal/Variable/RuntimeDxe/VariableRuntimeDxe.inf b/MdeModulePkg/Universal/Variable/RuntimeDxe/VariableRuntimeDxe.inf
index 8baa051f1a5..0586d7906ef 100644
--- a/MdeModulePkg/Universal/Variable/RuntimeDxe/VariableRuntimeDxe.inf
+++ b/MdeModulePkg/Universal/Variable/RuntimeDxe/VariableRuntimeDxe.inf
@@ -78,6 +78,7 @@
   VariablePolicyHelperLib
   SafeIntLib
   MemoryTypeInfoSecVarCheckLib  # MU_CHANGE TCBZ1086 - Mitigate potential system brick due to UEFI MemoryTypeInformation var changes
+  DeviceStateLib  # MU_CHANGE - Check device state before locking variable policy
 
 [Protocols]
   gEfiFirmwareVolumeBlockProtocolGuid           ## CONSUMES
diff --git a/MdeModulePkg/Universal/Variable/RuntimeDxe/VariableSmmRuntimeDxe.inf b/MdeModulePkg/Universal/Variable/RuntimeDxe/VariableSmmRuntimeDxe.inf
index 5cc42594535..3c02ded42d1 100644
--- a/MdeModulePkg/Universal/Variable/RuntimeDxe/VariableSmmRuntimeDxe.inf
+++ b/MdeModulePkg/Universal/Variable/RuntimeDxe/VariableSmmRuntimeDxe.inf
@@ -63,6 +63,7 @@
   SafeIntLib
   PcdLib
   MmUnblockMemoryLib
+  DeviceStateLib # MU_CHANGE - Check device state before locking variable policy
 
 [Protocols]
   gEfiVariableWriteArchProtocolGuid             ## PRODUCES