Refine the code to make it more safely.

[mirror_edk2.git] / ShellPkg / Library / UefiShellDebug1CommandsLib / DmpStore.c

diff --git a/ShellPkg/Library/UefiShellDebug1CommandsLib/DmpStore.c b/ShellPkg/Library/UefiShellDebug1CommandsLib/DmpStore.c
index be3d6f64724a51714d7479438c4c5a3818a896e1..2fed064f84eab1777e08e0ab452a4c95442d7a40 100644 (file)
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/DmpStore.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/DmpStore.c
@@ -17,7 +17,6 @@
 \r
 #define INIT_NAME_BUFFER_SIZE  128\r
 #define INIT_DATA_BUFFER_SIZE  1024\r
-#define INIT_ATTS_BUFFER_SIZE  64\r
 \r
 /**\r
   Base on the input attribute value to return the attribute string.\r
@@ -27,41 +26,45 @@
 \r
   @retval The attribute string info.\r
 **/\r
-CONST CHAR16 *\r
+CHAR16 *\r
 EFIAPI\r
 GetAttrType (\r
-  IN CONST UINT32 Atts,\r
-  IN OUT   CHAR16 *RetString\r
+  IN CONST UINT32 Atts\r
   )\r
 {\r
-  StrCpy(RetString, L"");\r
+  UINT32 BufLen;\r
+  CHAR16 *RetString;\r
 \r
+  BufLen      = 0;\r
+  RetString   = NULL;\r
+ \r
   if ((Atts & EFI_VARIABLE_NON_VOLATILE) != 0) {\r
-    StrCat(RetString, L"+NV");\r
+    StrnCatGrow (&RetString, &BufLen, L"+NV", 0);\r
   }\r
   if ((Atts & EFI_VARIABLE_RUNTIME_ACCESS) != 0) {\r
-    StrCat(RetString, L"+RS+BS");\r
+    StrnCatGrow (&RetString, &BufLen, L"+RS+BS", 0);\r
   } else if ((Atts & EFI_VARIABLE_BOOTSERVICE_ACCESS) != 0) {\r
-    StrCat(RetString, L"+BS");\r
+    StrnCatGrow (&RetString, &BufLen, L"+BS", 0);\r
   }\r
   if ((Atts & EFI_VARIABLE_HARDWARE_ERROR_RECORD) != 0) {\r
-    StrCat(RetString, L"+HR");\r
+    StrnCatGrow (&RetString, &BufLen, L"+HR", 0);\r
   }\r
   if ((Atts & EFI_VARIABLE_AUTHENTICATED_WRITE_ACCESS) != 0) {\r
-    StrCat(RetString, L"+AW");\r
+    StrnCatGrow (&RetString, &BufLen, L"+AW", 0);\r
   }\r
   if ((Atts & EFI_VARIABLE_TIME_BASED_AUTHENTICATED_WRITE_ACCESS) != 0) {\r
-    StrCat(RetString, L"+AT");\r
+    StrnCatGrow (&RetString, &BufLen, L"+AT", 0);\r
   }\r
 \r
-  if (RetString[0] == L'+') {\r
-    return (RetString+1);\r
+  if (RetString == NULL) {\r
+    RetString = StrnCatGrow(&RetString, &BufLen, L"Invalid", 0);\r
   }\r
-  if (RetString[0] == CHAR_NULL) {\r
-    StrCpy(RetString, L"invalid");\r
-    return (RetString);\r
+\r
+  if (RetString[0] == L'+') {\r
+    CopyMem(RetString, RetString + 1, StrSize(RetString + 1));\r
   }\r
-  return (RetString);\r
+\r
+  return RetString;\r
 }\r
 \r
 /**\r
@@ -98,7 +101,7 @@ ProcessVariables (
   CHAR16                    *OldName;\r
   UINTN                     OldNameBufferSize;\r
   UINTN                     DataBufferSize; // Allocated data buffer size\r
-  CHAR16                    RetString[INIT_ATTS_BUFFER_SIZE];\r
+  CHAR16                    *RetString;\r
 \r
   Found         = FALSE;\r
   ShellStatus   = SHELL_SUCCESS;\r
@@ -186,6 +189,7 @@ ProcessVariables (
     // do the print or delete\r
     //\r
     Found = TRUE;\r
+    RetString = GetAttrType(Atts);\r
     if (!Delete) {\r
       ShellPrintHiiEx(\r
         -1,\r
@@ -193,7 +197,7 @@ ProcessVariables (
         NULL,\r
         STRING_TOKEN(STR_DMPSTORE_HEADER_LINE),\r
         gShellDebug1HiiHandle,\r
-        GetAttrType(Atts, RetString),\r
+        RetString,\r
         &FoundVarGuid,\r
         FoundVarName,\r
         DataSize);\r
@@ -216,6 +220,10 @@ ProcessVariables (
         gRT->SetVariable(FoundVarName, &FoundVarGuid, Atts, 0, NULL));\r
         FoundVarName[0] = CHAR_NULL;\r
     }\r
+\r
+    if (RetString != NULL) {\r
+      FreePool (RetString);\r
+    }\r
   }\r
 \r
   if (FoundVarName != NULL) {\r