/** @file
 This file include all platform action which can be customized by IBV/OEM.

;******************************************************************************
;* Copyright (c) 2014 - 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 <UefiBootManagerLib.h>
#include <Guid/AdmiSecureBoot.h>
#include <Uefi.h>
#include "AmtLockUsbConInDxe.h"

UINT8 gTDSBlockConInProtocol = 0;
BOOLEAN mLockConsole = FALSE;

EFI_DRIVER_BINDING_PROTOCOL gAmtUsbConInLockrDriver = {
  AmtUsbConInLockrDriverBindingSupported,
  AmtUsbConInLockrDriverBindingStart,
  AmtUsbConInLockrDriverBindingStop,
  0x18,
  NULL,
  NULL
};

BOOLEAN
IsUsbKeyboard (
  IN  EFI_USB_IO_PROTOCOL       *UsbIo
  )
{
  EFI_STATUS                    Status;
  EFI_USB_INTERFACE_DESCRIPTOR  InterfaceDescriptor;

  //
  // Get the Default interface descriptor, currently we
  // assume it is interface 1
  //
  Status = UsbIo->UsbGetInterfaceDescriptor (
                    UsbIo,
                    &InterfaceDescriptor
                    );

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

  if (InterfaceDescriptor.InterfaceClass == CLASS_HID &&
      InterfaceDescriptor.InterfaceSubClass == SUBCLASS_BOOT &&
      InterfaceDescriptor.InterfaceProtocol == PROTOCOL_KEYBOARD
      ) {

    return TRUE;
  }

  return FALSE;
}

EFI_STATUS
EFIAPI
AmtUsbConInLockrDriverBindingSupported (
  IN EFI_DRIVER_BINDING_PROTOCOL    *This,
  IN EFI_HANDLE                     Controller,
  IN EFI_DEVICE_PATH_PROTOCOL       *RemainingDevicePath
  )
{
  EFI_STATUS                    OpenStatus;
  EFI_USB_IO_PROTOCOL           *UsbIo;
  EFI_STATUS                    Status;
  EFI_DEV_PATH_PTR              Node;
  EFI_DEVICE_PATH_PROTOCOL      *DevicePath;
  EFI_USB_CORE_PROTOCOL         *UsbCore;
  //
  // Check Device Path
  //
  if (RemainingDevicePath != NULL) {
    Node.DevPath = RemainingDevicePath;
    if (Node.DevPath->Type != MESSAGING_DEVICE_PATH ||
        Node.DevPath->SubType != MSG_USB_CLASS_DP   ||
        DevicePathNodeLength(Node.DevPath) != sizeof(USB_CLASS_DEVICE_PATH) ||
        Node.UsbClass->DeviceClass != CLASS_HID) {
      return EFI_UNSUPPORTED;
    }
  }
  //
  // Check if USB_IO protocol is attached on the controller handle.
  //
  OpenStatus = gBS->OpenProtocol (
                      Controller,
                      &gEfiUsbIoProtocolGuid,
                      (VOID **)&UsbIo,
                      This->DriverBindingHandle,
                      Controller,
                      EFI_OPEN_PROTOCOL_BY_DRIVER
                      );
  if (EFI_ERROR (OpenStatus)) {
    return OpenStatus;
  }

  //
  // Get the device path for CheckIgnoredDevice
  //
  Status = gBS->HandleProtocol (
                  Controller,
                  &gEfiDevicePathProtocolGuid,
                  (VOID **)&DevicePath
                  );
  if (EFI_ERROR (Status)) {
    goto ErrorExit;
  }
  //
  // Setup the data for UsbLegacy
  //
  Status = gBS->LocateProtocol (
                  &gEfiUsbCoreProtocolGuid,
                  NULL,
                  (VOID **)&UsbCore
                  );
  if (EFI_ERROR (Status)) {
    goto ErrorExit;
  }
#if (KEEP_USBIO_FOR_IGNORED_DEVICE == 1)
  //
  // Filter out the USB devices which in the UsbIgnoreDevice list
  //
  Status = UsbCore->CheckIgnoredDevice(DevicePath, UsbIo);
  if (EFI_ERROR (Status)) {
    goto ErrorExit;
  }
#endif
  //
  // Use the USB I/O protocol interface to check whether the Controller is
  // the Keyboard controller that can be managed by this driver.
  //
  Status = EFI_SUCCESS;

  if (!IsUsbKeyboard (UsbIo)) {
    Status = EFI_UNSUPPORTED;
  }
ErrorExit:
  gBS->CloseProtocol (
        Controller,
        &gEfiUsbIoProtocolGuid,
        This->DriverBindingHandle,
        Controller
        );

  return Status;
}

/**

  Start.

  @param  This                  EFI_DRIVER_BINDING_PROTOCOL
  @param  Controller            Controller handle
  @param  RemainingDevicePath   EFI_DEVICE_PATH_PROTOCOL

  @retval EFI_SUCCESS           Success
  @retval EFI_OUT_OF_RESOURCES  Can't allocate memory
  @retval EFI_UNSUPPORTED       The Start routine fail

**/
EFI_STATUS
EFIAPI
AmtUsbConInLockrDriverBindingStart (
  IN EFI_DRIVER_BINDING_PROTOCOL    *This,
  IN EFI_HANDLE                     Controller,
  IN EFI_DEVICE_PATH_PROTOCOL       *RemainingDevicePath
  )
{
  EFI_STATUS                    Status;
  EFI_USB_IO_PROTOCOL           *UsbIo;
  UINT8                         TDSBlockConIn;

  gBS->InstallMultipleProtocolInterfaces (&Controller, &gTDSBlockConInProtocolGuid, &gTDSBlockConInProtocol, NULL);

  if (PcdGetBool (PcdTDSBlockConInEnable) || mLockConsole) {
    //
    // It's TDS boot or MPM lock mode.
    // Open IsaIo protocol and install flag.
    //
    Status = EFI_SUCCESS;

    //
    // Open USB_IO Protocol
    //
    Status = gBS->OpenProtocol (
                    Controller,
                    &gEfiUsbIoProtocolGuid,
                    (VOID **)&UsbIo,
                    This->DriverBindingHandle,
                    Controller,
                    EFI_OPEN_PROTOCOL_BY_DRIVER
                    );
  }
  //
  // Create relationship for disconnecting.
  //
  Status = gBS->OpenProtocol (
                  Controller,
                  &gTDSBlockConInProtocolGuid,
                  (VOID **) &TDSBlockConIn,
                  This->DriverBindingHandle,
                  Controller,
                  EFI_OPEN_PROTOCOL_BY_DRIVER
                  );

  return Status;
}

