/** @file
  Implements CSE Variable Storage Services and installs
  an instance of the VariableStorage Runtime DXE protocol.

 @copyright
  INTEL CONFIDENTIAL
  Copyright 2016 - 2018 Intel Corporation.

  The source code contained or described herein and all documents related to the
  source code ("Material") are owned by Intel Corporation or its suppliers or
  licensors. Title to the Material remains with Intel Corporation or its suppliers
  and licensors. The Material may contain trade secrets and proprietary and
  confidential information of Intel Corporation and its suppliers and licensors,
  and is protected by worldwide copyright and trade secret laws and treaty
  provisions. No part of the Material may be used, copied, reproduced, modified,
  published, uploaded, posted, transmitted, distributed, or disclosed in any way
  without Intel's prior express written permission.

  No license under any patent, copyright, trade secret or other intellectual
  property right is granted to or conferred upon you by disclosure or delivery
  of the Materials, either expressly, by implication, inducement, estoppel or
  otherwise. Any license under such intellectual property rights must be
  express and approved by Intel in writing.

  Unless otherwise agreed by Intel in writing, you may not remove or alter
  this notice or any other notice embedded in Materials by Intel or
  Intel's suppliers or licensors in any way.

  This file contains an 'Intel Peripheral Driver' and is uniquely identified as
  "Intel Reference Module" and is licensed for Intel CPUs and chipsets under
  the terms of your license agreement with Intel or your vendor. This file may
  be modified by the user, subject to additional terms of the license agreement.

@par Specification Reference:
**/

#include "CseVariableStorage.h"

#include <CseVariableStorageCommandQueue.h>

#include <Library/BaseMemoryLib.h>
#include <Library/DebugLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/SmmMemLib.h>
#include <Library/SmmServicesTableLib.h>
#include <Library/UefiDriverEntryPoint.h>
#include <Library/VariableNvmStorageLib.h>
#include <Library/SmmServicesTableLib.h>

#include <Protocol/CseVariableStorageSupport.h>
#include <Protocol/SmmSxDispatch2.h>
#include <Protocol/TrustedChannel.h>

//
// Declaration of the SMM Command Queue Buffer
//
UINT8                *mCseCommandQueueBuffer                = NULL;
UINT8                *mCseCommandQueueVerificationBuffer    = NULL;
CONST UINTN          mCseCommandQueueVerificationBufferSize = CMD_QUEUE_SINGLE_CMD_BUFFER_SIZE;
EFI_HANDLE           mCseSmmVariableHandle                  = NULL;
EFI_CIRCULAR_BUFFER  mCseCommandQueue   = { 0, 0, CMD_QUEUE_TOTAL_BUFFER_SIZE, 0, NULL};

HECI2_TRUSTED_CHANNEL_BIOS_HEADER              mGetProxyStateMessageHeader;
HECI_PROXY_STATE                               mHeciProxyState                                  = HeciProxyStatePresent;
BOOLEAN                                        mIsProxyHeaderCreated                            = FALSE;
AUTHENTICATED_VARIABLE_NVM_HEADER              *mCurrentGetVariableHeader                       = NULL;
STATIC EDKII_VARIABLE_STORAGE_SUPPORT_PROTOCOL *mVariableStorageSupportProtocol                 = NULL;
STATIC HECI_TRUSTED_CHANNEL_PROTOCOL           *mHeciTrustedChannelVerificationProtocol         = NULL;
NVM_HEADER_GLOBALS                             mNvmHeaderGlobals;

UINT8                                          *mCseVaraiableDataBuffer                         = NULL;

STATIC CSE_VARIABLE_STORAGE_SUPPORT_PROTOCOL mCseVariableStorageSupportProtocol = {
  VariableTrustedChannelReadMessageComplete,
  VariableTrustedChannelWriteMessageComplete
};

STATIC
VOID
EFIAPI
PrintBinaryBuffer (
  IN        UINT8*      Buffer,
  IN        UINTN       BufferSize
  )
{
  UINTN    CurrentByte = 0;

  if (BufferSize == 0) {
    DEBUG ((EFI_D_INFO, "Skipping print of 0 size buffer\n"));
    return;
  }

  DEBUG ((EFI_D_INFO, "  Base  0  1  2  3  4  5  6  7  8  9  A  B  C  D  E  F\n"));
  DEBUG ((EFI_D_INFO, "  %4d %2x ", CurrentByte / 16, Buffer[0]));

  for (CurrentByte = 1; CurrentByte < BufferSize; CurrentByte++) {
    if ((CurrentByte % 16) == 0) {
      DEBUG ((EFI_D_INFO, "\n  %4d %2x ", CurrentByte / 16, Buffer[CurrentByte]));
    } else {
      DEBUG ((EFI_D_INFO, "%2x ", Buffer[CurrentByte]));
    }
  }
  DEBUG ((EFI_D_INFO, "\n"));
  return;
}

