/** @file

  Provides services to convert a BMP graphics image to a GOP BLT buffer and
  from a GOP BLT buffer to a BMP graphics image.

  Caution: This module requires additional review when modified.
  This module processes external input - BMP image.
  This external input must be validated carefully to avoid security issue such
  as buffer overflow, integer overflow.

  TranslateBmpToGopBlt() receives untrusted input and performs basic validation.

  Copyright (c) 2016-2017, Microsoft Corporation
  Copyright (c) 2018, Intel Corporation. All rights reserved.<BR>

  All rights reserved.
  Redistribution and use in source and binary forms, with or without
  modification, are permitted provided that the following conditions are met:
  1. Redistributions of source code must retain the above copyright notice,
  this list of conditions and the following disclaimer.
  2. Redistributions in binary form must reproduce the above copyright notice,
  this list of conditions and the following disclaimer in the documentation
  and/or other materials provided with the distribution.

  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
  IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
  INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
  BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
  OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
  ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

**/

#include <PiDxe.h>
#include <Library/DebugLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/SafeIntLib.h>
#include <IndustryStandard/Bmp.h>

#include <Library/BmpSupportLib.h>

//
// BMP Image header for an uncompressed 24-bit per pixel BMP image.
//
const BMP_IMAGE_HEADER  mBmpImageHeaderTemplate = {
  'B',    // CharB
  'M',    // CharM
  0,      // Size will be updated at runtime
  {0, 0}, // Reserved
  sizeof (BMP_IMAGE_HEADER), // ImageOffset
  sizeof (BMP_IMAGE_HEADER) - OFFSET_OF (BMP_IMAGE_HEADER, HeaderSize), // HeaderSize
  0,      // PixelWidth will be updated at runtime
  0,      // PixelHeight will be updated at runtime
  1,      // Planes
  24,     // BitPerPixel
  0,      // CompressionType
  0,      // ImageSize will be updated at runtime
  0,      // XPixelsPerMeter
  0,      // YPixelsPerMeter
  0,      // NumberOfColors
  0       // ImportantColors
};

