From 3860b1a2328a7562923c522e8eeef9f714e92370 Mon Sep 17 00:00:00 2001
From: Qiu Shumin <shumin.qiu@intel.com>
Date: Wed, 1 Jul 2015 08:28:13 +0000
Subject: [PATCH] MdeModulePkg/Universal/PlatformDriOverrideDxe: Use safe
 string functions to refine code.

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Qiu Shumin <shumin.qiu@intel.com>
Reviewed-by: Liming Gao <liming.gao@intel.com>

git-svn-id: https://svn.code.sf.net/p/edk2/code/trunk/edk2@17787 6f19259b-4bc3-4df7-8a09-765794883524
---
 .../PlatDriOverrideDxe.c                      | 21 ++++++++++++-------
 1 file changed, 13 insertions(+), 8 deletions(-)

diff --git a/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatDriOverrideDxe.c b/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatDriOverrideDxe.c
index 7027e9697e..1274d6720d 100644
--- a/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatDriOverrideDxe.c
+++ b/MdeModulePkg/Universal/PlatformDriOverrideDxe/PlatDriOverrideDxe.c
@@ -333,6 +333,7 @@ UpdateDeviceSelectPage (
   EFI_STATUS                                Status;
   UINTN                                     Index;
   UINTN                                     DevicePathHandleCount;
+  UINTN                                     NewStrSize;
   CHAR16                                    *NewString;
   EFI_STRING_ID                             NewStringToken;
   CHAR16                                    *ControllerName;
@@ -489,14 +490,15 @@ UpdateDeviceSelectPage (
     // Export the driver name string and create item in set options page
     //
     Len = StrSize (ControllerName);
-    NewString = AllocateZeroPool (Len + StrSize (L"--"));
+    NewStrSize = Len + StrSize (L"--");
+    NewString = AllocateZeroPool (NewStrSize);
     ASSERT (NewString != NULL);
     if (EFI_ERROR (CheckMapping (ControllerDevicePath,NULL, &mMappingDataBase, NULL, NULL))) {
-      StrCat (NewString, L"--");
+      StrCatS (NewString, NewStrSize/sizeof(CHAR16), L"--");
     } else {
-      StrCat (NewString, L"**");
+      StrCatS (NewString, NewStrSize/sizeof(CHAR16), L"**");
     }
-    StrCat (NewString, ControllerName);
+    StrCatS (NewString, NewStrSize/sizeof(CHAR16), ControllerName);
 
     NewStringToken = HiiSetString (Private->RegisteredHandle, mControllerToken[Index], NewString, NULL);
     ASSERT (NewStringToken != 0);
@@ -622,6 +624,7 @@ UpdateBindingDriverSelectPage (
 {
   EFI_STATUS                                Status;
   UINTN                                     Index;
+  UINTN                                     NewStrSize;
   CHAR16                                    *NewString;
   EFI_STRING_ID                             NewStringToken;
   EFI_STRING_ID                             NewStringHelpToken;
@@ -814,7 +817,8 @@ UpdateBindingDriverSelectPage (
     //
     // First create the driver image name
     //
-    NewString = AllocateZeroPool (StrSize (DriverName));
+    NewStrSize = StrSize (DriverName);
+    NewString = AllocateZeroPool (NewStrSize);
     ASSERT (NewString != NULL); 
     if (EFI_ERROR (CheckMapping (mControllerDevicePathProtocol[mSelectedCtrIndex], LoadedImageDevicePath, &mMappingDataBase, NULL, NULL))) {
       mDriSelection[Index] = FALSE;
@@ -822,7 +826,7 @@ UpdateBindingDriverSelectPage (
       mDriSelection[Index] = TRUE;
       mLastSavedDriverImageNum++;
     }
-    StrCat (NewString, DriverName);
+    StrCatS (NewString, NewStrSize/sizeof(CHAR16), DriverName);
     NewStringToken = HiiSetString (Private->RegisteredHandle, mDriverImageToken[Index], NewString, NULL);
     ASSERT (NewStringToken != 0);
     mDriverImageToken[Index] = NewStringToken;
@@ -836,9 +840,10 @@ UpdateBindingDriverSelectPage (
     //
     DriverName = DevicePathToStr (LoadedImageDevicePath);
 
-    NewString = AllocateZeroPool (StrSize (DriverName));
+    NewStrSize = StrSize (DriverName);
+    NewString = AllocateZeroPool (NewStrSize);
     ASSERT (NewString != NULL); 
-    StrCat (NewString, DriverName);
+    StrCatS (NewString, NewStrSize/sizeof(CHAR16), DriverName);
     NewStringHelpToken = HiiSetString (Private->RegisteredHandle, DriverImageFilePathToken[Index], NewString, NULL);
     ASSERT (NewStringHelpToken != 0);
     DriverImageFilePathToken[Index] = NewStringHelpToken;
-- 
2.39.2