From 2afd172dd50b09641ea203755d13bb5d35dd0a45 Mon Sep 17 00:00:00 2001
From: Vu Nguyen <Vu@amperecomputing.com>
Date: Thu, 21 Apr 2022 13:59:06 +0700
Subject: [PATCH] AmpereAltraPkg: Fix failed to update firmware via Capsule

Capsule update is relying on the gEfiMmCommunication2Protocol, meanwhile
some old firmwares might just expose gEfiMmCommunicationProtocol. This
change is to address this issue by trying to locate both protocols and
doing the update via the first found protocol.

Signed-off-by: Vu Nguyen <Vu@amperecomputing.com>
---
 .../PlatformFlashAccessLib.c                  | 43 +++++++++++++------
 .../PlatformFlashAccessLib.inf                |  4 +-
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/Silicon/Ampere/AmpereAltraPkg/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.c b/Silicon/Ampere/AmpereAltraPkg/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.c
index a9664c518a1..1662e3cecd5 100644
--- a/Silicon/Ampere/AmpereAltraPkg/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.c
+++ b/Silicon/Ampere/AmpereAltraPkg/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.c
@@ -16,7 +16,8 @@
 #include <Library/UefiBootServicesTableLib.h>
 #include <Protocol/MmCommunication2.h>
 
-EFI_MM_COMMUNICATION2_PROTOCOL *mMmCommunication = NULL;
+EFI_MM_COMMUNICATION2_PROTOCOL *mMmCommunication2 = NULL;
+EFI_MM_COMMUNICATION_PROTOCOL  *mMmCommunication  = NULL;
 
 #define EFI_MM_MAX_PAYLOAD_U64_E 10
 #define EFI_MM_MAX_PAYLOAD_SIZE  (EFI_MM_MAX_PAYLOAD_U64_E * sizeof (UINT64))
@@ -127,16 +128,23 @@ MmFlashUpdate (
     ImageSize
     ));
 
-  if (mMmCommunication == NULL) {
+  if (mMmCommunication2 == NULL && mMmCommunication == NULL) {
     Status = gBS->LocateProtocol (
                     &gEfiMmCommunication2ProtocolGuid,
                     NULL,
-                    (VOID **)&mMmCommunication
+                    (VOID **)&mMmCommunication2
                     );
-
     if (EFI_ERROR (Status)) {
-      DEBUG ((DEBUG_ERROR, "%a: Can't locate gEfiMmCommunication2ProtocolGuid.\n", __FUNCTION__));
-      return Status;
+      Status = gBS->LocateProtocol (
+                      &gEfiMmCommunicationProtocolGuid,
+                      NULL,
+                      (VOID **)&mMmCommunication
+                      );
+
+      if (EFI_ERROR (Status)) {
+        DEBUG ((DEBUG_ERROR, "%a: Can't locate MM Communication protocol.\n", __FUNCTION__));
+        return Status;
+      }
     }
   }
 
@@ -149,12 +157,23 @@ MmFlashUpdate (
     UefiMmCreateSysFwuReq ((VOID *)&MmData, sizeof (MmData));
 
     Size = sizeof (EFI_MM_COMM_HEADER_NOPAYLOAD) + sizeof (MmData);
-    Status = mMmCommunication->Communicate (
-                                 mMmCommunication,
-                                 (VOID *)&mEfiMmSysFwuReq,
-                                 (VOID *)&mEfiMmSysFwuReq,
-                                 &Size
-                                 );
+    if (mMmCommunication2 != NULL) {
+      Status = mMmCommunication2->Communicate (
+                                    mMmCommunication2,
+                                    (VOID *)&mEfiMmSysFwuReq,
+                                    (VOID *)&mEfiMmSysFwuReq,
+                                    &Size
+                                    );
+    } else if (mMmCommunication != NULL) {
+      Status =  mMmCommunication->Communicate (
+                                    mMmCommunication,
+                                    (VOID *)&mEfiMmSysFwuReq,
+                                    &Size
+                                    );
+    } else {
+      return EFI_UNSUPPORTED;
+    }
+
     if (EFI_ERROR (Status)) {
       DEBUG ((
         DEBUG_ERROR,
diff --git a/Silicon/Ampere/AmpereAltraPkg/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.inf b/Silicon/Ampere/AmpereAltraPkg/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.inf
index fb147b6da03..fddbbd8a1f7 100644
--- a/Silicon/Ampere/AmpereAltraPkg/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.inf
+++ b/Silicon/Ampere/AmpereAltraPkg/Library/PlatformFlashAccessLib/PlatformFlashAccessLib.inf
@@ -32,6 +32,8 @@
 
 [Protocols]
   gEfiMmCommunication2ProtocolGuid
+  gEfiMmCommunicationProtocolGuid
 
 [Depex]
-  TRUE
+  gEfiMmCommunication2ProtocolGuid OR
+  gEfiMmCommunicationProtocolGuid