//**********************************************************************
//**********************************************************************
//**                                                                  **
//**        (C)Copyright 1985-2016, American Megatrends, Inc.         **
//**                                                                  **
//**                       All Rights Reserved.                       **
//**                                                                  **
//**      5555 Oakbrook Parkway, Suite 200, Norcross, GA 30093        **
//**                                                                  **
//**                       Phone: (770)-246-8600                      **
//**                                                                  **
//**********************************************************************
//**********************************************************************


//**********************************************************************
//<AMI_FHDR_START>
//
// Name:  AmiUsbSmmGlobalDataValidationLib.c
//
// Description:	Source code for the AmiUsbSmmGlobalDataValidationLib library class.
//
//<AMI_FHDR_END>
//**********************************************************************

#include <Token.h>
#include <AmiDxeLib.h>
#include <AmiBufferValidationLib.h>
#include <AmiUsbSmmGlobalDataValidationLib.h>
#include <Protocol/SmmAccess2.h>
#include <Protocol/SmmBase2.h>
#include "../rt/UsbDef.h"
#include "../rt/Xhci.h"

extern UINT32  mCrcTable[256];
extern EFI_STATUS EFIAPI RuntimeDriverCalculateCrc32(VOID *, UINTN, UINT32*);
extern VOID RuntimeDriverInitializeCrc32Table(VOID);

AMI_USB_SMM_PROTOCOL	    *gAmiUsbSmmProtocol = NULL;
EFI_SMM_SYSTEM_TABLE2       *gSmst2 = NULL;

typedef struct {
    UINT32                  UsbLegSupOffSet;
    XHCI_DCBAA              *DcbaaPtr;
    TRB_RING                *XfrRings;
    UINTN                   XfrTrbs;
    VOID					*DeviceContext;
    VOID					*InputContext;
    UINT64                  *ScratchBufEntry;
    UINT8		            ContextSize;
    UINT32                  HashValue;
} USB3_HC_CONSTANT_DATA;

typedef struct {
    UINT8           HcType;
    UINT8           OpRegOffset;
    UINT8           ExtCapPtr;
    UINT16          BusDevFuncNum;
    UINT16          AsyncListSize;
    UINT32          *FrameList;
    UINTN           BaseAddressSize;
    VOID            *UsbBusData;
    VOID            *IsocTds;
    UINT32          HashValue;
} USB_HC_CONSTANT_DATA;

typedef struct {
    HCD_HEADER                  HcdriverTable[MAX_HC_TYPES];
    DEV_DRIVER                  DevDriverTable[MAX_DEVICE_TYPES];
    DEV_DRIVER                  DelayedDrivers[MAX_DEVICE_TYPES];
    CALLBACK_FUNC               CallBackFunctionTable[MAX_CALLBACK_FUNCTION];
    VOID                        *EfiKeyboardBuffer;
    UINT8                       *MemBlockStart;
    UINT32			            MemPages;
    UINT8                       *UsbTempBuffer;
    UINT8                       *UsbMassConsumeBuffer;
    UINT32                      *MemBlkSts;
    UINTN                       PciExpressBaseAddress;
    VOID                        *QueueCnnctDiscData;
    int                         QueueCnnctDiscMaxSize;
    VOID                        *IccQueueCnnctDiscData;
    int                         IccQueueCnnctDiscMaxSize;
    UINT32                      HashValue;
} USB_CONSTANT_DATA;

