From 9c2eb6d64280222990a17705eeb92eb0743ccd0c Mon Sep 17 00:00:00 2001
From: Mikhail Krichanov <mikhailkrichanov@gmail.com>
Date: Tue, 17 Sep 2024 11:37:46 +0300
Subject: [PATCH] CpuExceptionHandlerLib: Refactored IO Bit Map initialization.

---
 .../Ia32/ArchExceptionHandler.c               | 22 ++++++++++---------
 .../Ia32/ArchInterruptDefs.h                  |  2 +-
 .../X64/ArchExceptionHandler.c                | 22 ++++++++++---------
 .../X64/ArchInterruptDefs.h                   |  2 +-
 4 files changed, 26 insertions(+), 22 deletions(-)

diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
index ae39125a1c..dfd803c7a6 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchExceptionHandler.c
@@ -136,6 +136,8 @@ ArchSetupExceptionStack (
   UINTN                           NeedBufferSize;
   EXCEPTION_HANDLER_TEMPLATE_MAP  TemplateMap;
   UINT8                           *IOBitMap;
+  UINT8                           *IOBitMapPointer;
+  UINT8                           Offset;
 
   if (BufferSize == NULL) {
     return EFI_INVALID_PARAMETER;
@@ -234,17 +236,17 @@ ArchSetupExceptionStack (
   // and DebugIoPort = 0x402.
   //
   IOBitMap = (UINT8 *)((UINTN)Tss + Tss->IOMapBaseAddress);
-  for (Index = 0; Index < IO_BIT_MAP_SIZE; ++Index) {
-    if ((Index * 8) == FixedPcdGet16 (PcdUartBase)) {
-      *IOBitMap = 0x84;
-    } else if ((Index * 8) == (FixedPcdGet16 (PcdDebugIoPort) - 2)) {
-      *IOBitMap = 0xFB;
-    } else {
-      *IOBitMap = 0xFF;
-    }
+  SetMem (IOBitMap, IO_BIT_MAP_SIZE, 0xFF);
 
-    ++IOBitMap;
-  }
+  IOBitMapPointer             = (UINT8 *)((UINTN)IOBitMap + FixedPcdGet16 (PcdUartBase) / 8);
+  Offset                      = (UINT8)(FixedPcdGet16 (PcdUartBase) & 0x7U);
+  *(UINT16 *)IOBitMapPointer &= ~((1U << Offset) | (1U << (Offset + 1))
+                                | (1U << (Offset + 3)) | (1U << (Offset + 4))
+                                | (1U << (Offset + 5)) | (1U << (Offset + 6)));
+
+  IOBitMapPointer   = (UINT8 *)((UINTN)IOBitMap + FixedPcdGet16 (PcdDebugIoPort) / 8);
+  Offset            = (UINT8)(FixedPcdGet16 (PcdDebugIoPort) & 0x7U);
+  *IOBitMapPointer &= ~(1U << Offset);
 
   Tss = (IA32_TASK_STATE_SEGMENT *)((UINTN)Tss + sizeof (IA32_TASK_STATE_SEGMENT) + IO_BIT_MAP_SIZE);
   ++TssDesc;
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchInterruptDefs.h b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchInterruptDefs.h
index eeae24bdb3..f470cb2306 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchInterruptDefs.h
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/Ia32/ArchInterruptDefs.h
@@ -39,7 +39,7 @@ typedef struct {
   (sizeof (IA32_TSS_DESCRIPTOR) * \
    (FixedPcdGetSize (PcdCpuStackSwitchExceptionList) + 1))
 
-#define IO_BIT_MAP_SIZE    (ALIGN_VALUE (FixedPcdGet16 (PcdDebugIoPort) / 8 + 1, 16))
+#define IO_BIT_MAP_SIZE    (ALIGN_VALUE (MAX (FixedPcdGet16 (PcdUartBase) ,FixedPcdGet16 (PcdDebugIoPort)) / 8 + 1, 16))
 #define CPU_TSS_SIZE \
   (sizeof (IA32_TASK_STATE_SEGMENT) * \
    (FixedPcdGetSize (PcdCpuStackSwitchExceptionList) + 1) + IO_BIT_MAP_SIZE)
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
index 1b35108e97..236e72b78b 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchExceptionHandler.c
@@ -140,6 +140,8 @@ ArchSetupExceptionStack (
   UINT8                     *StackSwitchExceptions;
   UINTN                     NeedBufferSize;
   UINT8                     *IOBitMap;
+  UINT8                     *IOBitMapPointer;
+  UINT8                     Offset;
 
   if (BufferSize == NULL) {
     return EFI_INVALID_PARAMETER;
@@ -241,17 +243,17 @@ ArchSetupExceptionStack (
   // and DebugIoPort = 0x402.
   //
   IOBitMap = (UINT8 *)((UINTN)Tss + Tss->IOMapBaseAddress);
-  for (Index = 0; Index < IO_BIT_MAP_SIZE; ++Index) {
-    if ((Index * 8) == FixedPcdGet16 (PcdUartBase)) {
-      *IOBitMap = 0x84;
-    } else if ((Index * 8) == (FixedPcdGet16 (PcdDebugIoPort) - 2)) {
-      *IOBitMap = 0xFB;
-    } else {
-      *IOBitMap = 0xFF;
-    }
+  SetMem (IOBitMap, IO_BIT_MAP_SIZE, 0xFF);
 
-    ++IOBitMap;
-  }
+  IOBitMapPointer             = (UINT8 *)((UINTN)IOBitMap + FixedPcdGet16 (PcdUartBase) / 8);
+  Offset                      = (UINT8)(FixedPcdGet16 (PcdUartBase) & 0x7U);
+  *(UINT16 *)IOBitMapPointer &= ~((1U << Offset) | (1U << (Offset + 1))
+                                | (1U << (Offset + 3)) | (1U << (Offset + 4))
+                                | (1U << (Offset + 5)) | (1U << (Offset + 6)));
+
+  IOBitMapPointer   = (UINT8 *)((UINTN)IOBitMap + FixedPcdGet16 (PcdDebugIoPort) / 8);
+  Offset            = (UINT8)(FixedPcdGet16 (PcdDebugIoPort) & 0x7U);
+  *IOBitMapPointer &= ~(1U << Offset);
 
   //
   // Fixup IST and task-state segment
diff --git a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchInterruptDefs.h b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchInterruptDefs.h
index 7ae9ff7b0a..0122b3c07f 100644
--- a/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchInterruptDefs.h
+++ b/UefiCpuPkg/Library/CpuExceptionHandlerLib/X64/ArchInterruptDefs.h
@@ -38,7 +38,7 @@ typedef struct {
 } RESERVED_VECTORS_DATA;
 
 #define CPU_TSS_DESC_SIZE  sizeof (IA32_TSS_DESCRIPTOR)
-#define IO_BIT_MAP_SIZE    (ALIGN_VALUE (FixedPcdGet16 (PcdDebugIoPort) / 8 + 1, 16))
+#define IO_BIT_MAP_SIZE    (ALIGN_VALUE (MAX (FixedPcdGet16 (PcdUartBase) ,FixedPcdGet16 (PcdDebugIoPort)) / 8 + 1, 16))
 #define CPU_TSS_SIZE       (sizeof (IA32_TASK_STATE_SEGMENT) + IO_BIT_MAP_SIZE)
 
 #endif