/**
  Updates the proxy state to disabled on S3 entry.

  @param[in]      DispatchHandle   The dispatch handle.
  @param[in]      DispatchContext  A pointer to context for the callback function.
  @param[in]      CommBuffer       A pointer to the SMM communicate buffer.
  @param[in]      CommBufferSize   A pointer to the size of the SMM communicate buffer.

  @retval         EFI_SUCCESS      The proxy state was updated successfully.
**/
EFI_STATUS
EFIAPI
CseVariableStorageS3EntryProxyUpdateCallBack (
  IN  EFI_HANDLE                    DispatchHandle,
  IN  CONST VOID                    *DispatchContext,
  IN  OUT VOID                      *CommBuffer OPTIONAL,
  IN  UINTN                         *CommBufferSize  OPTIONAL
  )
{
  mHeciProxyState       = HeciProxyStatePresent;
  mIsProxyHeaderCreated = FALSE;
  return EFI_SUCCESS;
}

/**
  Updates the variable driver cache with data read from non-volatile storage.

  @param[in]  VariableHeader                  The address of AUTHENTICATED_VARIABLE_NVM_HEADER which has information about the variable
  @param[in]  Data                            The address of the data to be updated in cache
  @param[in]  DataSize                        The size in bytes of the data.

  @retval EFI_SUCCESS                         The variable cache was updated successfully.
  @retval EFI_NOT_FOUND                       Saved variable context information could not be found to update the cache.
  @retval Others                              An error occurred updating the variable cache.

**/
EFI_STATUS
EFIAPI
UpdateVariableCache (
  IN  AUTHENTICATED_VARIABLE_NVM_HEADER  *VariableHeader,
  IN  VOID                               *Data,
  IN  UINTN                              DataSize
  )
{
  if (VariableHeader == NULL) {
    DEBUG ((EFI_D_ERROR, "variable header is null in VerifyCommandResponseData()\n"));
    return EFI_NOT_FOUND;
  }
  DEBUG ((EFI_D_INFO, "  Variable Contents Going to UpdateNvCache():\n"));
  PrintBinaryBuffer (Data, DataSize);
  return mVariableStorageSupportProtocol->UpdateNvCache (
                                            GetVariableNamePtr (
                                              (VARIABLE_NVM_HEADER *) VariableHeader,
                                              TRUE
                                              ),
                                            &VariableHeader->VendorGuid,
                                            Data,
                                            DataSize,
                                            VariableHeader->Attributes,
                                            VariableHeader->PubKeyIndex,
                                            VariableHeader->MonotonicCount,
                                            &VariableHeader->TimeStamp
                                            );
}

