/** @file

  Copyright (c) 2013-2014, ARM Ltd. All rights reserved.<BR>
  Copyright (c) 2017, Linaro. All rights reserved.

  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include <libfdt.h>
#include <Library/AndroidBootImgLib.h>
#include <Library/PrintLib.h>
#include <Library/DevicePathLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiLib.h>

#include <Protocol/AndroidBootImg.h>
#include <Protocol/LoadFile2.h>
#include <Protocol/LoadedImage.h>

#include <Guid/LinuxEfiInitrdMedia.h>

#define FDT_ADDITIONAL_ENTRIES_SIZE  0x400

typedef struct {
  MEMMAP_DEVICE_PATH          Node1;
  EFI_DEVICE_PATH_PROTOCOL    End;
} MEMORY_DEVICE_PATH;

typedef struct {
  VENDOR_DEVICE_PATH          VendorMediaNode;
  EFI_DEVICE_PATH_PROTOCOL    EndNode;
} RAMDISK_DEVICE_PATH;

STATIC ANDROID_BOOTIMG_PROTOCOL  *mAndroidBootImg;
STATIC VOID                      *mRamdiskData          = NULL;
STATIC UINTN                     mRamdiskSize           = 0;
STATIC EFI_HANDLE                mRamDiskLoadFileHandle = NULL;

STATIC CONST MEMORY_DEVICE_PATH  mMemoryDevicePathTemplate =
{
  {
    {
      HARDWARE_DEVICE_PATH,
      HW_MEMMAP_DP,
      {
        (UINT8)(sizeof (MEMMAP_DEVICE_PATH)),
        (UINT8)((sizeof (MEMMAP_DEVICE_PATH)) >> 8),
      },
    }, // Header
    0, // StartingAddress (set at runtime)
    0  // EndingAddress   (set at runtime)
  }, // Node1
  {
    END_DEVICE_PATH_TYPE,
    END_ENTIRE_DEVICE_PATH_SUBTYPE,
    { sizeof (EFI_DEVICE_PATH_PROTOCOL), 0 }
  } // End
};

STATIC CONST RAMDISK_DEVICE_PATH  mRamdiskDevicePath =
{
  {
    {
      MEDIA_DEVICE_PATH,
      MEDIA_VENDOR_DP,
      { sizeof (VENDOR_DEVICE_PATH),       0 }
    },
    LINUX_EFI_INITRD_MEDIA_GUID
  },
  {
    END_DEVICE_PATH_TYPE,
    END_ENTIRE_DEVICE_PATH_SUBTYPE,
    { sizeof (EFI_DEVICE_PATH_PROTOCOL), 0 }
  }
};

/**
  Causes the driver to load a specified file.

  @param  This       Protocol instance pointer.
  @param  FilePath   The device specific path of the file to load.
  @param  BootPolicy Should always be FALSE.
  @param  BufferSize On input the size of Buffer in bytes. On output with a return
                     code of EFI_SUCCESS, the amount of data transferred to
                     Buffer. On output with a return code of EFI_BUFFER_TOO_SMALL,
                     the size of Buffer required to retrieve the requested file.
  @param  Buffer     The memory buffer to transfer the file to. IF Buffer is NULL,
                     then no the size of the requested file is returned in
                     BufferSize.

  @retval EFI_SUCCESS           The file was loaded.
  @retval EFI_UNSUPPORTED       BootPolicy is TRUE.
  @retval EFI_INVALID_PARAMETER FilePath is not a valid device path, or
                                BufferSize is NULL.
  @retval EFI_NO_MEDIA          No medium was present to load the file.
  @retval EFI_DEVICE_ERROR      The file was not loaded due to a device error.
  @retval EFI_NO_RESPONSE       The remote system did not respond.
  @retval EFI_NOT_FOUND         The file was not found
  @retval EFI_ABORTED           The file load process was manually canceled.
  @retval EFI_BUFFER_TOO_SMALL  The BufferSize is too small to read the current
                                directory entry. BufferSize has been updated with
                                the size needed to complete the request.


**/
EFI_STATUS
EFIAPI
AndroidBootImgLoadFile2 (
  IN EFI_LOAD_FILE2_PROTOCOL   *This,
  IN EFI_DEVICE_PATH_PROTOCOL  *FilePath,
  IN BOOLEAN                   BootPolicy,
  IN OUT UINTN                 *BufferSize,
  IN VOID                      *Buffer OPTIONAL
  )

