Fix up the MmuLib logic as we move to use generic interface (#191)

# Preface

Please ensure you have read the [contribution
docs](https://github.com/microsoft/mu/blob/master/CONTRIBUTING.md) prior
to submitting the pull request. In particular,
[pull request
guidelines](https://github.com/microsoft/mu/blob/master/CONTRIBUTING.md#pull-request-best-practices).

## Description

Upstream change removed declaration of specific individual interfaces
and moved to use the common interface `ArmSetMemoryAttributes`. This
change accommodates the MU_CHANGE with needed functions inside MmuLib.

For each item, place an "x" in between `[` and `]` if true. Example:
`[x]`.
_(you can also check items in the GitHub UI)_

- [x] Impacts functionality?
- **Functionality** - Does the change ultimately impact how firmware
functions?
- Examples: Add a new library, publish a new PPI, update an algorithm,
...
- [ ] 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, ...

## How This Was Tested

Tested bootable on QEMU SBSA.

## Integration Instructions

N/A

ArmPkg/Library/MmuLib/MmuLib.c

@@ -41,25 +41,22 @@ MmuSetAttributes (
 
   Status = EFI_UNSUPPORTED;
 
-  if (Attributes & EFI_MEMORY_XP) {
-    Status = ArmSetMemoryAttributes (BaseAddress, Length, EFI_MEMORY_XP, EFI_MEMORY_XP);
-    if (EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_ERROR, "%a - Failed to set NX.  Status = %r\n", __FUNCTION__, Status));
-    }
-
-    Attributes &= ~EFI_MEMORY_XP;
-  }
-
-  if (Attributes & EFI_MEMORY_RO) {
-    Status = ArmSetMemoryAttributes (BaseAddress, Length, EFI_MEMORY_RO, EFI_MEMORY_RO);
-    if (EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_ERROR, "%a - Failed to set RO.  Status = %r\n", __FUNCTION__, Status));
-    }
-
-    Attributes &= ~EFI_MEMORY_RO;
+  // MU_CHANGE - START
+  // Ensure that the attributes are valid
+  ASSERT ((Attributes & ~(EFI_MEMORY_XP | EFI_MEMORY_RO)) == 0);
+
+  // Use ArmSetMemoryAttributes because the individually called attribute updates have been removed
+  Status = ArmSetMemoryAttributes (
+              BaseAddress,
+              Length,
+              Attributes,
+              (EFI_MEMORY_XP | EFI_MEMORY_RO)
+              );
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a - Failed to clear attributes - 0x%llx.  Status = %r\n", __FUNCTION__, Attributes, Status));
   }
 
-  ASSERT (Attributes == 0);
+  // MU_CHANGE - END
   ASSERT_EFI_ERROR (Status);
   return Status;
 }
@@ -88,41 +85,25 @@ MmuClearAttributes (
 
   Status = EFI_UNSUPPORTED;
 
-  if (Attributes & EFI_MEMORY_XP) {
-    // MU_CHANGE - START
-    // Use ArmSetMemoryAttributes because the individually called attribute updates have been removed
-    Status = ArmSetMemoryAttributes (
-               BaseAddress,
-               Length,
-               0,
-               Attributes
-               );
-    // MU_CHANGE - END
-    if (EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_ERROR, "%a - Failed to clear NX.  Status = %r\n", __FUNCTION__, Status));
-    }
-
-    Attributes &= ~EFI_MEMORY_XP;
-  }
-
-  if (Attributes & EFI_MEMORY_RO) {
-    // MU_CHANGE - START
-    // Use ArmSetMemoryAttributes because the individually called attribute updates have been removed
-    Status = ArmSetMemoryAttributes (
-               BaseAddress,
-               Length,
-               0,
-               Attributes
-               );
-    // MU_CHANGE - END
-    if (EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_ERROR, "%a - Failed to clear RO.  Status = %r\n", __FUNCTION__, Status));
-    }
-
-    Attributes &= ~EFI_MEMORY_RO;
+  // MU_CHANGE - START
+  // Ensure that the attributes are valid
+  ASSERT ((Attributes & ~(EFI_MEMORY_XP | EFI_MEMORY_RO)) == 0);
+
+  // As we clear the attributes, we need to "set" the inverse of the attributes
+  Attributes = (~Attributes) & (EFI_MEMORY_XP | EFI_MEMORY_RO);
+
+  // Use ArmSetMemoryAttributes because the individually called attribute updates have been removed
+  Status = ArmSetMemoryAttributes (
+              BaseAddress,
+              Length,
+              Attributes,
+              (EFI_MEMORY_XP | EFI_MEMORY_RO)
+              );
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a - Failed to clear attributes - 0x%llx.  Status = %r\n", __FUNCTION__, Attributes, Status));
   }
 
