/** @file
  Implement TPM2 EnhancedAuthorization related command.

Copyright (c) 2014 - 2018, Intel Corporation. All rights reserved. <BR>
SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <IndustryStandard/UefiTcgPlatform.h>
#include <Library/Tpm2CommandLib.h>
#include <Library/Tpm2DeviceLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/BaseLib.h>
#include <Library/DebugLib.h>

#pragma pack(1)

typedef struct {
  TPM2_COMMAND_HEADER       Header;
  TPMI_DH_ENTITY            AuthHandle;
  TPMI_SH_POLICY            PolicySession;
  UINT32                    AuthSessionSize;
  TPMS_AUTH_COMMAND         AuthSession;
  TPM2B_NONCE               NonceTPM;
  TPM2B_DIGEST              CpHashA;
  TPM2B_NONCE               PolicyRef;
  INT32                     Expiration;
} TPM2_POLICY_SECRET_COMMAND;

typedef struct {
  TPM2_RESPONSE_HEADER      Header;
  UINT32                    AuthSessionSize;
  TPM2B_TIMEOUT             Timeout;
  TPMT_TK_AUTH              PolicyTicket;
  TPMS_AUTH_RESPONSE        AuthSession;
} TPM2_POLICY_SECRET_RESPONSE;

typedef struct {
  TPM2_COMMAND_HEADER       Header;
  TPMI_SH_POLICY            PolicySession;
  TPML_DIGEST               HashList;
} TPM2_POLICY_OR_COMMAND;

typedef struct {
  TPM2_RESPONSE_HEADER      Header;
} TPM2_POLICY_OR_RESPONSE;

typedef struct {
  TPM2_COMMAND_HEADER       Header;
  TPMI_SH_POLICY            PolicySession;
  TPM_CC                    Code;
} TPM2_POLICY_COMMAND_CODE_COMMAND;

typedef struct {
  TPM2_RESPONSE_HEADER      Header;
} TPM2_POLICY_COMMAND_CODE_RESPONSE;

typedef struct {
  TPM2_COMMAND_HEADER       Header;
  TPMI_SH_POLICY            PolicySession;
} TPM2_POLICY_GET_DIGEST_COMMAND;

typedef struct {
  TPM2_RESPONSE_HEADER      Header;
  TPM2B_DIGEST              PolicyHash;
} TPM2_POLICY_GET_DIGEST_RESPONSE;

#pragma pack()

/**
  This command includes a secret-based authorization to a policy.
  The caller proves knowledge of the secret value using an authorization
  session using the authValue associated with authHandle.

  @param[in]  AuthHandle         Handle for an entity providing the authorization
  @param[in]  PolicySession      Handle for the policy session being extended.
  @param[in]  AuthSession        Auth Session context
  @param[in]  NonceTPM           The policy nonce for the session.
  @param[in]  CpHashA            Digest of the command parameters to which this authorization is limited.
  @param[in]  PolicyRef          A reference to a policy relating to the authorization.
  @param[in]  Expiration         Time when authorization will expire, measured in seconds from the time that nonceTPM was generated.
  @param[out] Timeout            Time value used to indicate to the TPM when the ticket expires.
  @param[out] PolicyTicket       A ticket that includes a value indicating when the authorization expires.

  @retval EFI_SUCCESS            Operation completed successfully.
  @retval EFI_DEVICE_ERROR       The command was unsuccessful.
**/
EFI_STATUS
EFIAPI
Tpm2PolicySecret (
  IN      TPMI_DH_ENTITY            AuthHandle,
  IN      TPMI_SH_POLICY            PolicySession,
  IN      TPMS_AUTH_COMMAND         *AuthSession, OPTIONAL
  IN      TPM2B_NONCE               *NonceTPM,
  IN      TPM2B_DIGEST              *CpHashA,
  IN      TPM2B_NONCE               *PolicyRef,
  IN      INT32                     Expiration,
  OUT     TPM2B_TIMEOUT             *Timeout,
  OUT     TPMT_TK_AUTH              *PolicyTicket
  )
{
  EFI_STATUS                        Status;
  TPM2_POLICY_SECRET_COMMAND        SendBuffer;
  TPM2_POLICY_SECRET_RESPONSE       RecvBuffer;
  UINT32                            SendBufferSize;
  UINT32                            RecvBufferSize;
  UINT8                             *Buffer;
  UINT32                            SessionInfoSize;

  //
  // Construct command
  //
  SendBuffer.Header.tag = SwapBytes16(TPM_ST_SESSIONS);
  SendBuffer.Header.commandCode = SwapBytes32(TPM_CC_PolicySecret);
  SendBuffer.AuthHandle = SwapBytes32 (AuthHandle);
  SendBuffer.PolicySession = SwapBytes32 (PolicySession);

  //
  // Add in Auth session
  //
  Buffer = (UINT8 *)&SendBuffer.AuthSession;

  // sessionInfoSize
  SessionInfoSize = CopyAuthSessionCommand (AuthSession, Buffer);
  Buffer += SessionInfoSize;
  SendBuffer.AuthSessionSize = SwapBytes32(SessionInfoSize);

  //
  // Real data
  //
  WriteUnaligned16 ((UINT16 *)Buffer, SwapBytes16(NonceTPM->size));
  Buffer += sizeof(UINT16);
  CopyMem (Buffer, NonceTPM->buffer, NonceTPM->size);
  Buffer += NonceTPM->size;

  WriteUnaligned16 ((UINT16 *)Buffer, SwapBytes16(CpHashA->size));
  Buffer += sizeof(UINT16);
  CopyMem (Buffer, CpHashA->buffer, CpHashA->size);
  Buffer += CpHashA->size;

  WriteUnaligned16 ((UINT16 *)Buffer, SwapBytes16(PolicyRef->size));
  Buffer += sizeof(UINT16);
  CopyMem (Buffer, PolicyRef->buffer, PolicyRef->size);
  Buffer += PolicyRef->size;

  WriteUnaligned32 ((UINT32 *)Buffer, SwapBytes32((UINT32)Expiration));
  Buffer += sizeof(UINT32);

  SendBufferSize = (UINT32)((UINTN)Buffer - (UINTN)&SendBuffer);
  SendBuffer.Header.paramSize = SwapBytes32 (SendBufferSize);

  //
  // send Tpm command
  //
  RecvBufferSize = sizeof (RecvBuffer);
  Status = Tpm2SubmitCommand (SendBufferSize, (UINT8 *)&SendBuffer, &RecvBufferSize, (UINT8 *)&RecvBuffer);
  if (EFI_ERROR (Status)) {
    goto Done;
  }

  if (RecvBufferSize < sizeof (TPM2_RESPONSE_HEADER)) {
    DEBUG ((EFI_D_ERROR, "Tpm2PolicySecret - RecvBufferSize Error - %x\n", RecvBufferSize));
    Status = EFI_DEVICE_ERROR;
    goto Done;
  }
  if (SwapBytes32(RecvBuffer.Header.responseCode) != TPM_RC_SUCCESS) {
    DEBUG ((EFI_D_ERROR, "Tpm2PolicySecret - responseCode - %x\n", SwapBytes32(RecvBuffer.Header.responseCode)));
    Status = EFI_DEVICE_ERROR;
    goto Done;
  }

  //
  // Return the response
  //
  Buffer = (UINT8 *)&RecvBuffer.Timeout;
  Timeout->size = SwapBytes16(ReadUnaligned16 ((UINT16 *)Buffer));
  if (Timeout->size > sizeof(UINT64)) {
    DEBUG ((DEBUG_ERROR, "Tpm2PolicySecret - Timeout->size error %x\n", Timeout->size));
    Status = EFI_DEVICE_ERROR;
    goto Done;
  }

  Buffer += sizeof(UINT16);
  CopyMem (Timeout->buffer, Buffer, Timeout->size);

  PolicyTicket->tag = SwapBytes16(ReadUnaligned16 ((UINT16 *)Buffer));
  Buffer += sizeof(UINT16);
  PolicyTicket->hierarchy = SwapBytes32(ReadUnaligned32 ((UINT32 *)Buffer));
  Buffer += sizeof(UINT32);
  PolicyTicket->digest.size = SwapBytes16(ReadUnaligned16 ((UINT16 *)Buffer));
  Buffer += sizeof(UINT16);
  if (PolicyTicket->digest.size > sizeof(TPMU_HA)) {
    DEBUG ((DEBUG_ERROR, "Tpm2PolicySecret - digest.size error %x\n", PolicyTicket->digest.size));
    Status = EFI_DEVICE_ERROR;
    goto Done;
  }

  CopyMem (PolicyTicket->digest.buffer, Buffer, PolicyTicket->digest.size);

Done:
  //
  // Clear AuthSession Content
  //
  ZeroMem (&SendBuffer, sizeof(SendBuffer));
  ZeroMem (&RecvBuffer, sizeof(RecvBuffer));
  return Status;
}