{
  // Verify if the valid parameters
  if ((This == NULL) ||
      (BufferSize == NULL) ||
      (FilePath == NULL) ||
      !IsDevicePathValid (FilePath, 0))
  {
    return EFI_INVALID_PARAMETER;
  }

  if (BootPolicy) {
    return EFI_UNSUPPORTED;
  }

  // Check if the given buffer size is big enough
  // EFI_BUFFER_TOO_SMALL to allow caller to allocate a bigger buffer
  if (mRamdiskSize == 0) {
    return EFI_NOT_FOUND;
  }

  if ((Buffer == NULL) || (*BufferSize < mRamdiskSize)) {
    *BufferSize = mRamdiskSize;
    return EFI_BUFFER_TOO_SMALL;
  }

  // Copy InitRd
  CopyMem (Buffer, mRamdiskData, mRamdiskSize);
  *BufferSize = mRamdiskSize;

  return EFI_SUCCESS;
}

///
/// Load File Protocol instance
///
STATIC EFI_LOAD_FILE2_PROTOCOL  mAndroidBootImgLoadFile2 = {
  AndroidBootImgLoadFile2
};

EFI_STATUS
AndroidBootImgGetImgSize (
  IN  VOID   *BootImg,
  OUT UINTN  *ImgSize
  )
{
  ANDROID_BOOTIMG_HEADER  *Header;

  Header = (ANDROID_BOOTIMG_HEADER *)BootImg;

  if (AsciiStrnCmp (
        (CONST CHAR8 *)Header->BootMagic,
        ANDROID_BOOT_MAGIC,
        ANDROID_BOOT_MAGIC_LENGTH
        ) != 0)
  {
    return EFI_INVALID_PARAMETER;
  }

  /* The page size is not specified, but it should be power of 2 at least */
  ASSERT (IS_VALID_ANDROID_PAGE_SIZE (Header->PageSize));

  /* Get real size of abootimg */
  *ImgSize = ALIGN_VALUE (Header->KernelSize, Header->PageSize) +
             ALIGN_VALUE (Header->RamdiskSize, Header->PageSize) +
             ALIGN_VALUE (Header->SecondStageBootloaderSize, Header->PageSize) +
             Header->PageSize;
  return EFI_SUCCESS;
}

EFI_STATUS
AndroidBootImgGetKernelInfo (
  IN  VOID   *BootImg,
  OUT VOID   **Kernel,
  OUT UINTN  *KernelSize
  )
{
  ANDROID_BOOTIMG_HEADER  *Header;

  Header = (ANDROID_BOOTIMG_HEADER *)BootImg;

  if (AsciiStrnCmp (
        (CONST CHAR8 *)Header->BootMagic,
        ANDROID_BOOT_MAGIC,
        ANDROID_BOOT_MAGIC_LENGTH
        ) != 0)
  {
    return EFI_INVALID_PARAMETER;
  }

  if (Header->KernelSize == 0) {
    return EFI_NOT_FOUND;
  }

  ASSERT (IS_VALID_ANDROID_PAGE_SIZE (Header->PageSize));

  *KernelSize = Header->KernelSize;
  *Kernel     = (VOID *)((UINTN)BootImg + Header->PageSize);
  return EFI_SUCCESS;
}