-  ASSERT (Attributes == 0);
+  // MU_CHANGE - END
   ASSERT_EFI_ERROR (Status);
   return Status;
 }

ArmPkg/Library/StandaloneMmMmuLib/ArmMmuStandaloneMmLib.c

@@ -418,48 +418,65 @@ ArmSetMemoryAttributes (
 {
   // MU_CHANGE [START] - Add ArmSetMemoryAttributes functionality
   EFI_STATUS  Status;
+  UINT64      NeededAttributes;
 
   DEBUG ((
     DEBUG_INFO,
-    "%a: BaseAddress == 0x%lx, Length == 0x%lx, Attributes == 0x%lx\n",
+    "%a: BaseAddress == 0x%llx, Length == 0x%llx, Attributes == 0x%llx, Mask == 0x%llx\n",
     __FUNCTION__,
-    (UINTN)BaseAddress,
-    (UINTN)Length,
-    (UINTN)Attributes
+    BaseAddress,
+    Length,
+    Attributes,
+    AttributeMask
     ));
 
+  NeededAttributes = Attributes & AttributeMask;
+
   if ((Length == 0) ||
-      ((Attributes & ~(EFI_MEMORY_RO | EFI_MEMORY_RP | EFI_MEMORY_XP)) != 0))
+      ((NeededAttributes & ~(EFI_MEMORY_RO | EFI_MEMORY_RP | EFI_MEMORY_XP)) != 0))
   {
-    return EFI_INVALID_PARAMETER;
-  }
-
-  if (!RegionIsSystemMemory (BaseAddress, Length)) {
-    return EFI_UNSUPPORTED;
+    Status = EFI_INVALID_PARAMETER;
+    goto Done;
   }
 
-  if ((Attributes & EFI_MEMORY_RP) != 0) {
+  if ((NeededAttributes & EFI_MEMORY_RP) != 0) {
     Status = ArmSetMemoryRegionNoAccess (BaseAddress, Length);
     if (EFI_ERROR (Status)) {
-      return EFI_UNSUPPORTED;
+      goto Done;
+    }
+  } else {
+    Status = ArmClearMemoryRegionNoAccess (BaseAddress, Length);
+    if (EFI_ERROR (Status)) {
+      goto Done;
     }
   }
 
-  if ((Attributes & EFI_MEMORY_RO) != 0) {
+  if ((NeededAttributes & EFI_MEMORY_RO) != 0) {
     Status = ArmSetMemoryRegionReadOnly (BaseAddress, Length);
     if (EFI_ERROR (Status)) {
-      return EFI_UNSUPPORTED;
+      goto Done;
+    }
+  } else {
+    Status = ArmClearMemoryRegionReadOnly (BaseAddress, Length);
+    if (EFI_ERROR (Status)) {
+      goto Done;
     }
   }
 
-  if ((Attributes & EFI_MEMORY_XP) != 0) {
+  if ((NeededAttributes & EFI_MEMORY_XP) != 0) {
     Status = ArmSetMemoryRegionNoExec (BaseAddress, Length);
     if (EFI_ERROR (Status)) {
-      return EFI_UNSUPPORTED;
+      goto Done;
+    }
+  } else {
+    Status = ArmClearMemoryRegionNoExec (BaseAddress, Length);
+    if (EFI_ERROR (Status)) {
+      goto Done;
     }
   }
 
-  return EFI_SUCCESS;
+Done:
+  return Status;
   // MU_CHANGE [END] - Add ArmSetMemoryAttributes functionality
 }