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 d674f14c24c30c8050b875bc4b59d6216217ebf2..b94711142af6c3ff4532ed1b6127e2b70eea9e31 100644 (file)
--- a/ShellPkg/Application/Shell/ShellEnvVar.c
+++ b/ShellPkg/Application/Shell/ShellEnvVar.c
@@ -12,18 +12,7 @@
 \r
 **/\r
 \r
-#include <Uefi.h>\r
-#include <ShellBase.h>\r
-\r
-#include <Guid/ShellVariableGuid.h>\r
-\r
-#include <Library/BaseLib.h>\r
-#include <Library/UefiRuntimeServicesTableLib.h>\r
-#include <Library/MemoryAllocationLib.h>\r
-#include <Library/DebugLib.h>\r
-#include <Library/BaseMemoryLib.h>\r
-\r
-#include "ShellEnvVar.h"\r
+#include "Shell.h"\r
 \r
 /**\r
   Reports whether an environment variable is Volatile or Non-Volatile.\r
@@ -105,7 +94,7 @@ FreeEnvironmentVariableList(
   }\r
 \r
   for ( Node = (ENV_VAR_LIST*)GetFirstNode(List)\r
-      ; IsListEmpty(List)\r
+      ; !IsListEmpty(List)\r
       ; Node = (ENV_VAR_LIST*)GetFirstNode(List)\r
      ){\r
     ASSERT(Node != NULL);\r
@@ -123,8 +112,8 @@ FreeEnvironmentVariableList(
 /**\r
   Creates a list of all Shell-Guid-based environment variables.\r
 \r
-  @param[in,out] ListHead       The pointer to pointer to LIST ENTRY object for\r
-                                storing this list.\r
+  @param[in, out] ListHead       The pointer to pointer to LIST ENTRY object for\r
+                                 storing this list.\r
 \r
   @retval EFI_SUCCESS           the list was created sucessfully.\r
 **/\r
@@ -163,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
@@ -189,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
@@ -297,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
@@ -318,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