EFI_STATUS
AndroidBootImgGetRamdiskInfo (
  IN  VOID   *BootImg,
  OUT VOID   **Ramdisk,
  OUT UINTN  *RamdiskSize
  )
{
  ANDROID_BOOTIMG_HEADER  *Header;

  Header = (ANDROID_BOOTIMG_HEADER *)BootImg;

  if (AsciiStrnCmp (
        (CONST CHAR8 *)Header->BootMagic,
        ANDROID_BOOT_MAGIC,
        ANDROID_BOOT_MAGIC_LENGTH
        ) != 0)
  {
    return EFI_INVALID_PARAMETER;
  }

  ASSERT (IS_VALID_ANDROID_PAGE_SIZE (Header->PageSize));

  *RamdiskSize = Header->RamdiskSize;

  if (Header->RamdiskSize != 0) {
    *Ramdisk = (VOID *)((INTN)BootImg
                        + Header->PageSize
                        + ALIGN_VALUE (Header->KernelSize, Header->PageSize));
  }

  return EFI_SUCCESS;
}

EFI_STATUS
AndroidBootImgGetSecondBootLoaderInfo (
  IN  VOID   *BootImg,
  OUT VOID   **Second,
  OUT UINTN  *SecondSize
  )
{
  ANDROID_BOOTIMG_HEADER  *Header;

  Header = (ANDROID_BOOTIMG_HEADER *)BootImg;

  if (AsciiStrnCmp (
        (CONST CHAR8 *)Header->BootMagic,
        ANDROID_BOOT_MAGIC,
        ANDROID_BOOT_MAGIC_LENGTH
        ) != 0)
  {
    return EFI_INVALID_PARAMETER;
  }

  ASSERT (IS_VALID_ANDROID_PAGE_SIZE (Header->PageSize));

  *SecondSize = Header->SecondStageBootloaderSize;

  if (Header->SecondStageBootloaderSize != 0) {
    *Second = (VOID *)((UINTN)BootImg
                       + Header->PageSize
                       + ALIGN_VALUE (Header->KernelSize, Header->PageSize)
                       + ALIGN_VALUE (Header->RamdiskSize, Header->PageSize));
  }

  return EFI_SUCCESS;
}

EFI_STATUS
AndroidBootImgGetKernelArgs (
  IN  VOID   *BootImg,
  OUT CHAR8  *KernelArgs
  )
{
  ANDROID_BOOTIMG_HEADER  *Header;

  Header = (ANDROID_BOOTIMG_HEADER *)BootImg;
  AsciiStrnCpyS (
    KernelArgs,
    ANDROID_BOOTIMG_KERNEL_ARGS_SIZE,
    Header->KernelArgs,
    ANDROID_BOOTIMG_KERNEL_ARGS_SIZE
    );

  return EFI_SUCCESS;
}

EFI_STATUS
AndroidBootImgGetFdt (
  IN  VOID  *BootImg,
  IN  VOID  **FdtBase
  )
{
  UINTN       SecondLoaderSize;
  EFI_STATUS  Status;

  /* Check whether FDT is located in second boot region as some vendor do so,
   * because second loader is never used as far as I know. */
  Status = AndroidBootImgGetSecondBootLoaderInfo (
             BootImg,
             FdtBase,
             &SecondLoaderSize
             );
  return Status;
}

EFI_STATUS
AndroidBootImgUpdateArgs (
  IN  VOID  *BootImg,
  OUT VOID  **KernelArgs
  )
{
  CHAR8       ImageKernelArgs[ANDROID_BOOTIMG_KERNEL_ARGS_SIZE];
  EFI_STATUS  Status;
  UINT32      NewKernelArgSize;

  // Get kernel arguments from Android boot image
  Status = AndroidBootImgGetKernelArgs (BootImg, ImageKernelArgs);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  NewKernelArgSize = ANDROID_BOOTIMG_KERNEL_ARGS_SIZE + PcdGet32 (PcdAndroidKernelCommandLineOverflow);
  *KernelArgs      = AllocateZeroPool (sizeof (CHAR16) * NewKernelArgSize);
  if (*KernelArgs == NULL) {
    DEBUG ((DEBUG_ERROR, "Fail to allocate memory\n"));
    return EFI_OUT_OF_RESOURCES;
  }

  AsciiStrToUnicodeStrS (
    ImageKernelArgs,
    *KernelArgs,
    NewKernelArgSize
    );
  // Append platform kernel arguments
  if (mAndroidBootImg->AppendArgs) {
    Status = mAndroidBootImg->AppendArgs (
                                *KernelArgs,
                                NewKernelArgSize
                                );
  }

  return Status;
}

