/** @file
  function declarations for shell environment functions.

  Copyright (c) 2009 - 2018, Intel Corporation. All rights reserved.<BR>
  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include "Shell.h"

#define INIT_NAME_BUFFER_SIZE  128
#define INIT_DATA_BUFFER_SIZE  1024

//
// The list is used to cache the environment variables.
//
ENV_VAR_LIST  gShellEnvVarList;

/**
  Reports whether an environment variable is Volatile or Non-Volatile.

  @param EnvVarName             The name of the environment variable in question
  @param Volatile               Return TRUE if the environment variable is volatile

  @retval EFI_SUCCESS           The volatile attribute is returned successfully
  @retval others                Some errors happened.
**/
EFI_STATUS
IsVolatileEnv (
  IN CONST CHAR16  *EnvVarName,
  OUT BOOLEAN      *Volatile
  )
{
  EFI_STATUS  Status;
  UINTN       Size;
  VOID        *Buffer;
  UINT32      Attribs;

  ASSERT (Volatile != NULL);

  Size   = 0;
  Buffer = NULL;

  //
  // get the variable
  //
  Status = gRT->GetVariable (
                  (CHAR16 *)EnvVarName,
                  &gShellVariableGuid,
                  &Attribs,
                  &Size,
                  Buffer
                  );
  if (Status == EFI_BUFFER_TOO_SMALL) {
    Buffer = AllocateZeroPool (Size);
    if (Buffer == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }

    Status = gRT->GetVariable (
                    (CHAR16 *)EnvVarName,
                    &gShellVariableGuid,
                    &Attribs,
                    &Size,
                    Buffer
                    );
    FreePool (Buffer);
  }

  //
  // not found means volatile
  //
  if (Status == EFI_NOT_FOUND) {
    *Volatile = TRUE;
    return EFI_SUCCESS;
  }

  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // check for the Non Volatile bit
  //
  *Volatile = !(BOOLEAN)((Attribs & EFI_VARIABLE_NON_VOLATILE) == EFI_VARIABLE_NON_VOLATILE);
  return EFI_SUCCESS;
}

/**
  free function for ENV_VAR_LIST objects.

  @param[in] List               The pointer to pointer to list.
**/
VOID
FreeEnvironmentVariableList (
  IN LIST_ENTRY  *List
  )
{
  ENV_VAR_LIST  *Node;

  ASSERT (List != NULL);
  if (List == NULL) {
    return;
  }

  for ( Node = (ENV_VAR_LIST *)GetFirstNode (List)
        ; !IsListEmpty (List)
        ; Node = (ENV_VAR_LIST *)GetFirstNode (List)
        )
  {
    ASSERT (Node != NULL);
    RemoveEntryList (&Node->Link);
    if (Node->Key != NULL) {
      FreePool (Node->Key);
    }

    if (Node->Val != NULL) {
      FreePool (Node->Val);
    }

    FreePool (Node);
  }
}

