/** @file

;******************************************************************************
;* Copyright (c) 2018, 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 "HddSpinDownSmm.h"

EFI_L05_HDD_SPINDOWN_PROTOCOL           *mEfiL05HddSpinDownProtocol = NULL;

/**
  HDD Spin Down, STANDBY_IMMEDIATE command by AtaPassThru protocol.

  @param  AtaDevice                     A pointer to the EFI_ATA_PASS_THRU_PROTOCOL instance.
  @param  ControllerMode                Indicate that now is IDE mode or AHCI mode.
  @param  Port                          The port number of the ATA device to send the command.
  @param  PortMultiplierPort            The port multiplier port number of the ATA device to send the command.
                                        If there is no port multiplier, then specify 0.

  @retval EFI_SUCCESS                   The function completed successfully.
**/
EFI_STATUS
HddSpinDown (
  IN  EFI_ATA_PASS_THRU_PROTOCOL        *AtaDevice,
  IN  UINT8                             ControllerMode,
  IN  UINT16                            Port,
  IN  UINT16                            PortMultiplierPort
  )
{
  EFI_STATUS                            Status;
  EFI_ATA_PASS_THRU_COMMAND_PACKET      AtaPassThruCmdPacket;
  EFI_ATA_STATUS_BLOCK                  Asb;
  EFI_ATA_COMMAND_BLOCK                 Acb;

  //
  // Initialize Command packet
  //
  ZeroMem (&AtaPassThruCmdPacket, sizeof (EFI_ATA_PASS_THRU_COMMAND_PACKET));
  ZeroMem (&Asb, sizeof (EFI_ATA_STATUS_BLOCK));
  ZeroMem (&Acb, sizeof (EFI_ATA_COMMAND_BLOCK));

  AtaPassThruCmdPacket.Asb = &Asb;
  AtaPassThruCmdPacket.Acb = &Acb;

  if (ControllerMode == ATA_IDE_MODE) {
    //
    // IDE mode
    //
    AtaPassThruCmdPacket.Acb->AtaDeviceHead |= ((PortMultiplierPort << 4) | 0xE0);
  }

  //
  // Send "STANDBY IMMEDIATEˇ¨command to HDD.
  //
  AtaPassThruCmdPacket.Acb->AtaCommand = STANDBY_IMMEDIATE;
  AtaPassThruCmdPacket.Protocol = EFI_ATA_PASS_THRU_PROTOCOL_ATA_NON_DATA;
  AtaPassThruCmdPacket.Length = EFI_ATA_PASS_THRU_LENGTH_NO_DATA_TRANSFER;

  Status = AtaDevice->PassThru (
                        AtaDevice,
                        Port,
                        PortMultiplierPort,
                        &AtaPassThruCmdPacket,
                        NULL
                        );

  return Status;
}

/**
  HDD Spin Down All Ports.

  @param  None

  @retval EFI_SUCCESS                   The function completed successfully.
**/