EFI_STATUS
AndroidBootImgInstallLoadFile2 (
  IN  VOID   *RamdiskData,
  IN  UINTN  RamdiskSize
  )
{
  mRamDiskLoadFileHandle = NULL;
  mRamdiskData           = RamdiskData;
  mRamdiskSize           = RamdiskSize;
  return gBS->InstallMultipleProtocolInterfaces (
                &mRamDiskLoadFileHandle,
                &gEfiLoadFile2ProtocolGuid,
                &mAndroidBootImgLoadFile2,
                &gEfiDevicePathProtocolGuid,
                &mRamdiskDevicePath,
                NULL
                );
}

EFI_STATUS
AndroidBootImgUninstallLoadFile2 (
  VOID
  )
{
  EFI_STATUS  Status;

  Status       = EFI_SUCCESS;
  mRamdiskData = NULL;
  mRamdiskSize = 0;
  if (mRamDiskLoadFileHandle != NULL) {
    Status = gBS->UninstallMultipleProtocolInterfaces (
                    mRamDiskLoadFileHandle,
                    &gEfiLoadFile2ProtocolGuid,
                    &mAndroidBootImgLoadFile2,
                    &gEfiDevicePathProtocolGuid,
                    &mRamdiskDevicePath,
                    NULL
                    );
    mRamDiskLoadFileHandle = NULL;
  }

  return Status;
}

BOOLEAN
AndroidBootImgAcpiSupported (
  VOID
  )
{
  EFI_STATUS  Status;
  VOID        *AcpiTable;

  Status = EfiGetSystemConfigurationTable (&gEfiAcpiTableGuid, &AcpiTable);
  return !EFI_ERROR (Status);
}

EFI_STATUS
AndroidBootImgLocateFdt (
  IN  VOID  *BootImg,
  IN  VOID  **FdtBase
  )
{
  INTN        Err;
  EFI_STATUS  Status;

  Status = EfiGetSystemConfigurationTable (&gFdtTableGuid, FdtBase);
  if (!EFI_ERROR (Status)) {
    return EFI_SUCCESS;
  }

  Status = AndroidBootImgGetFdt (BootImg, FdtBase);
  if (EFI_ERROR (Status)) {
    return Status;
  }

  Err = fdt_check_header (*FdtBase);
  if (Err != 0) {
    DEBUG ((
      DEBUG_ERROR,
      "ERROR: Device Tree header not valid (Err:%d)\n",
      Err
      ));
    return EFI_INVALID_PARAMETER;
  }

  return EFI_SUCCESS;
}

INTN
AndroidBootImgGetChosenNode (
  IN  INTN  UpdatedFdtBase
  )
{
  INTN  ChosenNode;

  ChosenNode = fdt_subnode_offset ((CONST VOID *)UpdatedFdtBase, 0, "chosen");
  if (ChosenNode < 0) {
    ChosenNode = fdt_add_subnode ((VOID *)UpdatedFdtBase, 0, "chosen");
    if (ChosenNode < 0) {
      DEBUG ((DEBUG_ERROR, "Fail to find fdt node chosen!\n"));
      return 0;
    }
  }

  return ChosenNode;
}