/**
  Creates a list of all Shell-Guid-based environment variables.

  @param[in, out] ListHead       The pointer to pointer to LIST ENTRY object for
                                 storing this list.

  @retval EFI_SUCCESS           the list was created successfully.
**/
EFI_STATUS
GetEnvironmentVariableList (
  IN OUT LIST_ENTRY  *ListHead
  )
{
  CHAR16        *VariableName;
  UINTN         NameSize;
  UINTN         NameBufferSize;
  EFI_STATUS    Status;
  EFI_GUID      Guid;
  UINTN         ValSize;
  UINTN         ValBufferSize;
  ENV_VAR_LIST  *VarList;

  if (ListHead == NULL) {
    return (EFI_INVALID_PARAMETER);
  }

  Status = EFI_SUCCESS;

  ValBufferSize  = INIT_DATA_BUFFER_SIZE;
  NameBufferSize = INIT_NAME_BUFFER_SIZE;
  VariableName   = AllocateZeroPool (NameBufferSize);
  if (VariableName == NULL) {
    return (EFI_OUT_OF_RESOURCES);
  }

  *VariableName = CHAR_NULL;

  while (!EFI_ERROR (Status)) {
    NameSize = NameBufferSize;
    Status   = gRT->GetNextVariableName (&NameSize, VariableName, &Guid);
    if (Status == EFI_NOT_FOUND) {
      Status = EFI_SUCCESS;
      break;
    } else if (Status == EFI_BUFFER_TOO_SMALL) {
      NameBufferSize = NameSize > NameBufferSize * 2 ? NameSize : NameBufferSize * 2;
      SHELL_FREE_NON_NULL (VariableName);
      VariableName = AllocateZeroPool (NameBufferSize);
      if (VariableName == NULL) {
        Status = EFI_OUT_OF_RESOURCES;
        break;
      }

      NameSize = NameBufferSize;
      Status   = gRT->GetNextVariableName (&NameSize, VariableName, &Guid);
    }

    if (!EFI_ERROR (Status) && CompareGuid (&Guid, &gShellVariableGuid)) {
      VarList = AllocateZeroPool (sizeof (ENV_VAR_LIST));
      if (VarList == NULL) {
        Status = EFI_OUT_OF_RESOURCES;
      } else {
        ValSize = ValBufferSize;
        //
        // We need another CHAR16 to save '\0' in VarList->Val.
        //
        VarList->Val = AllocateZeroPool (ValSize + sizeof (CHAR16));
        if (VarList->Val == NULL) {
          SHELL_FREE_NON_NULL (VarList);
          Status = EFI_OUT_OF_RESOURCES;
          break;
        }

        Status = SHELL_GET_ENVIRONMENT_VARIABLE_AND_ATTRIBUTES (VariableName, &VarList->Atts, &ValSize, VarList->Val);
        if (Status == EFI_BUFFER_TOO_SMALL) {
          ValBufferSize = ValSize > ValBufferSize * 2 ? ValSize : ValBufferSize * 2;
          SHELL_FREE_NON_NULL (VarList->Val);
          //
          // We need another CHAR16 to save '\0' in VarList->Val.
          //
          VarList->Val = AllocateZeroPool (ValBufferSize + sizeof (CHAR16));
          if (VarList->Val == NULL) {
            SHELL_FREE_NON_NULL (VarList);
            Status = EFI_OUT_OF_RESOURCES;
            break;
          }

          ValSize = ValBufferSize;
          Status  = SHELL_GET_ENVIRONMENT_VARIABLE_AND_ATTRIBUTES (VariableName, &VarList->Atts, &ValSize, VarList->Val);
        }

        if (!EFI_ERROR (Status)) {
          VarList->Key = AllocateCopyPool (StrSize (VariableName), VariableName);
          if (VarList->Key == NULL) {
            SHELL_FREE_NON_NULL (VarList->Val);
            SHELL_FREE_NON_NULL (VarList);
            Status = EFI_OUT_OF_RESOURCES;
          } else {
            InsertTailList (ListHead, &VarList->Link);
          }
        } else {
          SHELL_FREE_NON_NULL (VarList->Val);
          SHELL_FREE_NON_NULL (VarList);
        }
      } // if (VarList == NULL) ... else ...
    } // compare guid
  } // while

  SHELL_FREE_NON_NULL (VariableName);

  if (EFI_ERROR (Status)) {
    FreeEnvironmentVariableList (ListHead);
  }

  return (Status);
}

/**
  Sets a list of all Shell-Guid-based environment variables.  this will
  also eliminate all existing shell environment variables (even if they
  are not on the list).

  This function will also deallocate the memory from List.

  @param[in] ListHead           The pointer to LIST_ENTRY from
                                GetShellEnvVarList().

  @retval EFI_SUCCESS           the list was Set successfully.
**/
EFI_STATUS
SetEnvironmentVariableList (
  IN LIST_ENTRY  *ListHead
  )
{
  ENV_VAR_LIST  VarList;
  ENV_VAR_LIST  *Node;
  EFI_STATUS    Status;
  UINTN         Size;

  InitializeListHead (&VarList.Link);

  //
  // Delete all the current environment variables
  //
  Status = GetEnvironmentVariableList (&VarList.Link);
  ASSERT_EFI_ERROR (Status);

  for ( Node = (ENV_VAR_LIST *)GetFirstNode (&VarList.Link)
        ; !IsNull (&VarList.Link, &Node->Link)
        ; Node = (ENV_VAR_LIST *)GetNextNode (&VarList.Link, &Node->Link)
        )
  {
    if (Node->Key != NULL) {
      Status = SHELL_DELETE_ENVIRONMENT_VARIABLE (Node->Key);
    }

    ASSERT_EFI_ERROR (Status);
  }

  FreeEnvironmentVariableList (&VarList.Link);

  //
  // set all the variables from the list
  //
  for ( Node = (ENV_VAR_LIST *)GetFirstNode (ListHead)
        ; !IsNull (ListHead, &Node->Link)
        ; Node = (ENV_VAR_LIST *)GetNextNode (ListHead, &Node->Link)
        )
  {
    Size = StrSize (Node->Val) - sizeof (CHAR16);
    if (Node->Atts & EFI_VARIABLE_NON_VOLATILE) {
      Status = SHELL_SET_ENVIRONMENT_VARIABLE_NV (Node->Key, Size, Node->Val);
    } else {
      Status = SHELL_SET_ENVIRONMENT_VARIABLE_V (Node->Key, Size, Node->Val);
    }

    ASSERT_EFI_ERROR (Status);
  }

  FreeEnvironmentVariableList (ListHead);

  return (Status);
}