/**
  Translate a *.BMP graphics image to a GOP blt buffer. If a NULL Blt buffer
  is passed in a GopBlt buffer will be allocated by this routine using
  EFI_BOOT_SERVICES.AllocatePool(). If a GopBlt buffer is passed in it will be
  used if it is big enough.

  @param[in]       BmpImage      Pointer to BMP file.
  @param[in]       BmpImageSize  Number of bytes in BmpImage.
  @param[in, out]  GopBlt        Buffer containing GOP version of BmpImage.
  @param[in, out]  GopBltSize    Size of GopBlt in bytes.
  @param[out]      PixelHeight   Height of GopBlt/BmpImage in pixels.
  @param[out]      PixelWidth    Width of GopBlt/BmpImage in pixels.

  @retval  RETURN_SUCCESS            GopBlt and GopBltSize are returned.
  @retval  RETURN_INVALID_PARAMETER  BmpImage is NULL.
  @retval  RETURN_INVALID_PARAMETER  GopBlt is NULL.
  @retval  RETURN_INVALID_PARAMETER  GopBltSize is NULL.
  @retval  RETURN_INVALID_PARAMETER  PixelHeight is NULL.
  @retval  RETURN_INVALID_PARAMETER  PixelWidth is NULL.
  @retval  RETURN_UNSUPPORTED        BmpImage is not a valid *.BMP image.
  @retval  RETURN_BUFFER_TOO_SMALL   The passed in GopBlt buffer is not big
                                     enough.  The required size is returned in
                                     GopBltSize.
  @retval  RETURN_OUT_OF_RESOURCES   The GopBlt buffer could not be allocated.

**/
RETURN_STATUS
EFIAPI
TranslateBmpToGopBlt (
  IN     VOID                           *BmpImage,
  IN     UINTN                          BmpImageSize,
  IN OUT EFI_GRAPHICS_OUTPUT_BLT_PIXEL  **GopBlt,
  IN OUT UINTN                          *GopBltSize,
  OUT    UINTN                          *PixelHeight,
  OUT    UINTN                          *PixelWidth
  )
{
  UINT8                          *Image;
  UINT8                          *ImageHeader;
  BMP_IMAGE_HEADER               *BmpHeader;
  BMP_COLOR_MAP                  *BmpColorMap;
  EFI_GRAPHICS_OUTPUT_BLT_PIXEL  *BltBuffer;
  EFI_GRAPHICS_OUTPUT_BLT_PIXEL  *Blt;
  UINT32                         BltBufferSize;
  UINTN                          Index;
  UINTN                          Height;
  UINTN                          Width;
  UINTN                          ImageIndex;
  UINT32                         DataSizePerLine;
  BOOLEAN                        IsAllocated;
  UINT32                         ColorMapNum;
  RETURN_STATUS                  Status;
  UINT32                         DataSize;
  UINT32                         Temp;

  if (BmpImage == NULL || GopBlt == NULL || GopBltSize == NULL) {
    return RETURN_INVALID_PARAMETER;
  }
  if (PixelHeight == NULL || PixelWidth == NULL) {
    return RETURN_INVALID_PARAMETER;
  }

  if (BmpImageSize < sizeof (BMP_IMAGE_HEADER)) {
    DEBUG ((DEBUG_ERROR, "TranslateBmpToGopBlt: BmpImageSize too small\n"));
    return RETURN_UNSUPPORTED;
  }

  BmpHeader = (BMP_IMAGE_HEADER *)BmpImage;

  if (BmpHeader->CharB != 'B' || BmpHeader->CharM != 'M') {
    DEBUG ((DEBUG_ERROR, "TranslateBmpToGopBlt: BmpHeader->Char fields incorrect\n"));
    return RETURN_UNSUPPORTED;
  }

  //
  // Doesn't support compress.
  //
  if (BmpHeader->CompressionType != 0) {
    DEBUG ((DEBUG_ERROR, "TranslateBmpToGopBlt: Compression Type unsupported.\n"));
    return RETURN_UNSUPPORTED;
  }

  //
  // Only support BITMAPINFOHEADER format.
  // BITMAPFILEHEADER + BITMAPINFOHEADER = BMP_IMAGE_HEADER
  //
  if (BmpHeader->HeaderSize != sizeof (BMP_IMAGE_HEADER) - OFFSET_OF (BMP_IMAGE_HEADER, HeaderSize)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateBmpToGopBlt: BmpHeader->Headership is not as expected.  Headersize is 0x%x\n",
      BmpHeader->HeaderSize
      ));
    return RETURN_UNSUPPORTED;
  }

  //
  // The data size in each line must be 4 byte alignment.
  //
  Status = SafeUint32Mult (
             BmpHeader->PixelWidth,
             BmpHeader->BitPerPixel,
             &DataSizePerLine
             );
  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateBmpToGopBlt: invalid BmpImage... PixelWidth:0x%x BitPerPixel:0x%x\n",
      BmpHeader->PixelWidth,
      BmpHeader->BitPerPixel
      ));
    return RETURN_UNSUPPORTED;
  }

  Status = SafeUint32Add (DataSizePerLine, 31, &DataSizePerLine);
  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateBmpToGopBlt: invalid BmpImage... DataSizePerLine:0x%x\n",
      DataSizePerLine
      ));

    return RETURN_UNSUPPORTED;
  }

  DataSizePerLine = (DataSizePerLine >> 3) &(~0x3);
  Status = SafeUint32Mult (
             DataSizePerLine,
             BmpHeader->PixelHeight,
             &BltBufferSize
             );

  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateBmpToGopBlt: invalid BmpImage... DataSizePerLine:0x%x PixelHeight:0x%x\n",
      DataSizePerLine, BmpHeader->PixelHeight
      ));

    return RETURN_UNSUPPORTED;
  }

  Status = SafeUint32Mult (
             BmpHeader->PixelHeight,
             DataSizePerLine,
             &DataSize
             );

  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateBmpToGopBlt: invalid BmpImage... PixelHeight:0x%x DataSizePerLine:0x%x\n",
      BmpHeader->PixelHeight, DataSizePerLine
      ));

    return RETURN_UNSUPPORTED;
  }

  if ((BmpHeader->Size != BmpImageSize) ||
      (BmpHeader->Size < BmpHeader->ImageOffset) ||
      (BmpHeader->Size - BmpHeader->ImageOffset != DataSize)) {

    DEBUG ((DEBUG_ERROR, "TranslateBmpToGopBlt: invalid BmpImage... \n"));
    DEBUG ((DEBUG_ERROR, "   BmpHeader->Size: 0x%x\n", BmpHeader->Size));
    DEBUG ((DEBUG_ERROR, "   BmpHeader->ImageOffset: 0x%x\n", BmpHeader->ImageOffset));
    DEBUG ((DEBUG_ERROR, "   BmpImageSize: 0x%lx\n", (UINTN)BmpImageSize));
    DEBUG ((DEBUG_ERROR, "   DataSize: 0x%lx\n", (UINTN)DataSize));

    return RETURN_UNSUPPORTED;
  }

  //
  // Calculate Color Map offset in the image.
  //
  Image = BmpImage;
  BmpColorMap = (BMP_COLOR_MAP *)(Image + sizeof (BMP_IMAGE_HEADER));
  if (BmpHeader->ImageOffset < sizeof (BMP_IMAGE_HEADER)) {
    return RETURN_UNSUPPORTED;
  }

  if (BmpHeader->ImageOffset > sizeof (BMP_IMAGE_HEADER)) {
    switch (BmpHeader->BitPerPixel) {
    case 1:
      ColorMapNum = 2;
      break;
    case 4:
      ColorMapNum = 16;
      break;
    case 8:
      ColorMapNum = 256;
      break;
    default:
      ColorMapNum = 0;
      break;
    }
    //
    // BMP file may has padding data between the bmp header section and the
    // bmp data section.
    //
    if (BmpHeader->ImageOffset - sizeof (BMP_IMAGE_HEADER) < sizeof (BMP_COLOR_MAP) * ColorMapNum) {
      return RETURN_UNSUPPORTED;
    }
  }

  //
  // Calculate graphics image data address in the image
  //
  Image = ((UINT8 *)BmpImage) + BmpHeader->ImageOffset;
  ImageHeader = Image;

  //
  // Calculate the BltBuffer needed size.
  //
  Status = SafeUint32Mult (
             BmpHeader->PixelWidth,
             BmpHeader->PixelHeight,
             &BltBufferSize
             );

  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateBmpToGopBlt: invalid BltBuffer needed size... PixelWidth:0x%x PixelHeight:0x%x\n",
      BltBufferSize
      ));

    return RETURN_UNSUPPORTED;
  }

  Temp = BltBufferSize;
  Status = SafeUint32Mult (
             BltBufferSize,
             sizeof (EFI_GRAPHICS_OUTPUT_BLT_PIXEL),
             &BltBufferSize
             );

  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateBmpToGopBlt: invalid BltBuffer needed size... BltBufferSize:0x%lx struct size:0x%x\n",
      Temp, sizeof (EFI_GRAPHICS_OUTPUT_BLT_PIXEL)
      ));

    return RETURN_UNSUPPORTED;
  }

  IsAllocated = FALSE;
  if (*GopBlt == NULL) {
    //
    // GopBlt is not allocated by caller.
    //
    DEBUG ((DEBUG_INFO, "Bmp Support: Allocating 0x%X bytes of memory\n", BltBufferSize));
    *GopBltSize = (UINTN)BltBufferSize;
    *GopBlt = AllocatePool (*GopBltSize);
    IsAllocated = TRUE;
    if (*GopBlt == NULL) {
      return RETURN_OUT_OF_RESOURCES;
    }
  } else {
    //
    // GopBlt has been allocated by caller.
    //
    if (*GopBltSize < (UINTN)BltBufferSize) {
      *GopBltSize = (UINTN)BltBufferSize;
      return RETURN_BUFFER_TOO_SMALL;
    }
  }

  *PixelWidth  = BmpHeader->PixelWidth;
  *PixelHeight = BmpHeader->PixelHeight;
  DEBUG ((DEBUG_INFO, "BmpHeader->ImageOffset 0x%X\n", BmpHeader->ImageOffset));
  DEBUG ((DEBUG_INFO, "BmpHeader->PixelWidth 0x%X\n", BmpHeader->PixelWidth));
  DEBUG ((DEBUG_INFO, "BmpHeader->PixelHeight 0x%X\n", BmpHeader->PixelHeight));
  DEBUG ((DEBUG_INFO, "BmpHeader->BitPerPixel 0x%X\n", BmpHeader->BitPerPixel));
  DEBUG ((DEBUG_INFO, "BmpHeader->ImageSize 0x%X\n", BmpHeader->ImageSize));
  DEBUG ((DEBUG_INFO, "BmpHeader->HeaderSize 0x%X\n", BmpHeader->HeaderSize));
  DEBUG ((DEBUG_INFO, "BmpHeader->Size 0x%X\n", BmpHeader->Size));

  //
  // Translate image from BMP to Blt buffer format
  //
  BltBuffer = *GopBlt;
  for (Height = 0; Height < BmpHeader->PixelHeight; Height++) {
    Blt = &BltBuffer[ (BmpHeader->PixelHeight - Height - 1) * BmpHeader->PixelWidth];
    for (Width = 0; Width < BmpHeader->PixelWidth; Width++, Image++, Blt++) {
      switch (BmpHeader->BitPerPixel) {
      case 1:
        //
        // Translate 1-bit (2 colors) BMP to 24-bit color
        //
        for (Index = 0; Index < 8 && Width < BmpHeader->PixelWidth; Index++) {
          Blt->Red   = BmpColorMap[ ((*Image) >> (7 - Index)) & 0x1].Red;
          Blt->Green = BmpColorMap[ ((*Image) >> (7 - Index)) & 0x1].Green;
          Blt->Blue  = BmpColorMap[ ((*Image) >> (7 - Index)) & 0x1].Blue;
          Blt++;
          Width++;
        }

        Blt--;
        Width--;
        break;

      case 4:
        //
        // Translate 4-bit (16 colors) BMP Palette to 24-bit color
        //
        Index = (*Image) >> 4;
        Blt->Red   = BmpColorMap[Index].Red;
        Blt->Green = BmpColorMap[Index].Green;
        Blt->Blue  = BmpColorMap[Index].Blue;
        if (Width < (BmpHeader->PixelWidth - 1)) {
          Blt++;
          Width++;
          Index = (*Image) & 0x0f;
          Blt->Red   = BmpColorMap[Index].Red;
          Blt->Green = BmpColorMap[Index].Green;
          Blt->Blue  = BmpColorMap[Index].Blue;
        }
        break;

      case 8:
        //
        // Translate 8-bit (256 colors) BMP Palette to 24-bit color
        //
        Blt->Red   = BmpColorMap[*Image].Red;
        Blt->Green = BmpColorMap[*Image].Green;
        Blt->Blue  = BmpColorMap[*Image].Blue;
        break;

      case 24:
        //
        // It is 24-bit BMP.
        //
        Blt->Blue  = *Image++;
        Blt->Green = *Image++;
        Blt->Red   = *Image;
        break;

      case 32:
        //
        //Conver 32 bit to 24bit bmp - just ignore the final byte of each pixel
        Blt->Blue  = *Image++;
        Blt->Green = *Image++;
        Blt->Red   = *Image++;
        break;

      default:
        //
        // Other bit format BMP is not supported.
        //
        if (IsAllocated) {
          FreePool (*GopBlt);
          *GopBlt = NULL;
        }
        DEBUG ((DEBUG_ERROR, "Bmp Bit format not supported.  0x%X\n", BmpHeader->BitPerPixel));
        return RETURN_UNSUPPORTED;
        break;
      };

    }

    ImageIndex = (UINTN)Image - (UINTN)ImageHeader;
    if ((ImageIndex % 4) != 0) {
      //
      // Bmp Image starts each row on a 32-bit boundary!
      //
      Image = Image + (4 - (ImageIndex % 4));
    }
  }

  return RETURN_SUCCESS;
}

