Restore original IDT entry if RegisterInterruptHandler() was used to unregister user...

[mirror_edk2.git] / UefiCpuPkg / CpuDxe / CpuDxe.c

diff --git a/UefiCpuPkg/CpuDxe/CpuDxe.c b/UefiCpuPkg/CpuDxe/CpuDxe.c
index fab4ad23cbdee5d804a1fdec417cb7bd862c0f34..89dd2df5be7a87112da436bf8345019d39a1c523 100644 (file)
--- a/UefiCpuPkg/CpuDxe/CpuDxe.c
+++ b/UefiCpuPkg/CpuDxe/CpuDxe.c
@@ -25,6 +25,8 @@ EFI_HANDLE                mCpuHandle = NULL;
 BOOLEAN                   mIsFlushingGCD;\r
 UINT64                    mValidMtrrAddressMask = MTRR_LIB_CACHE_VALID_ADDRESS;\r
 UINT64                    mValidMtrrBitsMask    = MTRR_LIB_MSR_VALID_MASK;\r
+IA32_IDT_GATE_DESCRIPTOR  *mOrigIdtEntry        = NULL;\r
+UINT16                    mOrigIdtEntryCount    = 0;\r
 \r
 FIXED_MTRR    mFixedMtrrTable[] = {\r
   {\r
@@ -520,7 +522,15 @@ CpuRegisterInterruptHandler (
     return EFI_ALREADY_STARTED;\r
   }\r
 \r
-  SetInterruptDescriptorTableHandlerAddress ((UINTN)InterruptType, NULL);\r
+  if (InterruptHandler != NULL) {\r
+    SetInterruptDescriptorTableHandlerAddress ((UINTN)InterruptType, NULL);\r
+  } else {\r
+    //\r
+    // Restore the original IDT handler address if InterruptHandler is NULL.\r
+    //\r
+    RestoreInterruptDescriptorTableHandlerAddress ((UINTN)InterruptType);\r
+  }\r
+\r
   ExternalVectorTable[InterruptType] = InterruptHandler;\r
   return EFI_SUCCESS;\r
 }\r
@@ -1081,6 +1091,25 @@ SetInterruptDescriptorTableHandlerAddress (
 #endif\r
 }\r
 \r
+/**\r
+  Restore original Interrupt Descriptor Table Handler Address.\r
+\r
+  @param Index        The Index of the interrupt descriptor table handle.\r
+\r
+**/\r
+VOID\r
+RestoreInterruptDescriptorTableHandlerAddress (\r
+  IN UINTN       Index\r
+  )\r
+{\r
+  if (Index < mOrigIdtEntryCount) {\r
+    gIdtTable[Index].Bits.OffsetLow   = mOrigIdtEntry[Index].Bits.OffsetLow;\r
+    gIdtTable[Index].Bits.OffsetHigh  = mOrigIdtEntry[Index].Bits.OffsetHigh;\r
+#if defined (MDE_CPU_X64)\r
+    gIdtTable[Index].Bits.OffsetUpper = mOrigIdtEntry[Index].Bits.OffsetUpper;\r
+#endif\r
+  }\r
+}\r
 \r
 /**\r
   Initialize Interrupt Descriptor Table for interrupt handling.\r
@@ -1111,6 +1140,12 @@ InitInterruptDescriptorTable (
   if ((OldIdtPtr.Base != 0) && ((OldIdtPtr.Limit & 7) == 7)) {\r
     OldIdt = (IA32_IDT_GATE_DESCRIPTOR*) OldIdtPtr.Base;\r
     OldIdtSize = (OldIdtPtr.Limit + 1) / sizeof (IA32_IDT_GATE_DESCRIPTOR);\r
+    //\r
+    // Save original IDT entry and IDT entry count.\r
+    //\r
+    mOrigIdtEntry = AllocateCopyPool (OldIdtPtr.Limit + 1, (VOID *) OldIdtPtr.Base);\r
+    ASSERT (mOrigIdtEntry != NULL);\r
+    mOrigIdtEntryCount = (UINT16) OldIdtSize;\r
   } else {\r
     OldIdt = NULL;\r
     OldIdtSize = 0;\r