1 files changed, 26 insertions, 10 deletions
diff --git a/SourceLevelDebugPkg/Library/DebugAgent/SecPeiDebugAgent/SecPeiDebugAgentLib.c b/SourceLevelDebugPkg/Library/DebugAgent/SecPeiDebugAgent/SecPeiDebugAgentLib.c
index 828407304f..0132942dea 100644
--- a/SourceLevelDebugPkg/Library/DebugAgent/SecPeiDebugAgent/SecPeiDebugAgentLib.c
+++ b/SourceLevelDebugPkg/Library/DebugAgent/SecPeiDebugAgent/SecPeiDebugAgentLib.c
@@ -174,19 +174,22 @@ GetMailboxPointer (
   MailboxLocationInIdt = GetLocationSavedMailboxPointerInIdtEntry ();
   Mailbox = (DEBUG_AGENT_MAILBOX *)(UINTN)(*MailboxLocationInIdt);
   //
-  // Check if mailbox was setup in PEI firstly, cannot used GetDebugFlag() to 
-  // get CheckMailboxInHob flag to avoid GetMailboxPointer() nesting.
+  // Cannot used GetDebugFlag() to get Debug Flag to avoid GetMailboxPointer() nested
   //
-  if (Mailbox->DebugFlag.Bits.CheckMailboxInHob != 1) {
-    //
-    // If mailbox in IDT entry has already been the final one
+  if (Mailbox->DebugFlag.Bits.CheckMailboxInHob != 1 ||
+      Mailbox->DebugFlag.Bits.InitArch != DEBUG_ARCH_SYMBOL) {
     //
+    // If mailbox was setup in SEC or the current CPU arch is different from the init arch
+    // Debug Agent initialized, return the mailbox from IDT entry directly.
+    // Otherwise, we need to check the mailbox location saved in GUIDed HOB further.
+    // 
     return Mailbox;
   }
 
   MailboxLocationInHob = GetMailboxLocationFromHob ();
   //
-  // Compare mailbox in IDT enry with mailbox in HOB
+  // Compare mailbox in IDT enry with mailbox in HOB,
+  // need to fix mailbox location if HOB moved by PEI CORE
   //
   if (MailboxLocationInHob != MailboxLocationInIdt && MailboxLocationInHob != NULL) {
     Mailbox = (DEBUG_AGENT_MAILBOX *)(UINTN)(*MailboxLocationInHob);
@@ -247,7 +250,7 @@ DebugAgentCallbackMemoryDiscoveredPpi (
   EFI_STATUS                     Status;
   DEBUG_AGENT_MAILBOX            *Mailbox;
   BOOLEAN                        InterruptStatus;
-  EFI_PHYSICAL_ADDRESS           Memory; 
+  EFI_PHYSICAL_ADDRESS           Address; 
   DEBUG_AGENT_MAILBOX            *NewMailbox;
   UINT64                         *MailboxLocationInHob;
 
@@ -262,10 +265,10 @@ DebugAgentCallbackMemoryDiscoveredPpi (
   Status = PeiServicesAllocatePages (
              EfiACPIMemoryNVS,
              EFI_SIZE_TO_PAGES (sizeof(DEBUG_AGENT_MAILBOX) + PcdGet16(PcdDebugPortHandleBufferSize)),
-             &Memory
+             &Address
              );
   ASSERT_EFI_ERROR (Status);
-  NewMailbox = (DEBUG_AGENT_MAILBOX *) (UINTN) Memory;
+  NewMailbox = (DEBUG_AGENT_MAILBOX *) (UINTN) Address;
   //
   // Copy Mailbox and Debug Port Handle buffer to new location in ACPI NVS memory, because original Mailbox
   // and Debug Port Handle buffer in the allocated pool that may be marked as free by DXE Core after DXE Core
@@ -369,7 +372,14 @@ InitializeDebugAgent (
     // Get and save debug port handle and set the length of memory block.
     //
     SetLocationSavedMailboxPointerInIdtEntry (&MailboxLocation);
+    //
+    // Force error message could be printed during the first shakehand between Target/HOST.
+    //
     SetDebugFlag (DEBUG_AGENT_FLAG_PRINT_ERROR_LEVEL, DEBUG_AGENT_ERROR);
+    //
+    // Save init arch type when debug agent initialized
+    //
+    SetDebugFlag (DEBUG_AGENT_FLAG_INIT_ARCH, sizeof (UINTN) / 4);
 
     InitializeDebugTimer ();
 
@@ -453,6 +463,10 @@ InitializeDebugAgent (
     //
     SetLocationSavedMailboxPointerInIdtEntry (MailboxLocationPointer);
     //
+    // Save init arch type when debug agent initialized
+    //
+    SetDebugFlag (DEBUG_AGENT_FLAG_INIT_ARCH, DEBUG_ARCH_SYMBOL);
+    //
     // Register for a callback once memory has been initialized.
     // If memery has been ready, the callback funtion will be invoked immediately
     //
@@ -483,6 +497,9 @@ InitializeDebugAgent (
       MailboxLocationPointer = (UINT64 *) (UINTN) (Ia32IdtEntry[DEBUG_MAILBOX_VECTOR].Bits.OffsetLow +
                                                 (Ia32IdtEntry[DEBUG_MAILBOX_VECTOR].Bits.OffsetHigh << 16));
       Mailbox = (DEBUG_AGENT_MAILBOX *) (UINTN)(*MailboxLocationPointer);
+      //
+      // Mailbox should valid and setup before executing thunk code
+      //
       VerifyMailboxChecksum (Mailbox);
 
       DebugPortHandle = (UINT64) (UINTN)DebugPortInitialize ((VOID *)(UINTN)Mailbox->DebugPortHandle, NULL);
@@ -508,7 +525,6 @@ InitializeDebugAgent (
     DEBUG ((EFI_D_ERROR, "Debug Agent: The InitFlag value is not allowed!\n"));
     CpuDeadLoop ();
     break;
-
   }
 
   EnableInterrupts ();