/**

  Stop.

  @param  This                  EFI_DRIVER_BINDING_PROTOCOL
  @param  Controller            Controller handle
  @param  NumberOfChildren      Child handle number
  @param  ChildHandleBuffer     Child handle buffer

  @retval EFI_SUCCESS           Success
  @retval EFI_UNSUPPORTED       Can't support

**/
EFI_STATUS
EFIAPI
AmtUsbConInLockrDriverBindingStop (
  IN  EFI_DRIVER_BINDING_PROTOCOL    *This,
  IN  EFI_HANDLE                     Controller,
  IN  UINTN                          NumberOfChildren,
  IN  EFI_HANDLE                     *ChildHandleBuffer
  )
{
  if (mLockConsole) {    
    //
    // MPM lock mode. Don't unlock Input device.
    //
    return EFI_UNSUPPORTED;
  }
  gBS->CloseProtocol (
                Controller,
                &gEfiUsbIoProtocolGuid,
                This->DriverBindingHandle,
                Controller
                );

  gBS->CloseProtocol (
                Controller,
                &gTDSBlockConInProtocolGuid,
                This->DriverBindingHandle,
                Controller
                );
  gBS->UninstallMultipleProtocolInterfaces (Controller, &gTDSBlockConInProtocolGuid, &gTDSBlockConInProtocol, NULL);
  return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
AmtLockUsbConInDxeEntryPoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS                    Status;
  AMT_WRAPPER_PROTOCOL          *AmtWrapper;
  EFI_USB_CORE_PROTOCOL         *UsbCore;
  BOOLEAN                       InSmram;
  UINT8                         SpecialCommandParamHighByte;
  SETUP_DATA                    SystemConfiguration;
  UINTN                         VarSize;
  BOOLEAN                       TrustedDeviceSetupSetupVar;
  EFI_STATUS                    TDSStatus;
  EFI_STATUS                    LockConStatus;

  SpecialCommandParamHighByte = 0;
  TrustedDeviceSetupSetupVar  = FALSE;
  TDSStatus                   = EFI_SUCCESS;
  LockConStatus               = EFI_SUCCESS;
  AmtWrapper = NULL;

  Status = gBS->LocateProtocol (&gAmtWrapperProtocolGuid, NULL, (VOID **) &AmtWrapper);

  //
  // Locate UsbCore protocol
  //
  Status = gBS->LocateProtocol (
                  &gEfiUsbCoreProtocolGuid,
                  NULL,
                  (VOID **)&UsbCore
                  );
  if (EFI_ERROR (Status)) {
    return EFI_UNSUPPORTED;
  }
  //
  // Register module in DXE/SMM instance
  //
  UsbCore->ModuleRegistration (ImageHandle);
  //
  // Check the phase of instance
  //
  UsbCore->IsInSmm (&InSmram);
  if (!InSmram) {
    if (AmtWrapper != NULL) {
      if (AmtWrapper->IsConsoleLocked()) {
        mLockConsole = TRUE;
      } else {
        LockConStatus = EFI_UNSUPPORTED;
      }
    }    
    //
    // Check ME status
    //
    SpecialCommandParamHighByte = AsfGetSpecialCmdParamHighByte ();
    //
    // Check SCU TDS option
    //
    VarSize = sizeof (SETUP_DATA);
    Status = gRT->GetVariable (
                    L"Setup",
                    &gSetupVariableGuid,
                    NULL,
                    &VarSize,
                    &SystemConfiguration
                    );
    
    if (!EFI_ERROR (Status)) {
      TrustedDeviceSetupSetupVar = SystemConfiguration.TrustedDeviceSetupBoot;
    }
    
    if ((SpecialCommandParamHighByte & TDS_ENABLE_BITS) != TDS_ENABLE_BITS && !TrustedDeviceSetupSetupVar) {
      TDSStatus = EFI_UNSUPPORTED;
    } else {
      Status = PcdSetBoolS (PcdTDSBlockConInEnable, TRUE);
      if (EFI_ERROR (Status)) {
        TDSStatus = EFI_UNSUPPORTED;
      }
    }
    
    if (EFI_ERROR (LockConStatus) && EFI_ERROR (TDSStatus)) {
      return EFI_UNSUPPORTED;
    }
    //
    // Now in boot service, install protocols
    //
    Status = EfiLibInstallDriverBindingComponentName2 (
               ImageHandle,
               SystemTable,
               &gAmtUsbConInLockrDriver,
               ImageHandle,
               NULL,
               NULL
               );
  } else {
    return EFI_UNSUPPORTED;
  }

  return Status;
}