ShellPkg/Application: Fix ">v" cannot update environment variable

When ">v" is used to redirect the command output to environment
variable (e.g.: "echo xxx >v yyy"), we only called SetVariable()
to update the variable storage but forgot to update the cached
environment variables in gShellEnvVarList.
When updating the variable storage, the existing code unnecessary
saved the ending NULL character into variable storage.

The patch fixes all the above issues.

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Chen A Chen <chen.a.chen@intel.com>
Reviewed-by: Ruiyu Ni <ruiyu.ni@intel.com>
Reviewed-by: Jaben Carsey <jaben.carsey@intel.com>
Reviewed-by: Tapan Shah <tapandshah@hpe.com>

diff --git a/ShellPkg/Application/Shell/FileHandleWrappers.c b/ShellPkg/Application/Shell/FileHandleWrappers.c
index 18f4169c7d3ae899fe0bdd7647113aaa9c8e96f0..3c11d82944e792ab0fda2474f1b2dbf22940e4dc 100644 (file)
--- a/ShellPkg/Application/Shell/FileHandleWrappers.c
+++ b/ShellPkg/Application/Shell/FileHandleWrappers.c
@@ -1040,6 +1040,7 @@ FileInterfaceEnvClose(
   UINTN       NewSize;\r
   EFI_STATUS  Status;\r
   BOOLEAN     Volatile;\r
+  UINTN       TotalSize;\r
 \r
   //\r
   // Most if not all UEFI commands will have an '\r\n' at the end of any output. \r
@@ -1049,6 +1050,7 @@ FileInterfaceEnvClose(
   //\r
   NewBuffer   = NULL;\r
   NewSize     = 0;\r
+  TotalSize   = 0;\r
 \r
   Status = IsVolatileEnv (((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &Volatile);\r
   if (EFI_ERROR (Status)) {\r
@@ -1057,7 +1059,8 @@ FileInterfaceEnvClose(
 \r
   Status = SHELL_GET_ENVIRONMENT_VARIABLE(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &NewSize, NewBuffer);\r
   if (Status == EFI_BUFFER_TOO_SMALL) {\r
-    NewBuffer = AllocateZeroPool(NewSize + sizeof(CHAR16));\r
+    TotalSize = NewSize + sizeof (CHAR16);\r
+    NewBuffer = AllocateZeroPool (TotalSize);\r
     if (NewBuffer == NULL) {\r
       return EFI_OUT_OF_RESOURCES;\r
     }  \r
@@ -1066,17 +1069,43 @@ FileInterfaceEnvClose(
   \r
   if (!EFI_ERROR(Status) && NewBuffer != NULL) {\r
     \r
-    if (StrSize(NewBuffer) > 6)\r
-    {\r
-      if ((((CHAR16*)NewBuffer)[(StrSize(NewBuffer)/2) - 2] == CHAR_LINEFEED) \r
-           && (((CHAR16*)NewBuffer)[(StrSize(NewBuffer)/2) - 3] == CHAR_CARRIAGE_RETURN)) {\r
-        ((CHAR16*)NewBuffer)[(StrSize(NewBuffer)/2) - 3] = CHAR_NULL;   \r
+    if (TotalSize / sizeof (CHAR16) >= 3) {\r
+      if ( (((CHAR16*)NewBuffer)[TotalSize / sizeof (CHAR16) - 2] == CHAR_LINEFEED) &&\r
+           (((CHAR16*)NewBuffer)[TotalSize / sizeof (CHAR16) - 3] == CHAR_CARRIAGE_RETURN)\r
+         ) {\r
+        ((CHAR16*)NewBuffer)[TotalSize / sizeof (CHAR16) - 3] = CHAR_NULL;\r
       }\r
 \r
       if (Volatile) {\r
-        Status = SHELL_SET_ENVIRONMENT_VARIABLE_V(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, StrSize(NewBuffer), NewBuffer);\r
+        Status = SHELL_SET_ENVIRONMENT_VARIABLE_V (\r
+                   ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
+                   TotalSize - sizeof (CHAR16),\r
+                   NewBuffer\r
+                   );\r
+\r
+        if (!EFI_ERROR(Status)) {\r
+          Status = ShellAddEnvVarToList (\r
+                     ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
+                     NewBuffer,\r
+                     TotalSize,\r
+                     EFI_VARIABLE_BOOTSERVICE_ACCESS\r
+                     );\r
+        }\r
       } else {\r
-        Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, StrSize(NewBuffer), NewBuffer);\r
+        Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV (\r
+                   ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
+                   TotalSize - sizeof (CHAR16),\r
+                   NewBuffer\r
+                   );\r
+\r
+        if (!EFI_ERROR(Status)) {\r
+          Status = ShellAddEnvVarToList (\r
+                     ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
+                     NewBuffer,\r
+                     TotalSize,\r
+                     EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS\r
+                     );\r
+        }\r
       }\r
     }\r
   } \r
@@ -1128,12 +1157,14 @@ FileInterfaceEnvRead(
 \r
 /**\r
   File style interface for Volatile Environment Variable (Write).\r
+  This function also caches the environment variable into gShellEnvVarList.\r
   \r
   @param[in] This              The pointer to the EFI_FILE_PROTOCOL object.\r
   @param[in, out] BufferSize   Size in bytes of Buffer.\r
   @param[in] Buffer            The pointer to the buffer to write.\r
   \r
-  @retval EFI_SUCCESS   The data was read.\r
+  @retval EFI_SUCCESS             The data was successfully write to variable.\r
+  @retval SHELL_OUT_OF_RESOURCES  A memory allocation failed.\r
 **/\r
 EFI_STATUS\r
 EFIAPI\r
@@ -1146,41 +1177,71 @@ FileInterfaceEnvVolWrite(
   VOID*       NewBuffer;\r
   UINTN       NewSize;\r
   EFI_STATUS  Status;\r
+  UINTN       TotalSize;\r
 \r
   NewBuffer   = NULL;\r
   NewSize     = 0;\r
+  TotalSize   = 0;\r
 \r
   Status = SHELL_GET_ENVIRONMENT_VARIABLE(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &NewSize, NewBuffer);\r
-  if (Status == EFI_BUFFER_TOO_SMALL){\r
-    NewBuffer = AllocateZeroPool(NewSize + *BufferSize + sizeof(CHAR16));\r
+  if (Status == EFI_BUFFER_TOO_SMALL) {\r
+    TotalSize = NewSize + *BufferSize + sizeof (CHAR16);\r
+  } else if (Status == EFI_NOT_FOUND) {\r
+    TotalSize = *BufferSize + sizeof(CHAR16);\r
+  } else {\r
+    return Status;\r
+  }\r
+\r
+  NewBuffer = AllocateZeroPool (TotalSize);\r
+  if (NewBuffer == NULL) {\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+\r
+  if (Status == EFI_BUFFER_TOO_SMALL) {\r
     Status = SHELL_GET_ENVIRONMENT_VARIABLE(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &NewSize, NewBuffer);\r
   }\r
-  if (!EFI_ERROR(Status) && NewBuffer != NULL) {\r
-    while (((CHAR16*)NewBuffer)[NewSize/2] == CHAR_NULL) {\r
-      //\r
-      // We want to overwrite the CHAR_NULL\r
-      //\r
-      NewSize -= 2;\r
-    }\r
-    CopyMem((UINT8*)NewBuffer + NewSize + 2, Buffer, *BufferSize);\r
-    Status = SHELL_SET_ENVIRONMENT_VARIABLE_V(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, StrSize(NewBuffer), NewBuffer);\r
-    FreePool(NewBuffer);\r
-    return (Status);\r
-  } else {\r
-    SHELL_FREE_NON_NULL(NewBuffer);\r
-    return (SHELL_SET_ENVIRONMENT_VARIABLE_V(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, *BufferSize, Buffer));\r
+\r
+  if (EFI_ERROR (Status) && Status != EFI_NOT_FOUND) {\r
+    FreePool (NewBuffer);\r
+    return Status;\r
+  }\r
+\r
+  CopyMem ((UINT8*)NewBuffer + NewSize, Buffer, *BufferSize);\r
+  Status = ShellAddEnvVarToList (\r
+             ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
+             NewBuffer,\r
+             TotalSize,\r
+             EFI_VARIABLE_BOOTSERVICE_ACCESS\r
+             );\r
+  if (EFI_ERROR(Status)) {\r
+    FreePool (NewBuffer);\r
+    return Status;\r
+  }\r
+\r
+  Status = SHELL_SET_ENVIRONMENT_VARIABLE_V (\r
+             ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
+             TotalSize - sizeof (CHAR16),\r
+             NewBuffer\r
+             );\r
+  if (EFI_ERROR(Status)) {\r
+    ShellRemvoeEnvVarFromList (((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name);\r
   }\r
+\r
+  FreePool (NewBuffer);\r
+  return Status;\r
 }\r
 \r
 \r
 /**\r
   File style interface for Non Volatile Environment Variable (Write).\r
+  This function also caches the environment variable into gShellEnvVarList.\r
   \r
   @param[in] This              The pointer to the EFI_FILE_PROTOCOL object.\r
   @param[in, out] BufferSize   Size in bytes of Buffer.\r
   @param[in] Buffer            The pointer to the buffer to write.\r
   \r
-  @retval EFI_SUCCESS   The data was read.\r
+  @retval EFI_SUCCESS             The data was successfully write to variable.\r
+  @retval SHELL_OUT_OF_RESOURCES  A memory allocation failed.\r
 **/\r
 EFI_STATUS\r
 EFIAPI\r
@@ -1193,27 +1254,58 @@ FileInterfaceEnvNonVolWrite(
   VOID*       NewBuffer;\r
   UINTN       NewSize;\r
   EFI_STATUS  Status;\r
+  UINTN       TotalSize;\r
 \r
   NewBuffer   = NULL;\r
   NewSize     = 0;\r
+  TotalSize   = 0;\r
 \r
   Status = SHELL_GET_ENVIRONMENT_VARIABLE(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &NewSize, NewBuffer);\r
-  if (Status == EFI_BUFFER_TOO_SMALL){\r
-    NewBuffer = AllocateZeroPool(NewSize + *BufferSize);\r
+  if (Status == EFI_BUFFER_TOO_SMALL) {\r
+    TotalSize = NewSize + *BufferSize + sizeof (CHAR16);\r
+  } else if (Status == EFI_NOT_FOUND) {\r
+    TotalSize = *BufferSize + sizeof (CHAR16);\r
+  } else {\r
+    return Status;\r
+  }\r
+\r
+  NewBuffer = AllocateZeroPool (TotalSize);\r
+  if (NewBuffer == NULL) {\r
+    return EFI_OUT_OF_RESOURCES;\r
+  }\r
+\r
+  if (Status == EFI_BUFFER_TOO_SMALL) {\r
     Status = SHELL_GET_ENVIRONMENT_VARIABLE(((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name, &NewSize, NewBuffer);\r
   }\r
-  if (!EFI_ERROR(Status)) {\r
-    CopyMem((UINT8*)NewBuffer + NewSize, Buffer, *BufferSize);\r
-    return (SHELL_SET_ENVIRONMENT_VARIABLE_NV(\r
-    ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
-    NewSize + *BufferSize,\r
-    NewBuffer));\r
-  } else {\r
-    return (SHELL_SET_ENVIRONMENT_VARIABLE_NV(\r
-    ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
-    *BufferSize,\r
-    Buffer));\r
+\r
+  if (EFI_ERROR(Status) && Status != EFI_NOT_FOUND) {\r
+    FreePool (NewBuffer);\r
+    return Status;\r
+  }\r
+\r
+  CopyMem ((UINT8*) NewBuffer + NewSize, Buffer, *BufferSize);\r
+  Status = ShellAddEnvVarToList (\r
+             ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
+             NewBuffer,\r
+             TotalSize,\r
+             EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS\r
+             );\r
+  if (EFI_ERROR (Status)) {\r
+    FreePool (NewBuffer);\r
+    return Status;\r
   }\r
+\r
+  Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV (\r
+             ((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name,\r
+             TotalSize - sizeof (CHAR16),\r
+             NewBuffer\r
+             );\r
+  if (EFI_ERROR (Status)) {\r
+    ShellRemvoeEnvVarFromList (((EFI_FILE_PROTOCOL_ENVIRONMENT*)This)->Name);\r
+  }\r
+\r
+  FreePool (NewBuffer);\r
+  return Status;\r
 }\r
 \r
 /**\r

diff --git a/ShellPkg/Application/Shell/ShellEnvVar.c b/ShellPkg/Application/Shell/ShellEnvVar.c
index 6d8c79b4ca6a21c42a8c597d215d58493e7fa6fa..4c49c1ca980e1df17f9d238974f3ccc46fa040a5 100644 (file)
--- a/ShellPkg/Application/Shell/ShellEnvVar.c
+++ b/ShellPkg/Application/Shell/ShellEnvVar.c
@@ -178,7 +178,10 @@ GetEnvironmentVariableList(
         Status = EFI_OUT_OF_RESOURCES;\r
       } else {\r
         ValSize = ValBufferSize;\r
-        VarList->Val = AllocateZeroPool(ValSize);\r
+        //\r
+        // We need another CHAR16 to save '\0' in VarList->Val.\r
+        //\r
+        VarList->Val = AllocateZeroPool (ValSize + sizeof (CHAR16));\r
         if (VarList->Val == NULL) {\r
             SHELL_FREE_NON_NULL(VarList);\r
             Status = EFI_OUT_OF_RESOURCES;\r
@@ -188,7 +191,10 @@ GetEnvironmentVariableList(
         if (Status == EFI_BUFFER_TOO_SMALL){\r
           ValBufferSize = ValSize > ValBufferSize * 2 ? ValSize : ValBufferSize * 2;\r
           SHELL_FREE_NON_NULL (VarList->Val);\r
-          VarList->Val = AllocateZeroPool(ValBufferSize);\r
+          //\r
+          // We need another CHAR16 to save '\0' in VarList->Val.\r
+          //\r
+          VarList->Val = AllocateZeroPool (ValBufferSize + sizeof (CHAR16));\r
           if (VarList->Val == NULL) {\r
             SHELL_FREE_NON_NULL(VarList);\r
             Status = EFI_OUT_OF_RESOURCES;\r
@@ -272,7 +278,7 @@ SetEnvironmentVariableList(
       ; !IsNull(ListHead, &Node->Link)\r
       ; Node = (ENV_VAR_LIST*)GetNextNode(ListHead, &Node->Link)\r
      ){\r
-    Size = StrSize(Node->Val);\r
+    Size = StrSize (Node->Val) - sizeof (CHAR16);\r
     if (Node->Atts & EFI_VARIABLE_NON_VOLATILE) {\r
       Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV(Node->Key, Size, Node->Val);\r
     } else {\r

diff --git a/ShellPkg/Application/Shell/ShellProtocol.c b/ShellPkg/Application/Shell/ShellProtocol.c
index 04b66c5acaae349977573b93e696d0a6b7e0efa1..12c7c40524ab6062555590d0d62ed351acaa2dfe 100644 (file)
--- a/ShellPkg/Application/Shell/ShellProtocol.c
+++ b/ShellPkg/Application/Shell/ShellProtocol.c
@@ -2872,8 +2872,8 @@ InternalEfiShellSetEnv(
                );\r
     if (!EFI_ERROR (Status)) {\r
       Status = Volatile\r
-             ? SHELL_SET_ENVIRONMENT_VARIABLE_V(Name, StrSize(Value), Value)\r
-             : SHELL_SET_ENVIRONMENT_VARIABLE_NV(Name, StrSize(Value), Value);\r
+             ? SHELL_SET_ENVIRONMENT_VARIABLE_V (Name, StrSize (Value) - sizeof (CHAR16), Value)\r
+             : SHELL_SET_ENVIRONMENT_VARIABLE_NV (Name, StrSize (Value) - sizeof (CHAR16), Value);\r
       if (EFI_ERROR (Status)) {\r
         ShellRemvoeEnvVarFromList(Name);\r
       }\r