MdeModulePkg/Universal/PlatformDriOverrideDxe: Use safe string functions to refine...

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

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