EFI_STATUS
HddSpinDownAllPort (
  )
{
  EFI_STATUS                            Status;
  UINTN                                 HandleCount;
  UINTN                                 BufferSize;
  EFI_HANDLE                            *HandleBuffer;
  UINTN                                 Index;
  EFI_ATA_PASS_THRU_PROTOCOL            *AtaDevice;
  EFI_DEVICE_PATH_PROTOCOL              *DevicePath;
  EFI_DEVICE_PATH_PROTOCOL              *DevicePathNode;
  BOOLEAN                               DevicePathMatch;
  UINT8                                 ControllerMode;
  UINT16                                Port;
  UINT16                                PortMultiplierPort;

  HandleCount        = 0;
  BufferSize         = 0;
  HandleBuffer       = NULL;
  AtaDevice          = NULL;
  DevicePath         = NULL;
  DevicePathNode     = NULL;
  ControllerMode     = 0XFF;

  //
  // Get the target AtaPassThruProtocol.
  //
  Status = gSmst->SmmLocateHandle (
                    ByProtocol,
                    &gEfiAtaPassThruProtocolGuid,
                    NULL,
                    &BufferSize,
                    HandleBuffer
                    );

  if ((Status != EFI_SUCCESS) && (Status != EFI_BUFFER_TOO_SMALL)) {
    return EFI_NOT_FOUND;
  }

  if (Status == EFI_BUFFER_TOO_SMALL) {
    Status = gSmst->SmmAllocatePool (
                      EfiRuntimeServicesData,
                      BufferSize,
                      (VOID *) &HandleBuffer
                      );

    if (HandleBuffer == NULL) {
      return EFI_NOT_FOUND;
    }

    Status = gSmst->SmmLocateHandle (
                      ByProtocol,
                      &gEfiAtaPassThruProtocolGuid,
                      NULL,
                      &BufferSize,
                      HandleBuffer
                      );
  }

  HandleCount = BufferSize / sizeof (EFI_HANDLE);

  for (Index = 0; Index < HandleCount; Index++) {
    Status = gSmst->SmmHandleProtocol (
                      HandleBuffer[Index],
                      &gEfiAtaPassThruProtocolGuid,
                      &AtaDevice
                      );

    if (EFI_ERROR (Status)) {
      continue;
    }

    //
    // Go through all of the ports and portmultiplierports and try something if a device is present
    //
    Port = 0xFFFF;

    while (TRUE) {

      Status = AtaDevice->GetNextPort (AtaDevice, &Port);

      if (EFI_ERROR (Status)) {
        break;
      }

      PortMultiplierPort = 0xFFFF;

      while (TRUE) {

        Status = AtaDevice->GetNextDevice (AtaDevice, Port, &PortMultiplierPort);

        if (EFI_ERROR (Status)) {
          break;
        }

        Status = AtaDevice->BuildDevicePath (AtaDevice, Port, PortMultiplierPort, &DevicePath);

        if (EFI_ERROR (Status)) {
          continue;
        }

        DevicePathMatch = FALSE;
        DevicePathNode = DevicePath;

        while (!IsDevicePathEnd (DevicePathNode)) {
          if ((DevicePathType (DevicePathNode) == MESSAGING_DEVICE_PATH) &&
              (DevicePathSubType (DevicePathNode) == MSG_ATAPI_DP)) {
            //
            // IDE mode
            //
            ControllerMode = ATA_IDE_MODE;
            DevicePathMatch = TRUE;
            break;
          }

          if ((DevicePathType (DevicePathNode) == MESSAGING_DEVICE_PATH) &&
              (DevicePathSubType (DevicePathNode) == MSG_SATA_DP)) {
            //
            // AHCI mode
            //
            ControllerMode = ATA_AHCI_MODE;
            DevicePathMatch = TRUE;
            break;
          }

          DevicePathNode = NextDevicePathNode (DevicePathNode);
        }

        if (!DevicePathMatch) {
          continue;
        }

        Status = HddSpinDown (AtaDevice, ControllerMode, Port, PortMultiplierPort);
      }
    }
  }

  return EFI_SUCCESS;
}

/**
  A callback function for SMM HDD spin down.

  @param  DispatchHandle                Unused.
  @param  DispatchContext               Unused.
  @param  CommBuffer                    Unused.
  @param  CommBufferSize                Unused.

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval EFI_UNSUPPORTED               The feature is not supported.
  @retval EFI_INVALID_PARAMETER         The parameter used for operation is not valid.
  @retval EFI_ACCESS_DENIED             The operation needs password and the password is not correct.
  @retval EFI_UNSUPPORTED               An unexpected error occurred.
**/
EFI_STATUS
EFIAPI
L05SmmHddSpinDownCallback (
  IN     EFI_HANDLE                     DispatchHandle,
  IN     CONST VOID                     *DispatchContext,
  IN OUT VOID                           *CommBuffer,
  IN OUT UINTN                          *CommBufferSize
  )
{
  EFI_STATUS                            Status;

  if (mEfiL05HddSpinDownProtocol == NULL) {
    return EFI_UNSUPPORTED;
  }

  Status =  mEfiL05HddSpinDownProtocol->HddSpinDownAllPort ();

  return Status;
}