/**
  sets a list of all Shell-Guid-based environment variables.

  @param Environment        Points to a NULL-terminated array of environment
                            variables with the format 'x=y', where x is the
                            environment variable name and y is the value.

  @retval EFI_SUCCESS       The command executed successfully.
  @retval EFI_INVALID_PARAMETER The parameter is invalid.
  @retval EFI_OUT_OF_RESOURCES Out of resources.

  @sa SetEnvironmentVariableList
**/
EFI_STATUS
SetEnvironmentVariables (
  IN CONST CHAR16  **Environment
  )
{
  CONST CHAR16  *CurrentString;
  UINTN         CurrentCount;
  ENV_VAR_LIST  *VarList;
  ENV_VAR_LIST  *Node;

  VarList = NULL;

  if (Environment == NULL) {
    return (EFI_INVALID_PARAMETER);
  }

  //
  // Build a list identical to the ones used for get/set list functions above
  //
  for ( CurrentCount = 0
        ;
        ; CurrentCount++
        )
  {
    CurrentString = Environment[CurrentCount];
    if (CurrentString == NULL) {
      break;
    }

    ASSERT (StrStr (CurrentString, L"=") != NULL);
    Node = AllocateZeroPool (sizeof (ENV_VAR_LIST));
    if (Node == NULL) {
      SetEnvironmentVariableList (&VarList->Link);
      return (EFI_OUT_OF_RESOURCES);
    }

    Node->Key = AllocateZeroPool ((StrStr (CurrentString, L"=") - CurrentString + 1) * sizeof (CHAR16));
    if (Node->Key == NULL) {
      SHELL_FREE_NON_NULL (Node);
      SetEnvironmentVariableList (&VarList->Link);
      return (EFI_OUT_OF_RESOURCES);
    }

    //
    // Copy the string into the Key, leaving the last character allocated as NULL to terminate
    //
    StrnCpyS (
      Node->Key,
      StrStr (CurrentString, L"=") - CurrentString + 1,
      CurrentString,
      StrStr (CurrentString, L"=") - CurrentString
      );

    //
    // ValueSize = TotalSize - already removed size - size for '=' + size for terminator (the last 2 items cancel each other)
    //
    Node->Val = AllocateCopyPool (StrSize (CurrentString) - StrSize (Node->Key), CurrentString + StrLen (Node->Key) + 1);
    if (Node->Val == NULL) {
      SHELL_FREE_NON_NULL (Node->Key);
      SHELL_FREE_NON_NULL (Node);
      SetEnvironmentVariableList (&VarList->Link);
      return (EFI_OUT_OF_RESOURCES);
    }

    Node->Atts = EFI_VARIABLE_BOOTSERVICE_ACCESS;

    if (VarList == NULL) {
      VarList = AllocateZeroPool (sizeof (ENV_VAR_LIST));
      if (VarList == NULL) {
        SHELL_FREE_NON_NULL (Node->Key);
        SHELL_FREE_NON_NULL (Node->Val);
        SHELL_FREE_NON_NULL (Node);
        return (EFI_OUT_OF_RESOURCES);
      }

      InitializeListHead (&VarList->Link);
    }

    InsertTailList (&VarList->Link, &Node->Link);
  } // for loop

  //
  // set this new list as the set of all environment variables.
  // this function also frees the memory and deletes all pre-existing
  // shell-guid based environment variables.
  //
  return (SetEnvironmentVariableList (&VarList->Link));
}