/**
  Updates the variable index area in memory with context in the backup area.

  @retval EFI_SUCCESS                         The index area was updated successfully.
  @retval EFI_NOT_FOUND                       Saved index context information could not be found.

**/
EFI_STATUS
EFIAPI
UpdateMemoryIndexArea (
  VOID
  )
{
  if (mNvmHeaderGlobals.BackupVariableStoreHeaderBase == 0) {
    return EFI_NOT_FOUND;
  }

  //
  // Ignore updating Entries in Variable store header for delete / update variable
  //
  if (
    ((VARIABLE_NVM_STORE_HEADER *) mNvmHeaderGlobals.BackupVariableStoreHeaderBase)->VariableHeaderTotalLength <
    mNvmHeaderGlobals.BackupVariableStoreHeader.VariableHeaderTotalLength
    ) {
    //
    // Update Variable store header values on local memory using Base and offset values from Global variables.
    //
    ((VARIABLE_NVM_STORE_HEADER *) mNvmHeaderGlobals.BackupVariableStoreHeaderBase)->VariableHeaderTotalEntries++;
  }

  ((VARIABLE_NVM_STORE_HEADER *) mNvmHeaderGlobals.BackupVariableStoreHeaderBase)->VariableHeaderTotalLength =
    mNvmHeaderGlobals.BackupVariableStoreHeader.VariableHeaderTotalLength;
  ((VARIABLE_NVM_STORE_HEADER *) mNvmHeaderGlobals.BackupVariableStoreHeaderBase)->VariableDataTotalLength =
    mNvmHeaderGlobals.BackupVariableStoreHeader.VariableDataTotalLength;
  ((VARIABLE_NVM_STORE_HEADER *) mNvmHeaderGlobals.BackupVariableStoreHeaderBase)->Size =
    mNvmHeaderGlobals.BackupVariableStoreHeader.Size;

  //
  // Update the Variable header value with global structure.
  //
  CopyMem (
    (UINT8 *) mNvmHeaderGlobals.BackupVariableStoreHeaderBase + mNvmHeaderGlobals.BackupVariableHeaderWriteOffset,
    &mNvmHeaderGlobals.AuthenticatedVariableHeader,
    mNvmHeaderGlobals.VariableHeaderTotalSize
    );
   
   //
   //Update the NV cache 
   //
   if (mNvmHeaderGlobals.AuthenticatedVariableHeader.State == VAR_ADDED) {
      UpdateVariableCache (
        &mNvmHeaderGlobals.AuthenticatedVariableHeader,
        mCseVaraiableDataBuffer,
        mNvmHeaderGlobals.AuthenticatedVariableHeader.DataSize
        );
      ZeroMem (mCseVaraiableDataBuffer, CMD_QUEUE_SINGLE_CMD_BUFFER_SIZE);
   }
  ZeroMem (&mNvmHeaderGlobals, sizeof (NVM_HEADER_GLOBALS));

  return EFI_SUCCESS;
}