/**
  This command allows options in authorizations without requiring that the TPM evaluate all of the options.
  If a policy may be satisfied by different sets of conditions, the TPM need only evaluate one set that
  satisfies the policy. This command will indicate that one of the required sets of conditions has been
  satisfied.

  @param[in] PolicySession      Handle for the policy session being extended.
  @param[in] HashList           the list of hashes to check for a match.

  @retval EFI_SUCCESS            Operation completed successfully.
  @retval EFI_DEVICE_ERROR       The command was unsuccessful.
**/
EFI_STATUS
EFIAPI
Tpm2PolicyOR (
  IN TPMI_SH_POLICY           PolicySession,
  IN TPML_DIGEST              *HashList
  )
{
  EFI_STATUS                        Status;
  TPM2_POLICY_OR_COMMAND            SendBuffer;
  TPM2_POLICY_OR_RESPONSE           RecvBuffer;
  UINT32                            SendBufferSize;
  UINT32                            RecvBufferSize;
  UINT8                             *Buffer;
  UINTN                             Index;

  //
  // Construct command
  //
  SendBuffer.Header.tag = SwapBytes16(TPM_ST_NO_SESSIONS);
  SendBuffer.Header.commandCode = SwapBytes32(TPM_CC_PolicyOR);

  SendBuffer.PolicySession = SwapBytes32 (PolicySession);
  Buffer = (UINT8 *)&SendBuffer.HashList;
  WriteUnaligned32 ((UINT32 *)Buffer, SwapBytes32 (HashList->count));
  Buffer += sizeof(UINT32);
  for (Index = 0; Index < HashList->count; Index++) {
    WriteUnaligned16 ((UINT16 *)Buffer, SwapBytes16 (HashList->digests[Index].size));
    Buffer += sizeof(UINT16);
    CopyMem (Buffer, HashList->digests[Index].buffer, HashList->digests[Index].size);
    Buffer += HashList->digests[Index].size;
  }

  SendBufferSize = (UINT32)((UINTN)Buffer - (UINTN)&SendBuffer);
  SendBuffer.Header.paramSize = SwapBytes32 (SendBufferSize);

  //
  // send Tpm command
  //
  RecvBufferSize = sizeof (RecvBuffer);
  Status = Tpm2SubmitCommand (SendBufferSize, (UINT8 *)&SendBuffer, &RecvBufferSize, (UINT8 *)&RecvBuffer);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  if (RecvBufferSize < sizeof (TPM2_RESPONSE_HEADER)) {
    DEBUG ((EFI_D_ERROR, "Tpm2PolicyOR - RecvBufferSize Error - %x\n", RecvBufferSize));
    return EFI_DEVICE_ERROR;
  }
  if (SwapBytes32(RecvBuffer.Header.responseCode) != TPM_RC_SUCCESS) {
    DEBUG ((EFI_D_ERROR, "Tpm2PolicyOR - responseCode - %x\n", SwapBytes32(RecvBuffer.Header.responseCode)));
    return EFI_DEVICE_ERROR;
  }

  return EFI_SUCCESS;
}

