/** @file
  Unit tests for the implementation of SecureBootVariableLib.

  Copyright (c) 2022, Intel Corporation. All rights reserved.
  SPDX-License-Identifier: BSD-2-Clause-Patent
**/
#include <Library/GoogleTestLib.h>
#include <GoogleTest/Library/MockUefiLib.h>
#include <GoogleTest/Library/MockUefiRuntimeServicesTableLib.h>

extern "C" {
  #include <Uefi.h>
  #include <UefiSecureBoot.h>
  #include <Guid/AuthenticatedVariableFormat.h>
  #include <Guid/ImageAuthentication.h>
  #include <Library/SecureBootVariableLib.h>
  #include <Library/MemoryAllocationLib.h>
}

using namespace testing;

//////////////////////////////////////////////////////////////////////////////
class SetSecureBootModeTest : public Test {
protected:
  MockUefiRuntimeServicesTableLib RtServicesMock;
  UINT8 SecureBootMode;
  EFI_STATUS Status;

  void
  SetUp (
    ) override
  {
    // Any random magic number can be used for these tests
    SecureBootMode = 0xAB;
  }
};

// Test SetSecureBootMode() API from SecureBootVariableLib to verify the
// expected error is returned when the call to gRT->SetVariable() fails.
TEST_F (SetSecureBootModeTest, SetVarError) {
  EXPECT_CALL (RtServicesMock, gRT_SetVariable)
    .WillOnce (Return (EFI_INVALID_PARAMETER));

  Status = SetSecureBootMode (SecureBootMode);
  EXPECT_EQ (Status, EFI_INVALID_PARAMETER);
}

// Test SetSecureBootMode() API from SecureBootVariableLib to verify the
// expected secure boot mode is written to the correct variable in the call
// to gRT->SetVariable().
TEST_F (SetSecureBootModeTest, PropogateModeToSetVar) {
  EXPECT_CALL (
    RtServicesMock,
    gRT_SetVariable (
      Char16StrEq (EFI_CUSTOM_MODE_NAME),
      BufferEq (&gEfiCustomModeEnableGuid, sizeof (EFI_GUID)),
      EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS,
      sizeof (SecureBootMode),
      BufferEq (&SecureBootMode, sizeof (SecureBootMode))
      )
    )
    .WillOnce (Return (EFI_SUCCESS));

  Status = SetSecureBootMode (SecureBootMode);
  EXPECT_EQ (Status, EFI_SUCCESS);
}

//////////////////////////////////////////////////////////////////////////////
class GetSetupModeTest : public Test {
protected:
  MockUefiRuntimeServicesTableLib RtServicesMock;
  UINT8 SetupMode;
  EFI_STATUS Status;
  UINT8 ExpSetupMode;

  void
  SetUp (
    ) override
  {
    // Any random magic number can be used for these tests
    ExpSetupMode = 0xAB;
  }
};

// Test GetSetupMode() API from SecureBootVariableLib to verify the expected
// error is returned when the call to gRT->GetVariable() fails.
TEST_F (GetSetupModeTest, GetVarError) {
  EXPECT_CALL (RtServicesMock, gRT_GetVariable)
    .WillOnce (Return (EFI_INVALID_PARAMETER));

  Status = GetSetupMode (&SetupMode);
  EXPECT_EQ (Status, EFI_INVALID_PARAMETER);
}

// Test GetSetupMode() API from SecureBootVariableLib to verify the expected
// setup mode is returned (and with a success return code) when the mode is
// successfully read from the call to gRT->GetVariable().
TEST_F (GetSetupModeTest, FetchModeFromGetVar) {
  EXPECT_CALL (
    RtServicesMock,
    gRT_GetVariable (
      Char16StrEq (EFI_SETUP_MODE_NAME),
      BufferEq (&gEfiGlobalVariableGuid, sizeof (EFI_GUID)),
      _,
      Pointee (Eq (sizeof (SetupMode))),
      NotNull ()
      )
    )
    .WillOnce (
       DoAll (
         SetArgPointee<3>(sizeof (ExpSetupMode)),
         SetArgBuffer<4>(&ExpSetupMode, sizeof (ExpSetupMode)),
         Return (EFI_SUCCESS)
         )
       );

  Status = GetSetupMode (&SetupMode);
  ASSERT_EQ (Status, EFI_SUCCESS);
  EXPECT_EQ (SetupMode, ExpSetupMode);
}

//////////////////////////////////////////////////////////////////////////////
class IsSecureBootEnabledTest : public Test {
protected:
  MockUefiLib UefiLibMock;
  BOOLEAN Enabled;
};

// Test IsSecureBootEnabled() API from SecureBootVariableLib to verify FALSE
// is returned when the call to GetEfiGlobalVariable2() fails.
TEST_F (IsSecureBootEnabledTest, GetVarError) {
  EXPECT_CALL (UefiLibMock, GetEfiGlobalVariable2)
    .WillOnce (Return (EFI_ABORTED));

  Enabled = IsSecureBootEnabled ();
  EXPECT_EQ (Enabled, FALSE);
}

//////////////////////////////////////////////////////////////////////////////
class IsSecureBootEnabledAllocTest : public IsSecureBootEnabledTest {
protected:
  UINT8 *BootEnabledBuffer;

  void
  SetUp (
    ) override
  {
    BootEnabledBuffer = (UINT8 *)AllocatePool (1);
    ASSERT_NE (BootEnabledBuffer, nullptr);
  }
};

// Test IsSecureBootEnabled() API from SecureBootVariableLib to verify TRUE
// is returned when the call to GetEfiGlobalVariable2() is successful and
// returns SECURE_BOOT_MODE_ENABLE.
TEST_F (IsSecureBootEnabledAllocTest, IsEnabled) {
  *BootEnabledBuffer = SECURE_BOOT_MODE_ENABLE;
  EXPECT_CALL (
    UefiLibMock,
    GetEfiGlobalVariable2 (
      Char16StrEq (EFI_SECURE_BOOT_MODE_NAME),
      NotNull (),
      _
      )
    )
    .WillOnce (
       DoAll (
         SetArgBuffer<1>(&BootEnabledBuffer, sizeof (VOID *)),
         Return (EFI_SUCCESS)
         )
       );

  Enabled = IsSecureBootEnabled ();
  EXPECT_EQ (Enabled, TRUE);
}

// Test IsSecureBootEnabled() API from SecureBootVariableLib to verify FALSE
// is returned when the call to GetEfiGlobalVariable2() is successful and
// returns SECURE_BOOT_MODE_DISABLE.
TEST_F (IsSecureBootEnabledAllocTest, IsDisabled) {
  *BootEnabledBuffer = SECURE_BOOT_MODE_DISABLE;
  EXPECT_CALL (
    UefiLibMock,
    GetEfiGlobalVariable2 (
      Char16StrEq (EFI_SECURE_BOOT_MODE_NAME),
      NotNull (),
      _
      )
    )
    .WillOnce (
       DoAll (
         SetArgBuffer<1>(&BootEnabledBuffer, sizeof (VOID *)),
         Return (EFI_SUCCESS)
         )
       );

  Enabled = IsSecureBootEnabled ();
  EXPECT_EQ (Enabled, FALSE);
}

int
main (
  int   argc,
  char  *argv[]
  )
{
  testing::InitGoogleTest (&argc, argv);
  return RUN_ALL_TESTS ();
}