MdePkg/BaseUefiDecompressLib: Improve performance of boundary validation

[mirror_edk2.git] / MdePkg / Library / BaseUefiDecompressLib / BaseUefiDecompressLib.c

diff --git a/MdePkg/Library/BaseUefiDecompressLib/BaseUefiDecompressLib.c b/MdePkg/Library/BaseUefiDecompressLib/BaseUefiDecompressLib.c
index dc89157c4efd47c2a1fac61757faf86f754c39a2..3d5b7a737ad6fa88d56d418b54536a582ec6b143 100644 (file)
--- a/MdePkg/Library/BaseUefiDecompressLib/BaseUefiDecompressLib.c
+++ b/MdePkg/Library/BaseUefiDecompressLib/BaseUefiDecompressLib.c
@@ -152,6 +152,7 @@ MakeTable (
   UINT16  Mask;\r
   UINT16  WordOfStart;\r
   UINT16  WordOfCount;\r
+  UINT16  MaxTableLength;\r
 \r
   //\r
   // The maximum mapping table width supported by this internal\r
@@ -164,6 +165,9 @@ MakeTable (
   }\r
 \r
   for (Index = 0; Index < NumOfChar; Index++) {\r
+    if (BitLen[Index] > 16) {\r
+      return (UINT16) BAD_TABLE;\r
+    }\r
     Count[BitLen[Index]]++;\r
   }\r
 \r
@@ -205,6 +209,7 @@ MakeTable (
 \r
   Avail = NumOfChar;\r
   Mask  = (UINT16) (1U << (15 - TableBits));\r
+  MaxTableLength = (UINT16) (1U << TableBits);\r
 \r
   for (Char = 0; Char < NumOfChar; Char++) {\r
 \r
@@ -217,6 +222,10 @@ MakeTable (
 \r
     if (Len <= TableBits) {\r
 \r
+      if (Start[Len] >= NextCode || NextCode > MaxTableLength){\r
+        return (UINT16) BAD_TABLE;\r
+      }\r
+\r
       for (Index = Start[Len]; Index < NextCode; Index++) {\r
         Table[Index] = Char;\r
       }\r
@@ -620,14 +629,25 @@ Decode (
       // Write BytesRemain of bytes into mDstBase\r
       //\r
       BytesRemain--;\r
+\r
       while ((INT16) (BytesRemain) >= 0) {\r
-        Sd->mDstBase[Sd->mOutBuf++] = Sd->mDstBase[DataIdx++];\r
         if (Sd->mOutBuf >= Sd->mOrigSize) {\r
           goto Done;\r
         }\r
+        if (DataIdx >= Sd->mOrigSize) {\r
+          Sd->mBadTableFlag = (UINT16) BAD_TABLE;\r
+          goto Done;\r
+        }\r
+        Sd->mDstBase[Sd->mOutBuf++] = Sd->mDstBase[DataIdx++];\r
 \r
         BytesRemain--;\r
       }\r
+      //\r
+      // Once mOutBuf is fully filled, directly return\r
+      //\r
+      if (Sd->mOutBuf >= Sd->mOrigSize) {\r
+        goto Done;\r
+      }\r
     }\r
   }\r
 \r
@@ -694,7 +714,7 @@ UefiDecompressGetInfo (
   }\r
 \r
   CompressedSize   = ReadUnaligned32 ((UINT32 *)Source);\r
-  if (SourceSize < (CompressedSize + 8)) {\r
+  if (SourceSize < (CompressedSize + 8) || (CompressedSize + 8) < 8) {\r
     return RETURN_INVALID_PARAMETER;\r
   }\r
 \r