5 files changed, 20 insertions, 3 deletions
diff --git a/MdeModulePkg/Universal/CapsulePei/Capsule.h b/MdeModulePkg/Universal/CapsulePei/Capsule.h
index 57af114c40..68d8e42745 100644
--- a/MdeModulePkg/Universal/CapsulePei/Capsule.h
+++ b/MdeModulePkg/Universal/CapsulePei/Capsule.h
@@ -1,6 +1,6 @@
 /** @file
 
-Copyright (c) 2006 - 2011, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2006 - 2012, Intel Corporation. All rights reserved.<BR>
 
 This program and the accompanying materials
 are licensed and made available under the terms and conditions
@@ -35,6 +35,7 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/PeCoffGetEntryPointLib.h>
 #include <Library/PcdLib.h>
 #include <Library/ReportStatusCodeLib.h>
+#include <Library/DebugAgentLib.h>
 #include <IndustryStandard/PeImage.h>
 #include "Common/CommonHeader.h"
 
diff --git a/MdeModulePkg/Universal/CapsulePei/CapsulePei.inf b/MdeModulePkg/Universal/CapsulePei/CapsulePei.inf
index 8f96d9ec8a..5ca9afa5b5 100644
--- a/MdeModulePkg/Universal/CapsulePei/CapsulePei.inf
+++ b/MdeModulePkg/Universal/CapsulePei/CapsulePei.inf
@@ -59,6 +59,7 @@
 [LibraryClasses.IA32]
   PeCoffGetEntryPointLib
   PcdLib
+  DebugAgentLib
 
 [Guids]
   gEfiCapsuleVendorGuid                         # ALWAYS_CONSUMED
diff --git a/MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf b/MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf
index 1c155ce5f3..6af8cbfc47 100644
--- a/MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf
+++ b/MdeModulePkg/Universal/CapsulePei/CapsuleX64.inf
@@ -47,6 +47,7 @@
   BaseLib
   DebugLib
   CpuExceptionHandlerLib
+  DebugAgentLib
 
 [Depex]
   FALSE
diff --git a/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c b/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
index 98b0d56457..f5dd85b2b6 100644
--- a/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
+++ b/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
@@ -1,7 +1,7 @@
 /** @file
   Capsule update PEIM for UEFI2.0
 
-Copyright (c) 2006 - 2011, Intel Corporation. All rights reserved.<BR>
+Copyright (c) 2006 - 2012, Intel Corporation. All rights reserved.<BR>
 
 This program and the accompanying materials
 are licensed and made available under the terms and conditions
@@ -350,6 +350,10 @@ Thunk32To64 (
     AsmWriteCr3 ((UINTN) PageTableAddress);
 
     //
+    // Disable interrupt of Debug timer, since the IDT table cannot work in long mode
+    //
+    SaveAndSetDebugTimerInterrupt (FALSE);
+    //
     // Transfer to long mode
     //
     AsmEnablePaging64 (
diff --git a/MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c b/MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c
index fa29e3dbfb..56913e6b03 100644
--- a/MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c
+++ b/MdeModulePkg/Universal/CapsulePei/X64/X64Entry.c
@@ -15,9 +15,10 @@ WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
 #include <Library/DebugLib.h>
 #include <Library/BaseMemoryLib.h>
 #include <Library/CpuExceptionHandlerLib.h>
+#include <Library/DebugAgentLib.h>
 #include "CommonHeader.h"
 
-#define EXCEPTION_VECTOR_NUMBER     0x20
+#define EXCEPTION_VECTOR_NUMBER     0x22
 
 /**
   The X64 entrypoint is used to process capsule in long mode then
@@ -58,6 +59,11 @@ _ModuleEntryPoint (
   // Setup the default CPU exception handlers
   //
   SetupCpuExceptionHandlers ();                
+  
+  //
+  // Initialize Debug Agent to support source level debug
+  //
+  InitializeDebugAgent (DEBUG_AGENT_INIT_THUNK_PEI_IA32TOX64, (VOID *) &Ia32Idtr, NULL);
 
   //
   // Call CapsuleDataCoalesce to process capsule.
@@ -72,6 +78,10 @@ _ModuleEntryPoint (
   ReturnContext->ReturnStatus = Status;
 
   //
+  // Disable interrupt of Debug timer, since the new IDT table cannot work in long mode
+  //
+  SaveAndSetDebugTimerInterrupt (FALSE);
+  //
   // Restore IA32 IDT table
   //
   AsmWriteIdtr ((IA32_DESCRIPTOR *) &Ia32Idtr);