/**
  This command indicates that the authorization will be limited to a specific command code.

  @param[in]  PolicySession      Handle for the policy session being extended.
  @param[in]  Code               The allowed commandCode.

  @retval EFI_SUCCESS            Operation completed successfully.
  @retval EFI_DEVICE_ERROR       The command was unsuccessful.
**/
EFI_STATUS
EFIAPI
Tpm2PolicyCommandCode (
  IN      TPMI_SH_POLICY            PolicySession,
  IN      TPM_CC                    Code
  )
{
  EFI_STATUS                        Status;
  TPM2_POLICY_COMMAND_CODE_COMMAND  SendBuffer;
  TPM2_POLICY_COMMAND_CODE_RESPONSE RecvBuffer;
  UINT32                            SendBufferSize;
  UINT32                            RecvBufferSize;

  //
  // Construct command
  //
  SendBuffer.Header.tag = SwapBytes16(TPM_ST_NO_SESSIONS);
  SendBuffer.Header.commandCode = SwapBytes32(TPM_CC_PolicyCommandCode);

  SendBuffer.PolicySession = SwapBytes32 (PolicySession);
  SendBuffer.Code = SwapBytes32 (Code);

  SendBufferSize = (UINT32) sizeof (SendBuffer);
  SendBuffer.Header.paramSize = SwapBytes32 (SendBufferSize);

  //
  // send Tpm command
  //
  RecvBufferSize = sizeof (RecvBuffer);
  Status = Tpm2SubmitCommand (SendBufferSize, (UINT8 *)&SendBuffer, &RecvBufferSize, (UINT8 *)&RecvBuffer);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  if (RecvBufferSize < sizeof (TPM2_RESPONSE_HEADER)) {
    DEBUG ((EFI_D_ERROR, "Tpm2PolicyCommandCode - RecvBufferSize Error - %x\n", RecvBufferSize));
    return EFI_DEVICE_ERROR;
  }
  if (SwapBytes32(RecvBuffer.Header.responseCode) != TPM_RC_SUCCESS) {
    DEBUG ((EFI_D_ERROR, "Tpm2PolicyCommandCode - responseCode - %x\n", SwapBytes32(RecvBuffer.Header.responseCode)));
    return EFI_DEVICE_ERROR;
  }

  return EFI_SUCCESS;
}

