From 231ea883ed13a43157e2ceca5869d02d5d727526 Mon Sep 17 00:00:00 2001
From: Qiu Shumin <shumin.qiu@intel.com>
Date: Thu, 21 Aug 2014 05:35:25 +0000
Subject: [PATCH] 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
---
 .../BdsDxe/DeviceMngr/DeviceManager.c         | 61 ++++++++++++-------
 1 file changed, 40 insertions(+), 21 deletions(-)

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