EFI_STATUS
CalculateUsbGlobalConstantDataCrc32(
    USB_GLOBAL_DATA         *UsbData,
    UINT32                  *Crc32
)
{
    USB_CONSTANT_DATA           UsbConStantData = {0};
    UINTN                       Index = 0;
    HC_STRUC                    *HcStruc;
    USB3_HOST_CONTROLLER        *Usb3Hc = NULL;
    UINTN                       HcInitializedNumber = 0;
    UINTN                       Usb3HcInitializedNumber = 0;
    UINTN                       HcStrucDataIndex = 0;
    UINTN                       Usb3HcStrucDataIndex = 0;
    USB_HC_CONSTANT_DATA        *HcStrucData = NULL;
    USB3_HC_CONSTANT_DATA       *Usb3HcStrucData = NULL;
    EFI_STATUS                  Status;
    UINT32                      Crc32Value[3] = {0};

    if (gAmiUsbSmmProtocol == NULL) {
        return EFI_NOT_READY;
    }
    
    MemCpy(UsbConStantData.HcdriverTable, UsbData->aHCDriverTable, sizeof(UsbConStantData.HcdriverTable));
    MemCpy(UsbConStantData.DevDriverTable, UsbData->aDevDriverTable, sizeof(UsbConStantData.DevDriverTable));
    MemCpy(UsbConStantData.DelayedDrivers, UsbData->aDelayedDrivers, sizeof(UsbConStantData.DelayedDrivers));
    MemCpy(UsbConStantData.CallBackFunctionTable, UsbData->aCallBackFunctionTable, sizeof(UsbConStantData.CallBackFunctionTable));
   
    UsbConStantData.EfiKeyboardBuffer = UsbData->EfiKeyboardBuffer;
    UsbConStantData.UsbTempBuffer = UsbData->fpUSBTempBuffer;
    UsbConStantData.UsbMassConsumeBuffer = UsbData->fpUSBMassConsumeBuffer;
    UsbConStantData.MemBlkSts = UsbData->aMemBlkSts;
    UsbConStantData.MemBlockStart = UsbData->fpMemBlockStart;
    UsbConStantData.MemPages = UsbData->MemPages;
    UsbConStantData.QueueCnnctDiscData = UsbData->QueueCnnctDisc.data;
    UsbConStantData.QueueCnnctDiscMaxSize = UsbData->QueueCnnctDisc.maxsize;
    UsbConStantData.IccQueueCnnctDiscData = UsbData->ICCQueueCnnctDisc.data;
    UsbConStantData.IccQueueCnnctDiscMaxSize = UsbData->ICCQueueCnnctDisc.maxsize;
    UsbConStantData.HashValue = gAmiUsbSmmProtocol->GlobalDataValidation.Crc32Hash;

    for (Index = 0; Index < UsbData->HcTableCount; Index++) {
        HcStruc = UsbData->HcTable[Index];
        if (HcStruc == NULL) {
            continue;
        }
        if (HcStruc->dHCFlag & (HC_STATE_RUNNING |HC_STATE_INITIALIZED)) {
            HcInitializedNumber++;
            if (HcStruc->usbbus_data) {
                Usb3HcInitializedNumber++;
            }
        }
    }

    if (HcInitializedNumber != 0) {
        Status = gSmst2->SmmAllocatePool(EfiRuntimeServicesData, 
                sizeof(USB_HC_CONSTANT_DATA) * HcInitializedNumber,
                (VOID**)&HcStrucData);

        if (EFI_ERROR(Status)) {
            return Status;
        }
        MemSet(HcStrucData, sizeof(USB_HC_CONSTANT_DATA) * HcInitializedNumber, 0);
    }

    if (Usb3HcInitializedNumber != 0) {

        Status = gSmst2->SmmAllocatePool(EfiRuntimeServicesData, 
                sizeof(USB3_HC_CONSTANT_DATA) * Usb3HcInitializedNumber,
                (VOID**)&Usb3HcStrucData);

        if (EFI_ERROR(Status)) {
            return Status;
        }
        MemSet(Usb3HcStrucData, sizeof(USB3_HC_CONSTANT_DATA) * Usb3HcInitializedNumber, 0);
    }

    for (Index = 0; Index < UsbData->HcTableCount; Index++) {
        HcStruc = UsbData->HcTable[Index];
        if (HcStruc == NULL) {
            continue;
        }
        Status = AmiValidateMemoryBuffer((VOID*)HcStruc, sizeof(HC_STRUC));
        if (EFI_ERROR(Status)) {
            return Status;
        }        
        if (HcStruc->dHCFlag & (HC_STATE_RUNNING |HC_STATE_INITIALIZED)) {
            HcStrucData[HcStrucDataIndex].HcType = HcStruc->bHCType;
            HcStrucData[HcStrucDataIndex].AsyncListSize = HcStruc->wAsyncListSize;
            HcStrucData[HcStrucDataIndex].FrameList = HcStruc->fpFrameList;
            HcStrucData[HcStrucDataIndex].BaseAddressSize = HcStruc->BaseAddressSize;
            HcStrucData[HcStrucDataIndex].BusDevFuncNum = HcStruc->wBusDevFuncNum;
            HcStrucData[HcStrucDataIndex].OpRegOffset = HcStruc->bOpRegOffset;
            HcStrucData[HcStrucDataIndex].ExtCapPtr = HcStruc->bExtCapPtr;
            HcStrucData[HcStrucDataIndex].UsbBusData = HcStruc->usbbus_data;
            HcStrucData[HcStrucDataIndex].HashValue = gAmiUsbSmmProtocol->GlobalDataValidation.Crc32Hash;
            HcStrucDataIndex++;
            if (HcStruc->usbbus_data) {
                Usb3Hc = HcStruc->usbbus_data;
                Status = AmiValidateMemoryBuffer((VOID*)Usb3Hc, sizeof(USB3_HOST_CONTROLLER));
                if (EFI_ERROR(Status)) {
                    return Status;
                }
                if (Usb3Hc->CapRegs) {
                    Status = AmiValidateMemoryBuffer((VOID*)Usb3Hc->CapRegs, sizeof(XHCI_HC_CAP_REGS));
                    if (EFI_ERROR(Status)) {
                        return Status;
                    }
                }
                if (Usb3Hc->OpRegs) {
                    Status = AmiValidateMemoryBuffer((VOID*)Usb3Hc->OpRegs, sizeof(XHCI_HC_OP_REGS));
                    if (EFI_ERROR(Status)) {
                        return Status;
                    }
                }
                if (Usb3Hc->RtRegs) {
                    Status = AmiValidateMemoryBuffer((VOID*)Usb3Hc->RtRegs, 
                        (sizeof(UINT32) * 8) + (sizeof(XHCI_INTERRUPTER_REGS) * Usb3Hc->CapRegs->HcsParams1.MaxIntrs));
                    if (EFI_ERROR(Status)) {
                        return Status;
                    }
                }
                if (Usb3Hc->ExtLegCap) {
                    Status = AmiValidateMemoryBuffer((VOID*)Usb3Hc->ExtLegCap, sizeof(XHCI_EXT_LEG_CAP));
                    if (EFI_ERROR(Status)) {
                        return Status;
                    }
                }
                if (Usb3Hc->Usb2Protocol) {
                    Status = AmiValidateMemoryBuffer((VOID*)Usb3Hc->Usb2Protocol, sizeof(XHCI_EXT_PROTOCOL));
                    if (EFI_ERROR(Status)) {
                        return Status;
                    }
                }
                if (Usb3Hc->Usb3Protocol) {
                    Status = AmiValidateMemoryBuffer((VOID*)Usb3Hc->Usb3Protocol, sizeof(XHCI_EXT_PROTOCOL));
                    if (EFI_ERROR(Status)) {
                        return Status;
                    }
                }
                if (Usb3Hc->DbCapRegs) {
                    Status = AmiValidateMemoryBuffer((VOID*)Usb3Hc->DbCapRegs, sizeof(XHCI_DB_CAP_REGS));
                    if (EFI_ERROR(Status)) {
                        return Status;
                    }
                }
                Usb3HcStrucData[Usb3HcStrucDataIndex].DcbaaPtr = Usb3Hc->DcbaaPtr;
                Usb3HcStrucData[Usb3HcStrucDataIndex].XfrRings = Usb3Hc->XfrRings;
                Usb3HcStrucData[Usb3HcStrucDataIndex].XfrTrbs = Usb3Hc->XfrTrbs;
                Usb3HcStrucData[Usb3HcStrucDataIndex].DeviceContext = Usb3Hc->DeviceContext;
                Usb3HcStrucData[Usb3HcStrucDataIndex].InputContext = Usb3Hc->InputContext;
                Usb3HcStrucData[Usb3HcStrucDataIndex].ScratchBufEntry = Usb3Hc->ScratchBufEntry;
                Usb3HcStrucData[Usb3HcStrucDataIndex].ContextSize = Usb3Hc->ContextSize;
                Usb3HcStrucData[Usb3HcStrucDataIndex].HashValue = gAmiUsbSmmProtocol->GlobalDataValidation.Crc32Hash;
                Usb3HcStrucDataIndex++;
            }
        }
    }

    if (HcInitializedNumber != 0) {
    	Status = RuntimeDriverCalculateCrc32((UINT8*)HcStrucData, 
    			(UINTN)(sizeof(USB_HC_CONSTANT_DATA) * HcInitializedNumber), &Crc32Value[0]);
        Status = gSmst2->SmmFreePool(HcStrucData);
        if (EFI_ERROR(Status)) {
            return Status;
        }
    }

    if (Usb3HcInitializedNumber != 0) {
    	Status = RuntimeDriverCalculateCrc32((UINT8*)Usb3HcStrucData, 
    			(UINTN)(sizeof(USB3_HC_CONSTANT_DATA) * Usb3HcInitializedNumber), &Crc32Value[1]);
        Status = gSmst2->SmmFreePool(Usb3HcStrucData);
        if (EFI_ERROR(Status)) {
            return Status;
        }
    }
    Status = RuntimeDriverCalculateCrc32((UINT8*)&UsbConStantData, sizeof(UsbConStantData), &Crc32Value[2]);
    
    RuntimeDriverCalculateCrc32((UINT8*)Crc32Value, sizeof(Crc32Value), Crc32);
 
    return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
UpdateAmiUsbSmmGlobalDataCrc32(
    VOID *UsbGlobdaData
)
{
    EFI_STATUS  Status;
    UINT32      Crc32;

    
    Status = CalculateUsbGlobalConstantDataCrc32(UsbGlobdaData, &Crc32);

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

    gAmiUsbSmmProtocol->GlobalDataValidation.ConstantDataCrc32 = Crc32;
    
    return EFI_SUCCESS;
}

EFI_STATUS
CheckAmiUsbSmmGlobalDataCrc32(
    VOID *UsbGlobdaData
)
{
    EFI_STATUS  Status;
    UINT32      Crc32;

    Status = CalculateUsbGlobalConstantDataCrc32(UsbGlobdaData, &Crc32);

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

    if (gAmiUsbSmmProtocol->GlobalDataValidation.ConstantDataCrc32 != Crc32) {
        return EFI_ACCESS_DENIED;
    }

    return EFI_SUCCESS;
}


EFI_STATUS
EFIAPI
AmiUsbSmmGlobalDataValidation(
    VOID *UsbGlobalData
)
{
    USB_GLOBAL_DATA *UsbData;
    EFI_STATUS          Status;
    UINT8               *MouseInputBufferEnd;
    UINT8               *KbcharacterBufferEnd;
    UINT8               *KbcScanCodeBufferEnd;
    UINT8               *KbcUsbDataBufferEnd;
    UINTN               Index;
    DEV_INFO            *DevInfoTableEnd;
    HC_STRUC            *HcStruc;

    UsbData = (USB_GLOBAL_DATA*)UsbGlobalData;
    
    MouseInputBufferEnd = UsbData->aMouseInputBuffer + sizeof(UsbData->aMouseInputBuffer);

    if ((UsbData->fpMouseInputBufferHeadPtr < UsbData->aMouseInputBuffer) ||
        (UsbData->fpMouseInputBufferHeadPtr > MouseInputBufferEnd) ||
        (UsbData->fpMouseInputBufferTailPtr < UsbData->aMouseInputBuffer) ||
        (UsbData->fpMouseInputBufferTailPtr > MouseInputBufferEnd)) {
        return EFI_ACCESS_DENIED;
    }

    KbcUsbDataBufferEnd = UsbData->aKBCUsbDataBufferStart + sizeof(UsbData->aKBCUsbDataBufferStart);

    if ((UsbData->aKBCUsbDataBufferHead < UsbData->aKBCUsbDataBufferStart) ||
        (UsbData->aKBCUsbDataBufferHead > KbcUsbDataBufferEnd) ||
        (UsbData->aKBCUsbDataBufferTail < UsbData->aKBCUsbDataBufferStart) ||
        (UsbData->aKBCUsbDataBufferTail > KbcUsbDataBufferEnd)) {
        return EFI_ACCESS_DENIED;
    }

    KbcharacterBufferEnd = UsbData->aKBCCharacterBufferStart + sizeof(UsbData->aKBCCharacterBufferStart);

    if ((UsbData->fpKBCCharacterBufferHead < UsbData->aKBCCharacterBufferStart) ||
        (UsbData->fpKBCCharacterBufferHead > KbcharacterBufferEnd) ||
        (UsbData->fpKBCCharacterBufferTail < UsbData->aKBCCharacterBufferStart) ||
        (UsbData->fpKBCCharacterBufferTail > KbcharacterBufferEnd)) {
        return EFI_ACCESS_DENIED;
    }

    KbcScanCodeBufferEnd = UsbData->aKBCScanCodeBufferStart + sizeof(UsbData->aKBCScanCodeBufferStart);

    if ((UsbData->fpKBCScanCodeBufferPtr < UsbData->aKBCScanCodeBufferStart) ||
        (UsbData->fpKBCScanCodeBufferPtr > KbcScanCodeBufferEnd)) {
        return EFI_ACCESS_DENIED;
    }

    DevInfoTableEnd = UsbData->aDevInfoTable + sizeof(UsbData->aDevInfoTable);

    if (UsbData->fpKeyRepeatDevInfo != NULL) { 
        if ((UsbData->fpKeyRepeatDevInfo < UsbData->aDevInfoTable) ||
            (UsbData->fpKeyRepeatDevInfo > DevInfoTableEnd)) {
            return EFI_ACCESS_DENIED;
        }
    }

    for (Index = 0; Index < USB_DEV_HID_COUNT; Index++) {
        if (UsbData->aUSBKBDeviceTable[Index] != NULL) {
            if ((UsbData->aUSBKBDeviceTable[Index] < UsbData->aDevInfoTable) ||
                (UsbData->aUSBKBDeviceTable[Index] > DevInfoTableEnd)) {
                return EFI_ACCESS_DENIED;
            }
        }
    }

    if (UsbData->HcTableCount != 0) {
        Status = AmiValidateMemoryBuffer((VOID*)UsbData->HcTable, sizeof(HC_STRUC) * UsbData->HcTableCount);
        if (EFI_ERROR(Status)) {
            return EFI_ACCESS_DENIED;
        }
        for (Index = 0; Index < UsbData->HcTableCount; Index++) {
            HcStruc = UsbData->HcTable[Index];
            if (HcStruc == NULL) {
                continue;
            }
           //Check if MMIO address space of usb controller resides in SMRAM region. If yes, don't proceed.
            if (HcStruc->bHCType != USB_HC_UHCI) {
                if (HcStruc->BaseAddress != 0) {
                    Status = AmiValidateMmioBuffer((VOID*)HcStruc->BaseAddress, HcStruc->BaseAddressSize);
                    if (EFI_ERROR(Status)) {
                        TRACE((-1, "Usb Mmio address is invalid, it is in SMRAM\n"));
                        return EFI_ACCESS_DENIED;
                    }
                }
            }
        }
    }

    Status = CheckAmiUsbSmmGlobalDataCrc32(UsbData);
    
    if (EFI_ERROR(Status)) {
        return EFI_ACCESS_DENIED;
    }

    return EFI_SUCCESS;
}

EFI_STATUS
EFIAPI
UsbSmmProtocolCallback(
    CONST EFI_GUID  *Protocol,
    VOID            *Interface,
    EFI_HANDLE      Handle
)
{
    EFI_STATUS  Status;
    
    Status = gSmst2->SmmLocateProtocol(
                    &gAmiUsbSmmProtocolGuid,
                    NULL,
                    &gAmiUsbSmmProtocol);

    ASSERT_EFI_ERROR(Status);
    
    return Status;
}

EFI_STATUS
EFIAPI
InitAmiUsbSmmGlobalDataValidationLib(
        IN EFI_HANDLE ImageHandle, 
        IN EFI_SYSTEM_TABLE *SystemTable)
{
    EFI_STATUS  Status;
    VOID        *ProtocolNotifyRegistration;
    EFI_SMM_BASE2_PROTOCOL          *SmmBase2;

    InitAmiSmmLib(ImageHandle,SystemTable);

    Status = pBS->LocateProtocol(&gEfiSmmBase2ProtocolGuid, NULL, &SmmBase2);
    
    if (EFI_ERROR(Status)) { 
        return Status;
    }
    
    Status = SmmBase2->GetSmstLocation(SmmBase2, &gSmst2);
    
    if (EFI_ERROR(Status)) {
        return EFI_UNSUPPORTED;
    }
    
    RuntimeDriverInitializeCrc32Table();
    
    Status = gSmst2->SmmLocateProtocol(
                    &gAmiUsbSmmProtocolGuid,
                    NULL,
                    &gAmiUsbSmmProtocol);

    if (EFI_ERROR(Status)) {
        Status = gSmst2->SmmRegisterProtocolNotify(
                    &gAmiUsbSmmProtocolGuid,
                    UsbSmmProtocolCallback,
                    &ProtocolNotifyRegistration
                    );
    }
    
    return EFI_SUCCESS;
}

//**********************************************************************
//**********************************************************************
//**                                                                  **
//**        (C)Copyright 1985-2016, American Megatrends, Inc.         **
//**                                                                  **
//**                       All Rights Reserved.                       **
//**                                                                  **
//**      5555 Oakbrook Parkway, Suite 200, Norcross, GA 30093        **
//**                                                                  **
//**                       Phone: (770)-246-8600                      **
//**                                                                  **
//**********************************************************************
//**********************************************************************