MdeModulePkg/Library/UefiBootManagerLib: 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: Ruiyu Ni <ruiyu.ni@intel.com>

git-svn-id: https://svn.code.sf.net/p/edk2/code/trunk/edk2@17782 6f19259b-4bc3-4df7-8a09-765794883524

2 files changed, 28 insertions, 16 deletions

diff --git a/MdeModulePkg/Library/UefiBootManagerLib/BmBoot.c b/MdeModulePkg/Library/UefiBootManagerLib/BmBoot.c
index 2d3d57b3e9..028edc35d5 100644
--- a/MdeModulePkg/Library/UefiBootManagerLib/BmBoot.c
+++ b/MdeModulePkg/Library/UefiBootManagerLib/BmBoot.c
@@ -552,6 +552,7 @@ BmGetUsbDescription (
   CHAR16                       *SerialNumber;

   CHAR16                       *Description;

   EFI_USB_DEVICE_DESCRIPTOR    DevDesc;

+  UINTN                        DescMaxSize;

 

   Status = gBS->HandleProtocol (

                   Handle,

@@ -606,15 +607,16 @@ BmGetUsbDescription (
     return NULL;

   }

 

-  Description = AllocateZeroPool (StrSize (Manufacturer) + StrSize (Product) + StrSize (SerialNumber));

+  DescMaxSize = StrSize (Manufacturer) + StrSize (Product) + StrSize (SerialNumber);

+  Description = AllocateZeroPool (DescMaxSize);

   ASSERT (Description != NULL);

-  StrCat (Description, Manufacturer);

-  StrCat (Description, L" ");

+  StrCatS (Description, DescMaxSize/sizeof(CHAR16), Manufacturer);

+  StrCatS (Description, DescMaxSize/sizeof(CHAR16), L" ");

 

-  StrCat (Description, Product);  

-  StrCat (Description, L" ");

+  StrCatS (Description, DescMaxSize/sizeof(CHAR16), Product);  

+  StrCatS (Description, DescMaxSize/sizeof(CHAR16), L" ");

 

-  StrCat (Description, SerialNumber);

+  StrCatS (Description, DescMaxSize/sizeof(CHAR16), SerialNumber);

 

   if (Manufacturer != &NullChar) {

     FreePool (Manufacturer);

@@ -774,8 +776,14 @@ BmGetBootDescription (
       //

       Temp = AllocatePool (StrSize (DefaultDescription) + sizeof (mBmUefiPrefix)); 

       ASSERT (Temp != NULL);

-      StrCpy (Temp, mBmUefiPrefix);

-      StrCat (Temp, DefaultDescription);

+      StrCpyS ( Temp, 

+                (StrSize (DefaultDescription) + sizeof (mBmUefiPrefix))/sizeof(CHAR16), 

+                mBmUefiPrefix

+                );

+      StrCatS ( Temp, 

+                (StrSize (DefaultDescription) + sizeof (mBmUefiPrefix))/sizeof(CHAR16), 

+                DefaultDescription

+                );

       FreePool (DefaultDescription);

       DefaultDescription = Temp;

       break;

diff --git a/MdeModulePkg/Library/UefiBootManagerLib/BmPerformance.c b/MdeModulePkg/Library/UefiBootManagerLib/BmPerformance.c
index 32229d0b0a..dd937a7396 100644
--- a/MdeModulePkg/Library/UefiBootManagerLib/BmPerformance.c
+++ b/MdeModulePkg/Library/UefiBootManagerLib/BmPerformance.c
@@ -26,12 +26,14 @@ EFI_PHYSICAL_ADDRESS      mBmAcpiLowMemoryBase = 0x0FFFFFFFFULL;
 

   @param PdbFileName     The long PDB file name.

   @param GaugeString     The output string to be logged by performance logger.

+  @param StringSize      The buffer size of GaugeString in bytes.

 

 **/

 VOID

 BmGetShortPdbFileName (

   IN  CONST CHAR8  *PdbFileName,

-  OUT       CHAR8  *GaugeString

+  OUT       CHAR8  *GaugeString,

+  IN        UINTN   StringSize

   )

 {

   UINTN Index;

@@ -40,7 +42,7 @@ BmGetShortPdbFileName (
   UINTN EndIndex;

 

   if (PdbFileName == NULL) {

-    AsciiStrCpy (GaugeString, " ");

+    AsciiStrCpyS (GaugeString, StringSize, " ");

   } else {

     StartIndex = 0;

     for (EndIndex = 0; PdbFileName[EndIndex] != 0; EndIndex++)

@@ -78,12 +80,14 @@ BmGetShortPdbFileName (
 

   @param Handle          Driver handle.

   @param GaugeString     The output string to be logged by performance logger.

+  @param StringSize      The buffer size of GaugeString in bytes.

 

 **/

 VOID

 BmGetNameFromHandle (

   IN  EFI_HANDLE     Handle,

-  OUT CHAR8          *GaugeString

+  OUT CHAR8          *GaugeString,

+  IN  UINTN          StringSize

   )

 {

   EFI_STATUS                  Status;

@@ -91,7 +95,7 @@ BmGetNameFromHandle (
   CHAR8                       *PdbFileName;

   EFI_DRIVER_BINDING_PROTOCOL *DriverBinding;

 

-  AsciiStrCpy (GaugeString, " ");

+  AsciiStrCpyS (GaugeString, StringSize, " ");

 

   //

   // Get handle name from image protocol

@@ -127,7 +131,7 @@ BmGetNameFromHandle (
   PdbFileName = PeCoffLoaderGetPdbPointer (Image->ImageBase);

 

   if (PdbFileName != NULL) {

-    BmGetShortPdbFileName (PdbFileName, GaugeString);

+    BmGetShortPdbFileName (PdbFileName, GaugeString, StringSize);

   }

 

   return ;

@@ -285,9 +289,9 @@ BmWriteBootToOsPerformanceData (
 

     if (Duration > 0) {

 

-      BmGetNameFromHandle (Handles[Index], GaugeString);

+      BmGetNameFromHandle (Handles[Index], GaugeString, PERF_TOKEN_LENGTH);

 

-      AsciiStrCpy (mBmPerfData.Token, GaugeString);

+      AsciiStrCpyS (mBmPerfData.Token, PERF_TOKEN_SIZE, GaugeString);

       mBmPerfData.Duration = Duration;

 

       CopyMem (Ptr, &mBmPerfData, sizeof (PERF_DATA));

@@ -316,7 +320,7 @@ BmWriteBootToOsPerformanceData (
 

       ZeroMem (&mBmPerfData, sizeof (PERF_DATA));

 

-      AsciiStrnCpy (mBmPerfData.Token, Token, PERF_TOKEN_LENGTH);

+      AsciiStrCpyS (mBmPerfData.Token, PERF_TOKEN_SIZE, Token);

       if (StartTicker == 1) {

         StartTicker = StartValue;

       }