/**
  Find an environment variable in the gShellEnvVarList.

  @param Key        The name of the environment variable.
  @param Value      The value of the environment variable, the buffer
                    shoule be freed by the caller.
  @param ValueSize  The size in bytes of the environment variable
                    including the tailing CHAR_NELL.
  @param Atts       The attributes of the variable.

  @retval EFI_SUCCESS       The command executed successfully.
  @retval EFI_NOT_FOUND     The environment variable is not found in
                            gShellEnvVarList.

**/
EFI_STATUS
ShellFindEnvVarInList (
  IN  CONST CHAR16  *Key,
  OUT CHAR16        **Value,
  OUT UINTN         *ValueSize,
  OUT UINT32        *Atts OPTIONAL
  )
{
  ENV_VAR_LIST  *Node;

  if ((Key == NULL) || (Value == NULL) || (ValueSize == NULL)) {
    return SHELL_INVALID_PARAMETER;
  }

  for ( Node = (ENV_VAR_LIST *)GetFirstNode (&gShellEnvVarList.Link)
        ; !IsNull (&gShellEnvVarList.Link, &Node->Link)
        ; Node = (ENV_VAR_LIST *)GetNextNode (&gShellEnvVarList.Link, &Node->Link)
        )
  {
    if ((Node->Key != NULL) && (StrCmp (Key, Node->Key) == 0)) {
      *Value     = AllocateCopyPool (StrSize (Node->Val), Node->Val);
      *ValueSize = StrSize (Node->Val);
      if (Atts != NULL) {
        *Atts = Node->Atts;
      }

      return EFI_SUCCESS;
    }
  }

  return EFI_NOT_FOUND;
}

/**
  Add an environment variable into gShellEnvVarList.

  @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_NULL
  @param Atts       The attributes of the variable.

  @retval EFI_SUCCESS  The environment variable was added to list successfully.
  @retval others       Some errors happened.

**/
EFI_STATUS
ShellAddEnvVarToList (
  IN CONST CHAR16  *Key,
  IN CONST CHAR16  *Value,
  IN UINTN         ValueSize,
  IN UINT32        Atts
  )
{
  ENV_VAR_LIST  *Node;
  CHAR16        *LocalKey;
  CHAR16        *LocalValue;

  if ((Key == NULL) || (Value == NULL) || (ValueSize == 0)) {
    return EFI_INVALID_PARAMETER;
  }

  LocalValue = AllocateCopyPool (ValueSize, Value);
  if (LocalValue == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  //
  // Update the variable value if it exists in gShellEnvVarList.
  //
  for ( Node = (ENV_VAR_LIST *)GetFirstNode (&gShellEnvVarList.Link)
        ; !IsNull (&gShellEnvVarList.Link, &Node->Link)
        ; Node = (ENV_VAR_LIST *)GetNextNode (&gShellEnvVarList.Link, &Node->Link)
        )
  {
    if ((Node->Key != NULL) && (StrCmp (Key, Node->Key) == 0)) {
      Node->Atts = Atts;
      SHELL_FREE_NON_NULL (Node->Val);
      Node->Val = LocalValue;
      return EFI_SUCCESS;
    }
  }

  //
  // If the environment variable 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));
  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 EFI_SUCCESS;
}

/**
  Remove a specified environment variable in gShellEnvVarList.

  @param Key        The name of the environment variable.

  @retval EFI_SUCCESS       The command executed successfully.
  @retval EFI_NOT_FOUND     The environment variable is not found in
                            gShellEnvVarList.
**/
EFI_STATUS
ShellRemvoeEnvVarFromList (
  IN CONST CHAR16  *Key
  )
{
  ENV_VAR_LIST  *Node;

  if (Key == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  for ( Node = (ENV_VAR_LIST *)GetFirstNode (&gShellEnvVarList.Link)
        ; !IsNull (&gShellEnvVarList.Link, &Node->Link)
        ; Node = (ENV_VAR_LIST *)GetNextNode (&gShellEnvVarList.Link, &Node->Link)
        )
  {
    if ((Node->Key != NULL) && (StrCmp (Key, Node->Key) == 0)) {
      SHELL_FREE_NON_NULL (Node->Key);
      SHELL_FREE_NON_NULL (Node->Val);
      RemoveEntryList (&Node->Link);
      SHELL_FREE_NON_NULL (Node);
      return EFI_SUCCESS;
    }
  }

  return EFI_NOT_FOUND;
}

/**
  Initialize the gShellEnvVarList and cache all Shell-Guid-based environment
  variables.

**/
EFI_STATUS
ShellInitEnvVarList (
  VOID
  )
{
  EFI_STATUS  Status;

  InitializeListHead (&gShellEnvVarList.Link);
  Status = GetEnvironmentVariableList (&gShellEnvVarList.Link);

  return Status;
}

/**
  Destructe the gShellEnvVarList.

**/
VOID
ShellFreeEnvVarList (
  VOID
  )
{
  FreeEnvironmentVariableList (&gShellEnvVarList.Link);
  InitializeListHead (&gShellEnvVarList.Link);

  return;
}