EFI_STATUS
AndroidBootImgSetProperty64 (
  IN  INTN    UpdatedFdtBase,
  IN  INTN    ChosenNode,
  IN  CHAR8   *PropertyName,
  IN  UINT64  Val
  )
{
  INTN                 Err;
  struct fdt_property  *Property;
  int                  Len;

  Property = fdt_get_property_w (
               (VOID *)UpdatedFdtBase,
               ChosenNode,
               PropertyName,
               &Len
               );
  if ((NULL == Property) && (Len == -FDT_ERR_NOTFOUND)) {
    Val = cpu_to_fdt64 (Val);
    Err = fdt_appendprop (
            (VOID *)UpdatedFdtBase,
            ChosenNode,
            PropertyName,
            &Val,
            sizeof (UINT64)
            );
    if (Err) {
      DEBUG ((DEBUG_ERROR, "fdt_appendprop() fail: %a\n", fdt_strerror (Err)));
      return EFI_INVALID_PARAMETER;
    }
  } else if (Property != NULL) {
    Err = fdt_setprop_u64 (
            (VOID *)UpdatedFdtBase,
            ChosenNode,
            PropertyName,
            Val
            );
    if (Err) {
      DEBUG ((DEBUG_ERROR, "fdt_setprop_u64() fail: %a\n", fdt_strerror (Err)));
      return EFI_INVALID_PARAMETER;
    }
  } else {
    DEBUG ((DEBUG_ERROR, "Failed to set fdt Property %a\n", PropertyName));
    return EFI_INVALID_PARAMETER;
  }

  return EFI_SUCCESS;
}

EFI_STATUS
AndroidBootImgUpdateFdt (
  IN  VOID   *BootImg,
  IN  VOID   *FdtBase,
  IN  VOID   *RamdiskData,
  IN  UINTN  RamdiskSize
  )
{
  INTN                  ChosenNode, Err, NewFdtSize;
  EFI_STATUS            Status;
  EFI_PHYSICAL_ADDRESS  UpdatedFdtBase, NewFdtBase;

  NewFdtSize = (UINTN)fdt_totalsize (FdtBase)
               + FDT_ADDITIONAL_ENTRIES_SIZE;
  Status = gBS->AllocatePages (
                  AllocateAnyPages,
                  EfiBootServicesData,
                  EFI_SIZE_TO_PAGES (NewFdtSize),
                  &UpdatedFdtBase
                  );
  if (EFI_ERROR (Status)) {
    DEBUG ((
      DEBUG_WARN,
      "Warning: Failed to reallocate FDT, err %d.\n",
      Status
      ));
    return Status;
  }

  // Load the Original FDT tree into the new region
  Err = fdt_open_into (FdtBase, (VOID *)(INTN)UpdatedFdtBase, NewFdtSize);
  if (Err) {
    DEBUG ((DEBUG_ERROR, "fdt_open_into(): %a\n", fdt_strerror (Err)));
    Status = EFI_INVALID_PARAMETER;
    goto Fdt_Exit;
  }

  if (FeaturePcdGet (PcdAndroidBootLoadFile2)) {
    Status = AndroidBootImgInstallLoadFile2 (RamdiskData, RamdiskSize);
    if (EFI_ERROR (Status)) {
      goto Fdt_Exit;
    }
  } else {
    ChosenNode = AndroidBootImgGetChosenNode (UpdatedFdtBase);
    if (!ChosenNode) {
      goto Fdt_Exit;
    }

    Status = AndroidBootImgSetProperty64 (
               UpdatedFdtBase,
               ChosenNode,
               "linux,initrd-start",
               (UINTN)RamdiskData
               );
    if (EFI_ERROR (Status)) {
      goto Fdt_Exit;
    }

    Status = AndroidBootImgSetProperty64 (
               UpdatedFdtBase,
               ChosenNode,
               "linux,initrd-end",
               (UINTN)RamdiskData + RamdiskSize
               );
    if (EFI_ERROR (Status)) {
      goto Fdt_Exit;
    }
  }

  if (mAndroidBootImg->UpdateDtb) {
    Status = mAndroidBootImg->UpdateDtb (UpdatedFdtBase, &NewFdtBase);
    if (EFI_ERROR (Status)) {
      goto Fdt_Exit;
    }
  } else {
    NewFdtBase = UpdatedFdtBase;
  }

  Status = gBS->InstallConfigurationTable (
                  &gFdtTableGuid,
                  (VOID *)(UINTN)NewFdtBase
                  );

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

Fdt_Exit:
  gBS->FreePages (UpdatedFdtBase, EFI_SIZE_TO_PAGES (NewFdtSize));
  return Status;
}