/**
  Translate a GOP blt buffer to an uncompressed 24-bit per pixel BMP graphics
  image. If a NULL BmpImage is passed in a BmpImage buffer will be allocated by
  this routine using EFI_BOOT_SERVICES.AllocatePool(). If a BmpImage buffer is
  passed in it will be used if it is big enough.

  @param [in]      GopBlt        Pointer to GOP blt buffer.
  @param [in]      PixelHeight   Height of GopBlt/BmpImage in pixels.
  @param [in]      PixelWidth    Width of GopBlt/BmpImage in pixels.
  @param [in, out] BmpImage      Buffer containing BMP version of GopBlt.
  @param [in, out] BmpImageSize  Size of BmpImage in bytes.

  @retval RETURN_SUCCESS            BmpImage and BmpImageSize are returned.
  @retval RETURN_INVALID_PARAMETER  GopBlt is NULL.
  @retval RETURN_INVALID_PARAMETER  BmpImage is NULL.
  @retval RETURN_INVALID_PARAMETER  BmpImageSize is NULL.
  @retval RETURN_UNSUPPORTED        GopBlt cannot be converted to a *.BMP image.
  @retval RETURN_BUFFER_TOO_SMALL   The passed in BmpImage buffer is not big
                                    enough.  The required size is returned in
                                    BmpImageSize.
  @retval RETURN_OUT_OF_RESOURCES   The BmpImage buffer could not be allocated.

**/
RETURN_STATUS
EFIAPI
TranslateGopBltToBmp (
  IN     EFI_GRAPHICS_OUTPUT_BLT_PIXEL  *GopBlt,
  IN     UINT32                         PixelHeight,
  IN     UINT32                         PixelWidth,
  IN OUT VOID                           **BmpImage,
  IN OUT UINT32                         *BmpImageSize
  )
{
  RETURN_STATUS                  Status;
  UINT32                         PaddingSize;
  UINT32                         BmpSize;
  BMP_IMAGE_HEADER               *BmpImageHeader;
  UINT8                          *Image;
  UINTN                          Col;
  UINTN                          Row;
  EFI_GRAPHICS_OUTPUT_BLT_PIXEL  *BltPixel;

  if (GopBlt == NULL || BmpImage == NULL || BmpImageSize == NULL) {
    return RETURN_INVALID_PARAMETER;
  }

  //
  // Allocate memory for BMP file.
  //
  PaddingSize = PixelWidth & 0x3;

  //
  // First check PixelWidth * 3 + PaddingSize doesn't overflow
  //
  Status = SafeUint32Mult (PixelWidth, 3, &BmpSize);
  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateGopBltToBmp: GopBlt is too large. PixelHeight:0x%x PixelWidth:0x%x\n",
      PixelHeight,
      PixelWidth
      ));
    return RETURN_UNSUPPORTED;
  }
  Status = SafeUint32Add (BmpSize, PaddingSize, &BmpSize);
  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateGopBltToBmp: GopBlt is too large. PixelHeight:0x%x PixelWidth:0x%x\n",
      PixelHeight,
      PixelWidth
      ));
    return RETURN_UNSUPPORTED;
  }

  //
  // Second check (mLogoWidth * 3 + PaddingSize) * mLogoHeight + sizeof (BMP_IMAGE_HEADER) doesn't overflow
  //
  Status = SafeUint32Mult (BmpSize, PixelHeight, &BmpSize);
  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateGopBltToBmp: GopBlt is too large. PixelHeight:0x%x PixelWidth:0x%x\n",
      PixelHeight,
      PixelWidth
      ));
    return RETURN_UNSUPPORTED;
  }
  Status = SafeUint32Add (BmpSize, sizeof (BMP_IMAGE_HEADER), &BmpSize);
  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_ERROR,
      "TranslateGopBltToBmp: GopBlt is too large. PixelHeight:0x%x PixelWidth:0x%x\n",
      PixelHeight,
      PixelWidth
      ));
    return RETURN_UNSUPPORTED;
  }

  //
  // The image should be stored in EfiBootServicesData, allowing the system to
  // reclaim the memory
  //
  if (*BmpImage == NULL) {
    *BmpImage = AllocateZeroPool (BmpSize);
    if (*BmpImage == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }
    *BmpImageSize = BmpSize;
  } else if (*BmpImageSize < BmpSize) {
    *BmpImageSize = BmpSize;
    return RETURN_BUFFER_TOO_SMALL;
  }

  BmpImageHeader = (BMP_IMAGE_HEADER *)*BmpImage;
  CopyMem (BmpImageHeader, &mBmpImageHeaderTemplate, sizeof (BMP_IMAGE_HEADER));
  BmpImageHeader->Size        = *BmpImageSize;
  BmpImageHeader->ImageSize   = *BmpImageSize - sizeof (BMP_IMAGE_HEADER);
  BmpImageHeader->PixelWidth  = PixelWidth;
  BmpImageHeader->PixelHeight = PixelHeight;

  //
  // Convert BLT buffer to BMP file.
  //
  Image = (UINT8 *)BmpImageHeader + sizeof (BMP_IMAGE_HEADER);
  for (Row = 0; Row < PixelHeight; Row++) {
    BltPixel = &GopBlt[(PixelHeight - Row - 1) * PixelWidth];

    for (Col = 0; Col < PixelWidth; Col++) {
      *Image++ = BltPixel->Blue;
      *Image++ = BltPixel->Green;
      *Image++ = BltPixel->Red;
      BltPixel++;
    }

    //
    // Padding for 4 byte alignment.
    //
    Image += PaddingSize;
  }

  return RETURN_SUCCESS;
}