/**
  This function should be called after each trusted channel read message completes. After the last message
  is received the read message data will be sent to the variable cache.

  @param[in]        IsLastCommand               Indicates whether this is the last command for a variable.
  @param[in]        TrustedChannelReadResponse  A pointer to the trusted channel read response for the message.

  @retval           EFI_SUCCESS                 The read message data was successfully processsed.
  @retval           EFI_INVALID_PARAMETER       The TrustedChannelHeader parameter is NULL.
  @retval           Others                      An error occurred attempting to update the variable cache.

**/
EFI_STATUS
EFIAPI
VariableTrustedChannelReadMessageComplete (
  IN  BOOLEAN                                     IsLastCommand,
  IN  CONST HECI2_TRUSTED_CHANNEL_BIOS_READ_RESP  *TrustedChannelReadResponse
  )
{
  EFI_STATUS    Status;
  STATIC UINTN  VerificationBufferAccumulatedDataSize = 0;

  if (TrustedChannelReadResponse == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  //
  // Continue appending command data in the module verification buffer
  //
  CopyMem (
    (VOID *) &mCseCommandQueueVerificationBuffer[VerificationBufferAccumulatedDataSize],
    (VOID *) (TrustedChannelReadResponse + 1),
    TrustedChannelReadResponse->DataSize
    );
  VerificationBufferAccumulatedDataSize += TrustedChannelReadResponse->DataSize;
  DEBUG ((EFI_D_INFO, "  Current data chunk (data size = %d bytes):\n", TrustedChannelReadResponse->DataSize));
  PrintBinaryBuffer ((UINT8 *)  (TrustedChannelReadResponse + 1), TrustedChannelReadResponse->DataSize);
  DEBUG ((EFI_D_INFO, "  VerificationBufferAccumulatedDataSize: %d\n", VerificationBufferAccumulatedDataSize));

  if (IsLastCommand) {
    Status = UpdateVariableCache (mCurrentGetVariableHeader,mCseCommandQueueVerificationBuffer,VerificationBufferAccumulatedDataSize);
    VerificationBufferAccumulatedDataSize = 0;
    ZeroMem (&mCseCommandQueueVerificationBuffer[0], mCseCommandQueueVerificationBufferSize);
  }

  return Status;
}

/**
  This function should be called after each trusted channel write message completes. After the last message
  is received the variable index area will be updated to reflect the details of the variable written.

  @param[in]        IsLastCommand             Indicates whether this is the last command for a variable.

  @retval           EFI_SUCCESS               The write message was successfully processed.
  @retval           Others                    An error occurred updating the index area.

**/
EFI_STATUS
EFIAPI
VariableTrustedChannelWriteMessageComplete (
  IN  BOOLEAN                             IsLastCommand
  )
{
  if (IsLastCommand) {
    //
    // When on last write command, update the local memory with new variable store header and variable header values.
    //
    DEBUG ((
      EFI_D_INFO,
      "Verfication success for new variable header at memory address 0x%x\nUpdating local memory index area...\n",
      mNvmHeaderGlobals.BackupVariableStoreHeaderBase + mNvmHeaderGlobals.BackupVariableHeaderWriteOffset
      ));

    return UpdateMemoryIndexArea ();
  }

  return EFI_SUCCESS;
}

/**
  CSE Variable Storage Response Command SMM verification.

  This function validates the CSE HMAC response data, concatenates chunked response data locally, and
  passes the accumulated data to the core variable cache on the final command in the GetVariable() sequence.

  @param[in]     VerifyResponse  A pointer to a Verify Response structure passed from a Runtime DXE wrapper.

  @retval EFI_SUCCESS                         The response data was verified and handles successfully.
  @retval EFI_INVALID_PARAMETER               The VerifyResponse actual parameter was NULL.
  @retval EFI_PROTOCOL_ERROR                  The HECI_TRUSTED_CHANNEL_PROTOCOL SMM protocol could not be located.
  @retval EFI_NOT_READY                       The EDKII_VARIABLE_STORAGE_SUPPORT_PROTOCOL could not be located.
  @retval EFI_OUT_OF_RESOURCES                An internal buffer could not be allocated using SMM memory services.
  @retval EFI_SECURITY_VIOLATION              The response message signature was computed and is invalid.
  @retval Others                              An error occurred attempting to update the non-volatile cache.
**/
EFI_STATUS
EFIAPI
VerifyCommandResponseData (
  IN   CONST SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_VERIFY_RESPONSE   *VerifyResponse
  )
{
  EFI_STATUS                                  Status;
  UINTN                                       CurrentVerificationCommandDataSize = 0;
  UINTN                                       CurrentVerificationCommandHeaderSize = 0;
  CONST HECI2_TRUSTED_CHANNEL_BIOS_HEADER     *VerificationTrustedHeader;

  BOOLEAN                                     IsVerificationSuccessful              = FALSE;


  VerificationTrustedHeader = (CONST HECI2_TRUSTED_CHANNEL_BIOS_HEADER *) VerifyResponse->Data;

  switch (VerificationTrustedHeader->CommandId) {
    case HECI2_READ_DATA_CMD_ID:
      // Read Command
      CurrentVerificationCommandDataSize = ((CONST HECI2_TRUSTED_CHANNEL_BIOS_READ_RESP *) VerificationTrustedHeader)->DataSize;
      CurrentVerificationCommandHeaderSize = sizeof (HECI2_TRUSTED_CHANNEL_BIOS_READ_RESP);

      //
      // Read commands must be able to update the variable NV cache
      //
      if (mVariableStorageSupportProtocol == NULL) {
        return EFI_NOT_READY;
      }
      break;
    case HECI2_WRITE_DATA_CMD_ID:
      // Write Command
      CurrentVerificationCommandDataSize = 0;
      CurrentVerificationCommandHeaderSize = sizeof (HECI2_TRUSTED_CHANNEL_BIOS_WRITE_RESP);

      break;
    default:
      // Invalid HECI2 command received
      ASSERT (FALSE);
      break;
  }

  DEBUG ((EFI_D_INFO, "  CurrentVerificationCommandHeaderSize: %d\n", CurrentVerificationCommandHeaderSize));
  DEBUG ((EFI_D_INFO, "  CurrentVerificationCommandDataSize: %d\n", CurrentVerificationCommandDataSize));

  if (mHeciTrustedChannelVerificationProtocol == NULL) {
    return EFI_PROTOCOL_ERROR;
  }
  Status = mHeciTrustedChannelVerificationProtocol->VerifyTrustedHeader (
                                                      (CONST UINT8 *) VerificationTrustedHeader,
                                                      (UINT32) (
                                                        CurrentVerificationCommandHeaderSize +
                                                        CurrentVerificationCommandDataSize
                                                        ),
                                                      &IsVerificationSuccessful
                                                      );
  if (!IsVerificationSuccessful) {
    DEBUG ((EFI_D_ERROR, "HMAC SHA256 signature verification failed in command queue!\n"));
    return EFI_SECURITY_VIOLATION;
  }

  if (VerificationTrustedHeader->CommandId == HECI2_READ_DATA_CMD_ID) {
    Status =  VariableTrustedChannelReadMessageComplete (
                VerifyResponse->IsLastCommand,
                (CONST HECI2_TRUSTED_CHANNEL_BIOS_READ_RESP *) VerificationTrustedHeader
                );
  } else if (VerificationTrustedHeader->CommandId == HECI2_WRITE_DATA_CMD_ID) {
    Status = VariableTrustedChannelWriteMessageComplete (VerifyResponse->IsLastCommand);
  }
  ASSERT_EFI_ERROR (Status);

  return Status;
}

/**
  CSE Variable Storage SMI handler entry.

  This SMI handler provides services for the Runtime DXE CSE Variable Storage Command Queue module

  Caution: This function may receive untrusted input.
  The communicate buffer is external input, so this function will do basic validation.
  Each sub function GetPayloadSize(), GetCommandQueue(),
  VerifyResponse(), GetProxyState(), SetProxyState(),
  ResetCommandQueue() should also do validation based on its own knowledge.

  @param[in]     DispatchHandle  The unique handle assigned to this handler by SmiHandlerRegister().
  @param[in]     RegisterContext Points to an optional handler context which was specified when the
                                 handler was registered.
  @param[in, out] CommBuffer     A pointer to a collection of data in memory that will
                                 be conveyed from a non-SMM environment into an SMM environment.
  @param[in, out] CommBufferSize The size of the CommBuffer.

  @retval EFI_SUCCESS                         The interrupt was handled and quiesced. No other handlers
                                              should still be called.
  @retval EFI_WARN_INTERRUPT_SOURCE_QUIESCED  The interrupt has been quiesced but other handlers should
                                              still be called.
  @retval EFI_WARN_INTERRUPT_SOURCE_PENDING   The interrupt is still pending and other handlers should still
                                              be called.
  @retval EFI_INTERRUPT_PENDING               The interrupt could not be quiesced.
**/
EFI_STATUS
EFIAPI
SmmCseVariableStorageHandler (
  IN       EFI_HANDLE    DispatchHandle,
  IN       CONST VOID    *RegisterContext,
  IN OUT   VOID          *CommBuffer,
  IN OUT   UINTN         *CommBufferSize
  )
{
  EFI_STATUS                                             Status;
  UINTN                                                  CommBufferPayloadSize;
  UINTN                                                  TempCommBufferSize;

  SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_HEADER                     *SmmVariableStorageFunctionHeader;
  SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_GET_COMMAND_QUEUE          *GetCommandQueue;
  SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_GET_PAYLOAD_SIZE           *GetPayloadSize;
  SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_GET_PROXYSTATE_MSG_HEADER  *GetProxyStateMessage;
  SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_SET_PROXYSTATE_MSG_HEADER  *SetProxyStateMessage;
  SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_DEBUG_MSG                  *DebugMessage;

  //
  // If input is invalid, stop processing this SMI
  //
  if (CommBuffer == NULL || CommBufferSize == NULL) {
    return EFI_SUCCESS;
  }

  TempCommBufferSize = *CommBufferSize;

  if (TempCommBufferSize < SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_HEADER_SIZE) {
    DEBUG ((EFI_D_ERROR, "CseVariableStorageSmmHandler: Communication buffer size is invalid!\n"));
    return EFI_SUCCESS;
  }
  CommBufferPayloadSize = TempCommBufferSize - SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_HEADER_SIZE;
  if (CommBufferPayloadSize > mCseCommandQueue.Capacity) {
    DEBUG ((EFI_D_ERROR, "CseVariableStorageSmmHandler: SMM communication buffer payload size invalid!\n"));
    return EFI_SUCCESS;
  }

  if (!SmmIsBufferOutsideSmmValid ((UINTN) CommBuffer, TempCommBufferSize)) {
    DEBUG ((EFI_D_ERROR, "CseVariableStorageSmmHandler: SMM communication buffer in SMRAM or overflow!\n"));
    return EFI_SUCCESS;
  }

  SmmVariableStorageFunctionHeader = (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_HEADER *) CommBuffer;

  switch (SmmVariableStorageFunctionHeader->Function) {
    case SMM_CSE_VARIABLE_STORAGE_FUNCTION_GET_PAYLOAD_SIZE:
      if (CommBufferPayloadSize < sizeof (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_GET_PAYLOAD_SIZE)) {
        DEBUG ((EFI_D_ERROR, "SMM communication buffer size for GetPayloadSize is invalid!\n"));
        Status = EFI_INVALID_PARAMETER;
      } else {
        GetPayloadSize = (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_GET_PAYLOAD_SIZE *) SmmVariableStorageFunctionHeader->Data;
        GetPayloadSize->PayloadSize = mCseCommandQueue.Capacity;
        Status = EFI_SUCCESS;
      }
      break;

    case SMM_CSE_VARIABLE_STORAGE_FUNCTION_GET_COMMAND_QUEUE:
      if (CommBufferPayloadSize < sizeof (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_GET_COMMAND_QUEUE)) {
        DEBUG ((EFI_D_ERROR, "SMM communication buffer for GetCommandQueue is invalid!\n"));
        Status = EFI_INVALID_PARAMETER;
      } else if (CommBufferPayloadSize < mCseCommandQueue.Capacity) {
        // The communicate buffer should have been allocated large enough to hold the entire SMM command queue
        DEBUG ((EFI_D_ERROR, "SMM communication buffer size for GetCommandQueue is too small!\n"));
        Status = EFI_BUFFER_TOO_SMALL;
      } else {
        DEBUG ((EFI_D_INFO, "Getting command queue in SMM.\n"));
        GetCommandQueue = (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_GET_COMMAND_QUEUE *) SmmVariableStorageFunctionHeader->Data;
        CopyMem (
          &GetCommandQueue->CommandQueueDataBuffer[0],
          &mCseCommandQueueBuffer[0],
          mCseCommandQueue.Capacity
          );
        GetCommandQueue->CommandQueueCount = mCseCommandQueue.Count;
        Status = EFI_SUCCESS;
      }
      break;

    case SMM_CSE_VARIABLE_STORAGE_FUNCTION_VERIFY_RESPONSE:
      if (CommBufferPayloadSize < sizeof (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_VERIFY_RESPONSE)) {
        DEBUG ((EFI_D_ERROR, "SMM communication buffer size for VerifyResponse is invalid!\n"));
        Status = EFI_INVALID_PARAMETER;
      } else {
        Status = VerifyCommandResponseData (
                   (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_VERIFY_RESPONSE *) SmmVariableStorageFunctionHeader->Data
                   );
      }
      break;

    case SMM_CSE_VARIABLE_STORAGE_FUNCTION_GET_PROXY_STATE:
      if (CommBufferPayloadSize < sizeof (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_GET_PROXYSTATE_MSG_HEADER)) {
        DEBUG ((EFI_D_ERROR, "SMM communication buffer for GetProxyStateMsgHeader is invalid!\n"));
        Status = EFI_INVALID_PARAMETER;
      } else {
        if (mGetProxyStateMessageHeader.MonotonicCounter == 0) {
          DEBUG ((EFI_D_ERROR, "GetProxyState message was not created properly.\n"));
          Status = EFI_NOT_READY;
        } else {
          GetProxyStateMessage = (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_GET_PROXYSTATE_MSG_HEADER *) SmmVariableStorageFunctionHeader->Data;
          CopyMem (
            &GetProxyStateMessage->MessageHeader,
            &mGetProxyStateMessageHeader,
            sizeof (HECI2_TRUSTED_CHANNEL_BIOS_HEADER)
            );
          DEBUG ((EFI_D_INFO, "GetProxyState message returned to RT DXE from SMM global:\n"));
          PrintBinaryBuffer ((UINT8 *) &mGetProxyStateMessageHeader, sizeof (HECI2_TRUSTED_CHANNEL_BIOS_HEADER));
          GetProxyStateMessage->ProxyState = mHeciProxyState;
          Status = EFI_SUCCESS;
        }
      }
      break;

    case SMM_CSE_VARIABLE_STORAGE_FUNCTION_SET_PROXY_STATE:
      if (CommBufferPayloadSize < sizeof (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_SET_PROXYSTATE_MSG_HEADER)) {
        DEBUG ((DEBUG_ERROR, "SMM communication buffer for SetProxyState is invalid!\n"));
        Status = EFI_INVALID_PARAMETER;
      } else {
        SetProxyStateMessage = (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_SET_PROXYSTATE_MSG_HEADER *) SmmVariableStorageFunctionHeader->Data;
        mHeciProxyState = SetProxyStateMessage->ProxyState;
        DEBUG ((DEBUG_INFO, "RT DXE set the ProxyState to: %d\n", mHeciProxyState));
        Status = EFI_SUCCESS;
      }
      break;

    case SMM_CSE_VARIABLE_STORAGE_FUNCTION_RESET_COMMAND_QUEUE:
      Status = CircularBufferReset (&mCseCommandQueue);
      break;

    case SMM_CSE_VARIABLE_STORAGE_FUNCTION_DEBUG_MESSAGE:
      DebugMessage = (SMM_CSE_VARIABLE_STORAGE_COMMUNICATE_DEBUG_MSG *) SmmVariableStorageFunctionHeader->Data;
      DEBUG ((EFI_D_INFO, "RT DXE Message:\n"));
      DEBUG ((EFI_D_INFO, "  %a", DebugMessage->DebugMsg));

      if (DebugMessage->Datalength != 0) {
        DEBUG ((EFI_D_INFO, "\nRT DXE Message Data:\n"));
        PrintBinaryBuffer (&DebugMessage->Data[0], DebugMessage->Datalength);
      }
      DEBUG ((EFI_D_INFO, "\n"));
      Status = EFI_SUCCESS;
      break;

    default:
      Status = EFI_UNSUPPORTED;
  }

  SmmVariableStorageFunctionHeader->ReturnStatus = Status;

  return EFI_SUCCESS;
}

/**
  Protocol notification event handler.

  @param   Protocol      Points to the protocol's unique identifier
  @param   Interface     Points to the interface instance
  @param   Handle        The handle on which the interface was installed

  @retval  EFI_SUCCESS   The callback function ran successfully
  @retval  EFI_NOT_FOUND The EFI_SMM_SX_DISPATCH2_PROTOCOL could not be found
**/
EFI_STATUS
EFIAPI
SmmSxDispatch2Callback (
  IN CONST EFI_GUID      *Protocol,
  IN VOID                *Interface,
  IN EFI_HANDLE          Handle
  )
{
  EFI_STATUS                    Status;
  EFI_SMM_SX_REGISTER_CONTEXT   EntryDispatchContext;
  EFI_SMM_SX_DISPATCH2_PROTOCOL *SxDispatch = NULL;

  Status = gSmst->SmmLocateProtocol (
                    &gEfiSmmSxDispatch2ProtocolGuid,
                    NULL,
                    (VOID **) &SxDispatch
                    );
  ASSERT_EFI_ERROR (Status);
  if (SxDispatch == NULL) {
    return EFI_NOT_FOUND;
  }

  EntryDispatchContext.Type  = SxS3;
  EntryDispatchContext.Phase = SxEntry;

  Status = SxDispatch->Register (
                         SxDispatch,
                         CseVariableStorageS3EntryProxyUpdateCallBack,
                         &EntryDispatchContext,
                         &mCseSmmVariableHandle
                         );
  ASSERT_EFI_ERROR (Status);
  return Status;
}

/**
  Initializes module buffers used in this module.

  @retval        EFI_SUCCESS                 The buffers were initialized successfully.
  @retval        Others                      The buffers could not be initialized.

**/
EFI_STATUS
EFIAPI
InitializeModuleBuffers (
  VOID
  )
{
  mCseCommandQueueBuffer = AllocateZeroPool (CMD_QUEUE_TOTAL_BUFFER_SIZE);
  if (mCseCommandQueueBuffer == NULL) {
    ASSERT_EFI_ERROR (mCseCommandQueueBuffer != NULL);
    return EFI_OUT_OF_RESOURCES;
  }

  mCseCommandQueueVerificationBuffer = AllocateZeroPool (CMD_QUEUE_SINGLE_CMD_BUFFER_SIZE);
  if (mCseCommandQueueVerificationBuffer == NULL) {
    ASSERT_EFI_ERROR (mCseCommandQueueVerificationBuffer != NULL);
    return EFI_OUT_OF_RESOURCES;
  }

  
  mCseVaraiableDataBuffer = AllocateZeroPool (CMD_QUEUE_SINGLE_CMD_BUFFER_SIZE);
  if (mCseVaraiableDataBuffer == NULL) {
    ASSERT_EFI_ERROR (mCseVaraiableDataBuffer != NULL);
    return EFI_OUT_OF_RESOURCES;
  }

  return EFI_SUCCESS;
}

/**
  EDKII Variable Storage Support Protocol notification callback.

  @param   Protocol      Points to the protocol's unique identifier
  @param   Interface     Points to the interface instance
  @param   Handle        The handle on which the interface was installed

  @retval  EFI_SUCCESS   The callback function ran successfully
  @retval  EFI_NOT_FOUND The EDKII_VARIABLE_STORAGE_SUPPORT_PROTOCOL could not be found
**/
EFI_STATUS
EFIAPI
VariableStorageSupportProtocolCallback (
  IN CONST EFI_GUID      *Protocol,
  IN VOID                *Interface,
  IN EFI_HANDLE          Handle
  )
{
  EFI_STATUS         Status;

  Status = gSmst->SmmLocateProtocol (
                    &gEdkiiVariableStorageSupportProtocolGuid,
                    NULL,
                    (VOID **) &mVariableStorageSupportProtocol
                    );
  ASSERT_EFI_ERROR (Status);
  if (mVariableStorageSupportProtocol == NULL) {
    return EFI_NOT_FOUND;
  }
  mVariableStorageSupportProtocol->NotifyWriteServiceReady ();

  return Status;
}

/**
  Provide SMM functionality of CSE NVM variable storage services.

  @param  ImageHandle  The image handle.
  @param  SystemTable  The system table.

  @retval EFI_SUCCESS  The protocol was installed successfully.
  @retval Others       Protocol could not be installed.
**/
EFI_STATUS
EFIAPI
CseVariableStorageSmmInitialize (
  IN EFI_HANDLE         ImageHandle,
  IN EFI_SYSTEM_TABLE   *SystemTable
  )
{
  EFI_STATUS    Status;
  VOID          *EdkiiVariableStorageSupportProtocolRegistration;
  VOID          *SmmSxDispatch2ProtocolRegistration;
  EFI_HANDLE    Handle                 = NULL;
  EFI_HANDLE    SmiHandle              = NULL;

  if (!PcdGetBool (PcdEnableCseVariableStorage) || PcdGetBool (PcdNvVariableEmulationMode)) {
    DEBUG ((EFI_D_INFO, "CSE Variable Storage Protocol is disabled.\n"));
    return EFI_SUCCESS;
  }

  Status = CseVariableStorageCommonInitialize ();
  if (EFI_ERROR (Status)) {
    ASSERT_EFI_ERROR (Status);
    return Status;
  }

  Status = InitializeModuleBuffers ();
  if (EFI_ERROR (Status)) {
    return Status;
  }
  mInSmm = TRUE;
  mCseCommandQueue.Buffer = mCseCommandQueueBuffer;

  Status = gSmst->SmiHandlerRegister (
                    SmmCseVariableStorageHandler,
                    &gCseVariableStorageProtocolInstanceGuid,
                    &SmiHandle
                    );
  ASSERT_EFI_ERROR (Status);

  Status = gSmst->SmmInstallProtocolInterface (
                    &Handle,
                    &gEdkiiVariableStorageProtocolGuid,
                    EFI_NATIVE_INTERFACE,
                    &mCseVariableStorageProtocol
                    );
  ASSERT_EFI_ERROR (Status);

  Status = gSmst->SmmInstallProtocolInterface (
                    &Handle,
                    &gCseVariableStorageSupportProtocolGuid,
                    EFI_NATIVE_INTERFACE,
                    &mCseVariableStorageSupportProtocol
                    );
  ASSERT_EFI_ERROR (Status);

  Status = gSmst->SmmLocateProtocol (
                    &gEfiHeciTrustedChannelSmmProtocolGuid,
                    NULL,
                    (VOID **) &mHeciTrustedChannelVerificationProtocol
                    );
  ASSERT_EFI_ERROR (Status);

  ZeroMem (&mGetProxyStateMessageHeader, sizeof (mGetProxyStateMessageHeader));

  Status = gSmst->SmmRegisterProtocolNotify (
                    &gEdkiiVariableStorageSupportProtocolGuid,
                    VariableStorageSupportProtocolCallback,
                    &EdkiiVariableStorageSupportProtocolRegistration
                    );
  ASSERT_EFI_ERROR (Status);

  Status = gSmst->SmmRegisterProtocolNotify (
                    &gEfiSmmSxDispatch2ProtocolGuid,
                    SmmSxDispatch2Callback,
                    &SmmSxDispatch2ProtocolRegistration
                    );
  ASSERT_EFI_ERROR (Status);

  //
  // Notify the Runtime DXE wrapper driver CSE SMM variable services are ready
  //
  DEBUG ((EFI_D_INFO, "Notifying modules CSE Variable Storage SMM services are available.\n"));
  Status = SystemTable->BootServices->InstallProtocolInterface (
                                        &mCseSmmVariableHandle,
                                        &gEfiCseVariableSmmReadyProtocolGuid,
                                        EFI_NATIVE_INTERFACE,
                                        NULL
                                        );
  ASSERT_EFI_ERROR (Status);

  return Status;
}
