ShellPkg: Refactor string manipulation

[mirror_edk2.git] / ShellPkg / Application / Shell / ShellEnvVar.c

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.c b/ShellPkg/Application/Shell/ShellEnvVar.c
index f40a867cc7ae564a6f4e607e36e4ed56205fab5f..b94711142af6c3ff4532ed1b6127e2b70eea9e31 100644 (file)
--- a/ShellPkg/Application/Shell/ShellEnvVar.c
+++ b/ShellPkg/Application/Shell/ShellEnvVar.c
@@ -152,7 +152,7 @@ GetEnvironmentVariableList(
   if (VariableName == NULL) {\r
     return (EFI_OUT_OF_RESOURCES);\r
   }\r
-  StrCpy(VariableName, L"");\r
+  *VariableName = CHAR_NULL;\r
 \r
   while (!EFI_ERROR(Status)) {\r
     NameSize = (UINTN)MaxVarSize;\r
@@ -178,13 +178,12 @@ GetEnvironmentVariableList(
           }\r
         }\r
         if (!EFI_ERROR(Status) && VarList != NULL) {\r
-          VarList->Key = AllocateZeroPool(StrSize(VariableName));\r
+          VarList->Key = AllocateCopyPool(StrSize(VariableName), VariableName);\r
           if (VarList->Key == NULL) {\r
             SHELL_FREE_NON_NULL(VarList->Val);\r
             SHELL_FREE_NON_NULL(VarList);\r
             Status = EFI_OUT_OF_RESOURCES;\r
           } else {\r
-            StrCpy(VarList->Key, VariableName);\r
             InsertTailList(ListHead, &VarList->Link);\r
           }\r
         }\r
@@ -286,7 +285,6 @@ SetEnvironmentVariables(
   UINTN         CurrentCount;\r
   ENV_VAR_LIST  *VarList;\r
   ENV_VAR_LIST  *Node;\r
-  UINTN         NewSize;\r
 \r
   VarList = NULL;\r
 \r
@@ -307,20 +305,44 @@ SetEnvironmentVariables(
     }\r
     ASSERT(StrStr(CurrentString, L"=") != NULL);\r
     Node = AllocateZeroPool(sizeof(ENV_VAR_LIST));\r
-    ASSERT(Node != NULL);\r
+    if (Node == NULL) {\r
+      SetEnvironmentVariableList(&VarList->Link);\r
+      return (EFI_OUT_OF_RESOURCES);\r
+    }\r
+\r
     Node->Key = AllocateZeroPool((StrStr(CurrentString, L"=") - CurrentString + 1) * sizeof(CHAR16));\r
-    ASSERT(Node->Key != NULL);\r
+    if (Node->Key == NULL) {\r
+      SHELL_FREE_NON_NULL(Node);\r
+      SetEnvironmentVariableList(&VarList->Link);\r
+      return (EFI_OUT_OF_RESOURCES);\r
+    }\r
+\r
+    //\r
+    // Copy the string into the Key, leaving the last character allocated as NULL to terminate\r
+    //\r
     StrnCpy(Node->Key, CurrentString, StrStr(CurrentString, L"=") - CurrentString);\r
-    NewSize = StrSize(CurrentString);\r
-    NewSize -= StrLen(Node->Key) - 1;\r
-    Node->Val = AllocateZeroPool(NewSize);\r
-    ASSERT(Node->Val != NULL);\r
-    StrCpy(Node->Val, CurrentString + StrLen(Node->Key) + 1);\r
+\r
+    //\r
+    // ValueSize = TotalSize - already removed size - size for '=' + size for terminator (the last 2 items cancel each other)\r
+    //\r
+    Node->Val = AllocateCopyPool(StrSize(CurrentString) - StrSize(Node->Key), CurrentString + StrLen(Node->Key) + 1);\r
+    if (Node->Val == NULL) {\r
+      SHELL_FREE_NON_NULL(Node->Key);\r
+      SHELL_FREE_NON_NULL(Node);\r
+      SetEnvironmentVariableList(&VarList->Link);\r
+      return (EFI_OUT_OF_RESOURCES);\r
+    }\r
+\r
     Node->Atts = EFI_VARIABLE_BOOTSERVICE_ACCESS;\r
 \r
     if (VarList == NULL) {\r
       VarList = AllocateZeroPool(sizeof(ENV_VAR_LIST));\r
-      ASSERT(VarList != NULL);\r
+      if (VarList == NULL) {\r
+        SHELL_FREE_NON_NULL(Node->Key);\r
+        SHELL_FREE_NON_NULL(Node->Val);\r
+        SHELL_FREE_NON_NULL(Node);\r
+        return (EFI_OUT_OF_RESOURCES);\r
+      }\r
       InitializeListHead(&VarList->Link);\r
     }\r
     InsertTailList(&VarList->Link, &Node->Link);\r