/**
  Install L05 HDD spin down protocol.

  @param  None

  @retval EFI_SUCCESS                   The function completed successfully.
**/
EFI_STATUS
EFIAPI
InstallL05HddSpindownProtocol (
  )
{
  EFI_STATUS                            Status;
  EFI_HANDLE                            Handle;

  Status = gSmst->SmmAllocatePool (
                    EfiRuntimeServicesData,
                    sizeof (EFI_L05_HDD_SPINDOWN_PROTOCOL),
                    &mEfiL05HddSpinDownProtocol
                    );

  if (Status != EFI_SUCCESS) {
    return EFI_OUT_OF_RESOURCES;
  }

  mEfiL05HddSpinDownProtocol->HddSpinDownAllPort = HddSpinDownAllPort;

  Handle = NULL;
  Status = gSmst->SmmInstallProtocolInterface (
                    &Handle,
                    &gEfiL05HddSpindownProtocolGuid,
                    EFI_NATIVE_INTERFACE,
                    mEfiL05HddSpinDownProtocol
                    );

  return Status;
}

/**
  HddSpindown SMM AtaPassThru Protocol callback function.

  @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                   This function execute successfully.
**/
EFI_STATUS
EFIAPI
HddSpindownAtaPassThruCallBack (
  IN CONST EFI_GUID                     *Protocol,
  IN       VOID                         *Interface,
  IN       EFI_HANDLE                   Handle
  )
{
  EFI_STATUS                            Status;

  Status = InstallL05HddSpindownProtocol ();

  return Status;
}

/**
  HDD spin down SMM entry

  @param  ImageHandle                   The firmware allocated handle for the UEFI image.
  @param  SystemTable                   A pointer to the EFI System Table.

  @retval EFI_SUCCESS                   The operation completed successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
L05HddSpinDownSmmEntry (
  IN EFI_HANDLE                         ImageHandle,
  IN EFI_SYSTEM_TABLE                   *SystemTable
  )
{
  EFI_STATUS                            Status;
  UINTN                                 BufferSize;
  EFI_HANDLE                            *HandleBuffer;
  VOID                                  *Registration;
  EFI_HANDLE                            SwHandle;
  EFI_SMM_SW_DISPATCH2_PROTOCOL         *SwDispatch;
  EFI_SMM_SW_REGISTER_CONTEXT           SwContext;
  EFI_HANDLE                            Handle;

  BufferSize   = 0;
  HandleBuffer = NULL;
  Registration = NULL;
  SwHandle     = NULL;
  SwDispatch   = NULL;
  Handle       = NULL;

  //
  // Check AtaPassThruProtocol.
  //
  Status = gSmst->SmmLocateHandle (
                    ByProtocol,
                    &gEfiAtaPassThruProtocolGuid,
                    NULL,
                    &BufferSize,
                    HandleBuffer
                    );

  if (Status == EFI_NOT_FOUND) {
    Status = gSmst->SmmRegisterProtocolNotify (
                      &gEfiAtaPassThruProtocolGuid,
                      HddSpindownAtaPassThruCallBack,
                      &Registration
                      );

  } else {
    Status = InstallL05HddSpindownProtocol ();
  }

  if (Status != EFI_SUCCESS) {
    return Status;
  }

  //
  // Get the Sw dispatch protocol
  //
  Status = gSmst->SmmLocateProtocol (
                    &gEfiSmmSwDispatch2ProtocolGuid,
                    NULL,
                    &SwDispatch
                    );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Register Software SMI function
  //
  SwContext.SwSmiInputValue = EFI_L05_HDD_SPIN_DOWN_CALLBACK;
  Status = SwDispatch->Register (
                         SwDispatch,
                         L05SmmHddSpinDownCallback,
                         &SwContext,
                         &SwHandle
                         );

  if (!EFI_ERROR (Status)) {
    //
    // Publish HddSpinDown SW SMI Ready Protocol
    //
    Status = gBS->InstallProtocolInterface (
                    &Handle,
                    &gEfiL05HddSpindownSwSmiReadyProtocolGuid,
                    EFI_NATIVE_INTERFACE,
                    NULL
                    );
  }
  return Status;
}