From de4b64f714332d80e9f1e0ada42c1998b5fbb81c Mon Sep 17 00:00:00 2001
From: vanjeff <vanjeff@6f19259b-4bc3-4df7-8a09-765794883524>
Date: Wed, 16 Nov 2011 02:31:31 +0000
Subject: [PATCH] Restore original IDT entry if RegisterInterruptHandler() was
 used to unregister user defined interrupt handler.

Signed-off-by: vanjeff
Reviewed-by: rsun3



git-svn-id: https://edk2.svn.sourceforge.net/svnroot/edk2/trunk/edk2@12719 6f19259b-4bc3-4df7-8a09-765794883524
---
 UefiCpuPkg/CpuDxe/CpuDxe.c | 37 ++++++++++++++++++++++++++++++++++++-
 UefiCpuPkg/CpuDxe/CpuDxe.h | 10 ++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

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