Contributed-under: TianoCore Contribution Agreement 1.0 Signed-off-by: Ruiyu Ni <ruiyu...@intel.com> Cc: Jaben Carsey <jaben.car...@intel.com> --- ShellPkg/Application/Shell/ShellEnvVar.c | 39 +++++++++++++++++++++--------- ShellPkg/Application/Shell/ShellEnvVar.h | 7 ++++-- ShellPkg/Application/Shell/ShellProtocol.c | 30 +++++++++-------------- 3 files changed, 43 insertions(+), 33 deletions(-)
diff --git a/ShellPkg/Application/Shell/ShellEnvVar.c b/ShellPkg/Application/Shell/ShellEnvVar.c index 9f87b90..77bf599 100644 --- a/ShellPkg/Application/Shell/ShellEnvVar.c +++ b/ShellPkg/Application/Shell/ShellEnvVar.c @@ -442,8 +442,11 @@ ShellFindEnvVarInList ( including the tailing CHAR_NULL @param Atts The attributes of the variable. + @retval EFI_SUCCESS The environment variable was added to list successfully. + @retval others Some errors happened. + **/ -VOID +EFI_STATUS ShellAddEnvVarToList ( IN CONST CHAR16 *Key, IN CONST CHAR16 *Value, @@ -452,9 +455,16 @@ ShellAddEnvVarToList ( ) { ENV_VAR_LIST *Node; + CHAR16 *LocalKey; + CHAR16 *LocalValue; if (Key == NULL || Value == NULL || ValueSize == 0) { - return; + return EFI_INVALID_PARAMETER; + } + + LocalValue = AllocateCopyPool (ValueSize, Value); + if (LocalValue == NULL) { + return EFI_OUT_OF_RESOURCES; } // @@ -467,10 +477,8 @@ ShellAddEnvVarToList ( if (Node->Key != NULL && StrCmp(Key, Node->Key) == 0) { Node->Atts = Atts; SHELL_FREE_NON_NULL(Node->Val); - Node->Val = AllocateZeroPool (ValueSize); - ASSERT (Node->Val != NULL); - CopyMem(Node->Val, Value, ValueSize); - return; + Node->Val = LocalValue; + return EFI_SUCCESS; } } @@ -478,16 +486,23 @@ ShellAddEnvVarToList ( // If the environment varialbe key doesn't exist in list just insert // a new node. // + LocalKey = AllocateCopyPool (StrSize(Key), Key); + if (LocalKey == NULL) { + FreePool (LocalValue); + return EFI_OUT_OF_RESOURCES; + } Node = (ENV_VAR_LIST*)AllocateZeroPool (sizeof(ENV_VAR_LIST)); - ASSERT (Node != NULL); - Node->Key = AllocateCopyPool(StrSize(Key), Key); - ASSERT (Node->Key != NULL); - Node->Val = AllocateCopyPool(ValueSize, Value); - ASSERT (Node->Val != NULL); + if (Node == NULL) { + FreePool (LocalKey); + FreePool (LocalValue); + return EFI_OUT_OF_RESOURCES; + } + Node->Key = LocalKey; + Node->Val = LocalValue; Node->Atts = Atts; InsertTailList(&gShellEnvVarList.Link, &Node->Link); - return; + return EFI_SUCCESS; } /** diff --git a/ShellPkg/Application/Shell/ShellEnvVar.h b/ShellPkg/Application/Shell/ShellEnvVar.h index 5356580..f05ef2b 100644 --- a/ShellPkg/Application/Shell/ShellEnvVar.h +++ b/ShellPkg/Application/Shell/ShellEnvVar.h @@ -240,11 +240,14 @@ ShellFindEnvVarInList ( @param Key The name of the environment variable. @param Value The value of environment variable. @param ValueSize The size in bytes of the environment variable - including the tailing CHAR_NELL + including the tailing CHAR_NULL @param Atts The attributes of the variable. + @retval EFI_SUCCESS The environment variable was added to list successfully. + @retval others Some errors happened. + **/ -VOID +EFI_STATUS ShellAddEnvVarToList ( IN CONST CHAR16 *Key, IN CONST CHAR16 *Value, diff --git a/ShellPkg/Application/Shell/ShellProtocol.c b/ShellPkg/Application/Shell/ShellProtocol.c index 39b0e78..a95da00 100644 --- a/ShellPkg/Application/Shell/ShellProtocol.c +++ b/ShellPkg/Application/Shell/ShellProtocol.c @@ -2852,36 +2852,28 @@ InternalEfiShellSetEnv( ) { EFI_STATUS Status; - UINT32 Atts; - Atts = 0x0; - if (Value == NULL || StrLen(Value) == 0) { Status = SHELL_DELETE_ENVIRONMENT_VARIABLE(Name); if (!EFI_ERROR(Status)) { ShellRemvoeEnvVarFromList(Name); } - return Status; } else { SHELL_DELETE_ENVIRONMENT_VARIABLE(Name); - if (Volatile) { - Status = SHELL_SET_ENVIRONMENT_VARIABLE_V(Name, StrSize(Value), Value); - if (!EFI_ERROR(Status)) { - Atts &= ~EFI_VARIABLE_NON_VOLATILE; - Atts |= EFI_VARIABLE_BOOTSERVICE_ACCESS; - ShellAddEnvVarToList(Name, Value, StrSize(Value), Atts); + Status = ShellAddEnvVarToList( + Name, Value, StrSize(Value), + EFI_VARIABLE_BOOTSERVICE_ACCESS | (Volatile ? 0 : EFI_VARIABLE_NON_VOLATILE) + ); + if (!EFI_ERROR (Status)) { + Status = Volatile + ? SHELL_SET_ENVIRONMENT_VARIABLE_V(Name, StrSize(Value), Value) + : SHELL_SET_ENVIRONMENT_VARIABLE_NV(Name, StrSize(Value), Value); + if (EFI_ERROR (Status)) { + ShellRemvoeEnvVarFromList(Name); } - return Status; - } else { - Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV(Name, StrSize(Value), Value); - if (!EFI_ERROR(Status)) { - Atts |= EFI_VARIABLE_NON_VOLATILE; - Atts |= EFI_VARIABLE_BOOTSERVICE_ACCESS; - ShellAddEnvVarToList(Name, Value, StrSize(Value), Atts); - } - return Status; } } + return Status; } /** -- 2.9.0.windows.1 _______________________________________________ edk2-devel mailing list edk2-devel@lists.01.org https://lists.01.org/mailman/listinfo/edk2-devel