Allocate exact memory size for string buffer to avoid buffer overflow.

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Qiu Shumin <shumin.qiu@intel.com>
Reviewed-by: Ni Ruiyu <ruiyu.ni@intel.com>
git-svn-id: https://svn.code.sf.net/p/edk2/code/trunk/edk2@15863 6f19259b-4bc3-4df7-8a09-765794883524

diff --git a/IntelFrameworkModulePkg/Universal/BdsDxe/DeviceMngr/DeviceManager.c b/IntelFrameworkModulePkg/Universal/BdsDxe/DeviceMngr/DeviceManager.c
index 866df574f6e7c7550a55478d82b9004948e855d6..85cbe8951d5234801abec757a8279bcd03f5f065 100644 (file)
--- a/IntelFrameworkModulePkg/Universal/BdsDxe/DeviceMngr/DeviceManager.c
+++ b/IntelFrameworkModulePkg/Universal/BdsDxe/DeviceMngr/DeviceManager.c
@@ -1260,11 +1260,13 @@ CallDriverHealth (
   LIST_ENTRY                  *Link;\r
   EFI_DEVICE_PATH_PROTOCOL    *DriverDevicePath;\r
   BOOLEAN                     RebootRequired;\r
+  BOOLEAN                     IsControllerNameEmpty;\r
   UINTN                       StringSize;\r
 \r
   Index               = 0;\r
   DriverHealthInfo    = NULL;  \r
   DriverDevicePath    = NULL;\r
+  IsControllerNameEmpty = FALSE;\r
   InitializeListHead (&DriverHealthList);\r
 \r
   HiiHandle = gDeviceManagerPrivate.DriverHealthHiiHandle;\r
@@ -1337,14 +1339,7 @@ CallDriverHealth (
   Link = GetFirstNode (&DriverHealthList);\r
 \r
   while (!IsNull (&DriverHealthList, Link)) {   \r
-    DriverHealthInfo = DEVICE_MANAGER_HEALTH_INFO_FROM_LINK (Link);\r
-    \r
-    //\r
-    // Assume no line strings is longer than 512 bytes.\r
-    //\r
-    StringSize = 0x200;\r
-    String = (EFI_STRING) AllocateZeroPool (StringSize);\r
-    ASSERT (String != NULL);\r
+    DriverHealthInfo = DEVICE_MANAGER_HEALTH_INFO_FROM_LINK (Link);    \r
 \r
     Status = DriverHealthGetDriverName (DriverHealthInfo->DriverHandle, &DriverName);\r
     if (EFI_ERROR (Status)) {\r
@@ -1354,11 +1349,7 @@ CallDriverHealth (
       DriverDevicePath = DevicePathFromHandle (DriverHealthInfo->DriverHandle);\r
       DriverName       = DevicePathToStr (DriverDevicePath);\r
     }\r
-    //\r
-    // Add the Driver name & Controller name into FormSetTitle string\r
-    // \r
-    StrnCat (String, DriverName, StrLen (DriverName));\r
-\r
+    StringSize = StrSize (DriverName);\r
 \r
     Status = DriverHealthGetControllerName (\r
                DriverHealthInfo->DriverHandle, \r
@@ -1368,23 +1359,39 @@ CallDriverHealth (
                );\r
 \r
     if (!EFI_ERROR (Status)) {\r
-      //\r
-      // Can not get the Controller name, just let it empty.\r
-      //\r
-      StrnCat (String, L"    ", StrLen (L"    "));\r
-      StrnCat (String, ControllerName, StrLen (ControllerName));   \r
+      IsControllerNameEmpty = FALSE;\r
+      StringSize += StrLen (L"    ") * sizeof(CHAR16);\r
+      StringSize += StrLen (ControllerName) * sizeof(CHAR16);\r
+    } else {\r
+      IsControllerNameEmpty = TRUE;\r
     }\r
    \r
     //\r
     // Add the message of the Module itself provided after the string item.\r
     //\r
     if ((DriverHealthInfo->MessageList != NULL) && (DriverHealthInfo->MessageList->StringId != 0)) {\r
-       StrnCat (String, L"    ", StrLen (L"    "));\r
        TmpString = HiiGetString (\r
                      DriverHealthInfo->MessageList->HiiHandle, \r
                      DriverHealthInfo->MessageList->StringId, \r
                      NULL\r
                      );\r
+       ASSERT (TmpString != NULL);\r
+       \r
+       StringSize += StrLen (L"    ") * sizeof(CHAR16);\r
+       StringSize += StrLen (TmpString) * sizeof(CHAR16);\r
+\r
+       String = (EFI_STRING) AllocateZeroPool (StringSize);\r
+       ASSERT (String != NULL);\r
+       \r
+       StrnCpy (String, DriverName, StringSize / sizeof(CHAR16));\r
+       if (!IsControllerNameEmpty) {\r
+        StrnCat (String, L"    ", StringSize / sizeof(CHAR16) - StrLen(String) - 1);\r
+        StrnCat (String, ControllerName, StringSize / sizeof(CHAR16) - StrLen(String) - 1);\r
+       }\r
+\r
+       StrnCat (String, L"    ", StringSize / sizeof(CHAR16) - StrLen(String) - 1);\r
+       StrnCat (String, TmpString, StringSize / sizeof(CHAR16) - StrLen(String) - 1);\r
+       \r
     } else {\r
       //\r
       // Update the string will be displayed base on the driver's health status\r
@@ -1409,10 +1416,22 @@ CallDriverHealth (
         TmpString = GetStringById (STRING_TOKEN (STR_DRIVER_HEALTH_HEALTHY));\r
         break;\r
       }\r
+      ASSERT (TmpString != NULL);\r
+\r
+      StringSize += StrLen (TmpString) * sizeof(CHAR16);\r
+\r
+      String = (EFI_STRING) AllocateZeroPool (StringSize);\r
+      ASSERT (String != NULL);\r
+      \r
+      StrnCpy (String, DriverName, StringSize / sizeof(CHAR16));\r
+      if (!IsControllerNameEmpty) {\r
+        StrnCat (String, L"    ", StringSize / sizeof(CHAR16) - StrLen(String) - 1);\r
+        StrnCat (String, ControllerName, StringSize / sizeof(CHAR16) - StrLen(String) - 1);\r
+      }\r
+\r
+      StrnCat (String, TmpString, StringSize / sizeof(CHAR16) - StrLen(String) - 1);\r
     }\r
 \r
-    ASSERT (TmpString != NULL);\r
-    StrnCat (String, TmpString, StringSize / sizeof (CHAR16) - StrLen (String) - 1);\r
     FreePool (TmpString);\r
 \r
     Token = HiiSetString (HiiHandle, 0, String, NULL);\r