EFI_STATUS
AndroidBootImgBoot (
  IN VOID   *Buffer,
  IN UINTN  BufferSize
  )
{
  EFI_STATUS                 Status;
  VOID                       *Kernel;
  UINTN                      KernelSize;
  MEMORY_DEVICE_PATH         KernelDevicePath;
  EFI_HANDLE                 ImageHandle;
  VOID                       *NewKernelArg;
  EFI_LOADED_IMAGE_PROTOCOL  *ImageInfo;
  VOID                       *RamdiskData;
  UINTN                      RamdiskSize;
  IN  VOID                   *FdtBase;

  if ((Buffer == NULL) || (BufferSize == 0)) {
    return EFI_INVALID_PARAMETER;
  }

  NewKernelArg = NULL;
  ImageHandle  = NULL;

  Status = gBS->LocateProtocol (
                  &gAndroidBootImgProtocolGuid,
                  NULL,
                  (VOID **)&mAndroidBootImg
                  );
  if (EFI_ERROR (Status)) {
    goto Exit;
  }

  Status = AndroidBootImgGetKernelInfo (
             Buffer,
             &Kernel,
             &KernelSize
             );
  if (EFI_ERROR (Status)) {
    goto Exit;
  }

  Status = AndroidBootImgUpdateArgs (Buffer, &NewKernelArg);
  if (EFI_ERROR (Status)) {
    goto Exit;
  }

  Status = AndroidBootImgGetRamdiskInfo (
             Buffer,
             &RamdiskData,
             &RamdiskSize
             );
  if (EFI_ERROR (Status)) {
    goto Exit;
  }

  if (AndroidBootImgAcpiSupported ()) {
    Status = AndroidBootImgInstallLoadFile2 (RamdiskData, RamdiskSize);
    if (EFI_ERROR (Status)) {
      goto Exit;
    }
  } else {
    Status = AndroidBootImgLocateFdt (Buffer, &FdtBase);
    if (EFI_ERROR (Status)) {
      goto Exit;
    }

    Status = AndroidBootImgUpdateFdt (Buffer, FdtBase, RamdiskData, RamdiskSize);
    if (EFI_ERROR (Status)) {
      goto Exit;
    }
  }

  KernelDevicePath = mMemoryDevicePathTemplate;

  KernelDevicePath.Node1.StartingAddress = (EFI_PHYSICAL_ADDRESS)(UINTN)Kernel;
  KernelDevicePath.Node1.EndingAddress   = (EFI_PHYSICAL_ADDRESS)(UINTN)Kernel
                                           + KernelSize;

  Status = gBS->LoadImage (
                  TRUE,
                  gImageHandle,
                  (EFI_DEVICE_PATH *)&KernelDevicePath,
                  (VOID *)(UINTN)Kernel,
                  KernelSize,
                  &ImageHandle
                  );
  if (EFI_ERROR (Status)) {
    goto Exit;
  }

  // Set kernel arguments
  Status = gBS->HandleProtocol (
                  ImageHandle,
                  &gEfiLoadedImageProtocolGuid,
                  (VOID **)&ImageInfo
                  );
  if (EFI_ERROR (Status)) {
    goto Exit;
  }

  ImageInfo->LoadOptions     = NewKernelArg;
  ImageInfo->LoadOptionsSize = StrLen (NewKernelArg) * sizeof (CHAR16);

  // Before calling the image, enable the Watchdog Timer for  the 5 Minute period
  gBS->SetWatchdogTimer (5 * 60, 0x10000, 0, NULL);
  // Start the image
  Status = gBS->StartImage (ImageHandle, NULL, NULL);
  // Clear the Watchdog Timer if the image returns
  gBS->SetWatchdogTimer (0, 0x10000, 0, NULL);

Exit:
  // Unload image as it will not be used anymore
  if (ImageHandle != NULL) {
    gBS->UnloadImage (ImageHandle);
    ImageHandle = NULL;
  }

  if (EFI_ERROR (Status)) {
    if (NewKernelArg != NULL) {
      FreePool (NewKernelArg);
      NewKernelArg = NULL;
    }
  }

  AndroidBootImgUninstallLoadFile2 ();
  return Status;
}