/**
  This command returns the current policyDigest of the session. This command allows the TPM
  to be used to perform the actions required to precompute the authPolicy for an object.

  @param[in]  PolicySession      Handle for the policy session.
  @param[out] PolicyHash         the current value of the policyHash of policySession.

  @retval EFI_SUCCESS            Operation completed successfully.
  @retval EFI_DEVICE_ERROR       The command was unsuccessful.
**/
EFI_STATUS
EFIAPI
Tpm2PolicyGetDigest (
  IN      TPMI_SH_POLICY            PolicySession,
     OUT  TPM2B_DIGEST              *PolicyHash
  )
{
  EFI_STATUS                        Status;
  TPM2_POLICY_GET_DIGEST_COMMAND    SendBuffer;
  TPM2_POLICY_GET_DIGEST_RESPONSE   RecvBuffer;
  UINT32                            SendBufferSize;
  UINT32                            RecvBufferSize;

  //
  // Construct command
  //
  SendBuffer.Header.tag = SwapBytes16(TPM_ST_NO_SESSIONS);
  SendBuffer.Header.commandCode = SwapBytes32(TPM_CC_PolicyGetDigest);

  SendBuffer.PolicySession = SwapBytes32 (PolicySession);

  SendBufferSize = (UINT32) sizeof (SendBuffer);
  SendBuffer.Header.paramSize = SwapBytes32 (SendBufferSize);

  //
  // send Tpm command
  //
  RecvBufferSize = sizeof (RecvBuffer);
  Status = Tpm2SubmitCommand (SendBufferSize, (UINT8 *)&SendBuffer, &RecvBufferSize, (UINT8 *)&RecvBuffer);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  if (RecvBufferSize < sizeof (TPM2_RESPONSE_HEADER)) {
    DEBUG ((EFI_D_ERROR, "Tpm2PolicyGetDigest - RecvBufferSize Error - %x\n", RecvBufferSize));
    return EFI_DEVICE_ERROR;
  }
  if (SwapBytes32(RecvBuffer.Header.responseCode) != TPM_RC_SUCCESS) {
    DEBUG ((EFI_D_ERROR, "Tpm2PolicyGetDigest - responseCode - %x\n", SwapBytes32(RecvBuffer.Header.responseCode)));
    return EFI_DEVICE_ERROR;
  }

  //
  // Return the response
  //
  PolicyHash->size = SwapBytes16 (RecvBuffer.PolicyHash.size);
  if (PolicyHash->size > sizeof(TPMU_HA)) {
    DEBUG ((DEBUG_ERROR, "Tpm2PolicyGetDigest - PolicyHash->size error %x\n", PolicyHash->size));
    return EFI_DEVICE_ERROR;
  }

  CopyMem (PolicyHash->buffer, &RecvBuffer.PolicyHash.buffer, PolicyHash->size);

  return EFI_SUCCESS;
}