/** @file
    Implement the test SMM driver for status code to RAM.

***************************************************************************
* Copyright (c) 2021, Insyde Software Corp. All Rights Reserved.
*
* You may not reproduce, distribute, publish, display, perform, modify, adapt,
* transmit, broadcast, present, recite, release, license or otherwise exploit
* any part of this publication in any form, by any means, without the prior
* written permission of Insyde Software Corporation.
*
******************************************************************************
*/

#include <Uefi.h>
#include <Pi/PiStatusCode.h>
#include <Guid/MemoryStatusCodeRecord.h>
#include <Library/BaseLib.h>
#include <Library/UefiDriverEntryPoint.h>
#include <Library/SmmServicesTableLib.h>
#include <Library/DebugLib.h>
#include <Library/UefiLib.h>
#include <Library/BaseMemoryLib.h>
#include <Protocol/SmmSwDispatch2.h>

//
// Global variables
//
GLOBAL_REMOVE_IF_UNREFERENCED EFI_SMM_SYSTEM_TABLE2                      *mSmst;

RUNTIME_MEMORY_STATUSCODE_HEADER                                         *mMmMemoryStatusCodeTable;

/**
  Software SMI callback to dump the status code.

  @param[in]      DispatchHandle  The unique handle assigned to this handler by SmiHandlerRegister().
  @param[in]      DispatchContext 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 successfully.

**/
EFI_STATUS
EFIAPI
StatusCodeHandlerSmmCallback (
  IN  EFI_HANDLE                    DispatchHandle,
  IN  CONST VOID                    *DispatchContext,
  IN  OUT VOID                      *CommBuffer  OPTIONAL,
  IN  OUT UINTN                     *CommBufferSize  OPTIONAL
  )
{
  MEMORY_STATUSCODE_RECORD              *Record;
  UINTN                                 Index;
  UINTN                                 TempRecordIndex;
  EFI_SMM_SYSTEM_TABLE2                 *SmmSystemTable;


  DEBUG ((EFI_D_INFO ,"StatusCodeHandlerTestSmm: Start.\n"));

  SmmSystemTable = mSmst;
  for (Index = 0; Index < SmmSystemTable->NumberOfTableEntries; Index++) {
    if (CompareGuid (&gMemoryStatusCodeRecordGuid, &(SmmSystemTable->SmmConfigurationTable[Index].VendorGuid))) {
      // Workaround, and remove it if EDK is ever revised
      mMmMemoryStatusCodeTable = *((VOID **) SmmSystemTable->SmmConfigurationTable[Index].VendorTable);
      // Unmark if EDK is ever revised
      //mMmMemoryStatusCodeTable = SmmSystemTable->SmmConfigurationTable[Index].VendorTable;
      break;
    }
  }

  if (Index == SmmSystemTable->NumberOfTableEntries) {
    DEBUG ((EFI_D_INFO ,"Get Smm System Configuration Table: Not Found.\n"));
    return EFI_SUCCESS;
  }

  DEBUG ((EFI_D_INFO ,"mMmMemoryStatusCodeTable address: 0x%X\n", mMmMemoryStatusCodeTable));
  DEBUG ((EFI_D_INFO ,"mMmMemoryStatusCodeTable->RecordIndex = %d\n", mMmMemoryStatusCodeTable->RecordIndex));
  DEBUG ((EFI_D_INFO ,"mMmMemoryStatusCodeTable->NumberOfRecords = %d\n", mMmMemoryStatusCodeTable->NumberOfRecords));
  DEBUG ((EFI_D_INFO ,"mMmMemoryStatusCodeTable->MaxRecordsNumber = %d\n", mMmMemoryStatusCodeTable->MaxRecordsNumber));
  TempRecordIndex = mMmMemoryStatusCodeTable->RecordIndex;
  DEBUG ((EFI_D_INFO ,"TempRecordIndex = %d\n", TempRecordIndex));

  if (mMmMemoryStatusCodeTable != NULL) {
    Record = (MEMORY_STATUSCODE_RECORD *) (mMmMemoryStatusCodeTable + 1);
    for (Index = 0; Index < TempRecordIndex; Index++) {
      DEBUG ((EFI_D_INFO ,"Smm[%d] CodeType 0x%X V%08X I%X\n", Index + 1, Record->CodeType, Record->Value, Record->Instance));
    }
  }
  DEBUG ((EFI_D_INFO ,"StatusCodeHandlerTestSmm: End.\n"));

  return EFI_SUCCESS;
}

/**

  @param[in]  ImageHandle     The image handle.
  @param[in]  SystemTable     The system table.

  @retval EFI_SUCCESS            Register SW SMM successfully.
**/
EFI_STATUS
EFIAPI
StatusCodeHandlerTestSmmEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS                            Status;
  EFI_SMM_SW_DISPATCH2_PROTOCOL         *SwDispatch;
  EFI_HANDLE                            SwHandle;
  EFI_SMM_SW_REGISTER_CONTEXT           SwContext;

  mSmst = gSmst;
  //
  // Locate the SmmSxDispatch protocol
  //
  Status = mSmst->SmmLocateProtocol (&gEfiSmmSwDispatch2ProtocolGuid, NULL, (VOID**)&SwDispatch);
  if (EFI_ERROR (Status)) {
    return Status;
  }
  //
  // Register StatusCode handler
  //
  SwContext.SwSmiInputValue = (UINTN) -1;
  Status = SwDispatch->Register (
                         SwDispatch,
                         StatusCodeHandlerSmmCallback,
                         &SwContext,
                         &SwHandle
                         );
  DEBUG ((EFI_D_INFO ,"SwContext.SwSmiInputValue: 0x%X, Status = %r\n", SwContext.SwSmiInputValue, Status));

  return Status;
}