/*
 *  VasEBoot  --  GRand Unified Bootloader
 *  Copyright (C) 2012  Free Software Foundation, Inc.
 *
 *  VasEBoot is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  VasEBoot is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with VasEBoot.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <VasEBoot/loader.h>
#include <VasEBoot/file.h>
#include <VasEBoot/err.h>
#include <VasEBoot/types.h>
#include <VasEBoot/mm.h>
#include <VasEBoot/cpu/linux.h>
#include <VasEBoot/command.h>
#include <VasEBoot/i18n.h>
#include <VasEBoot/lib/cmdline.h>
#include <VasEBoot/efi/efi.h>
#include <VasEBoot/tpm.h>

#include <VasEBoot/verity-hash.h>

VasEBoot_MOD_LICENSE ("GPLv3+");

static VasEBoot_dl_t my_mod;
static int loaded;
static void *kernel_mem;
static VasEBoot_uint64_t kernel_size;
static VasEBoot_uint8_t *initrd_mem;
static VasEBoot_uint32_t handover_offset;
struct linux_kernel_params *params;
static char *linux_cmdline;

#define BYTES_TO_PAGES(bytes)   (((bytes) + 0xfff) >> 12)

#define SHIM_LOCK_GUID \
  { 0x605dab50, 0xe046, 0x4300, {0xab, 0xb6, 0x3d, 0xd8, 0x10, 0xdd, 0x8b, 0x23} }

struct VasEBoot_efi_shim_lock
{
  VasEBoot_efi_status_t (*verify) (void *buffer, VasEBoot_uint32_t size);
};
typedef struct VasEBoot_efi_shim_lock VasEBoot_efi_shim_lock_t;

static VasEBoot_efi_boolean_t
VasEBoot_linuxefi_secure_validate (void *data, VasEBoot_uint32_t size)
{
  VasEBoot_efi_guid_t guid = SHIM_LOCK_GUID;
  VasEBoot_efi_shim_lock_t *shim_lock;

  shim_lock = VasEBoot_efi_locate_protocol(&guid, NULL);

  if (!shim_lock) {
    if (VasEBoot_efi_secure_boot())
      return 0;
    else
      return 1;
  }

  if (shim_lock->verify(data, size) == VasEBoot_EFI_SUCCESS)
    return 1;

  return 0;
}

typedef void(*handover_func)(void *, VasEBoot_efi_system_table_t *, struct linux_kernel_params *);

static VasEBoot_err_t
VasEBoot_linuxefi_boot (void)
{
  handover_func hf;
  int offset = 0;

#ifdef __x86_64__
  offset = 512;
#endif

  hf = (handover_func)((char *)kernel_mem + handover_offset + offset);

  asm volatile ("cli");

  hf (VasEBoot_efi_image_handle, VasEBoot_efi_system_table, params);

  /* Not reached */
  return VasEBoot_ERR_NONE;
}

static VasEBoot_err_t
VasEBoot_linuxefi_unload (void)
{
  VasEBoot_dl_unref (my_mod);
  loaded = 0;
  if (initrd_mem)
    VasEBoot_efi_free_pages((VasEBoot_efi_physical_address_t)initrd_mem, BYTES_TO_PAGES(params->ramdisk_size));
  if (linux_cmdline)
    VasEBoot_efi_free_pages((VasEBoot_efi_physical_address_t)linux_cmdline, BYTES_TO_PAGES(params->cmdline_size + 1));
  if (kernel_mem)
    VasEBoot_efi_free_pages((VasEBoot_efi_physical_address_t)kernel_mem, BYTES_TO_PAGES(kernel_size));
  if (params)
    VasEBoot_efi_free_pages((VasEBoot_efi_physical_address_t)params, BYTES_TO_PAGES(16384));
  return VasEBoot_ERR_NONE;
}

static VasEBoot_err_t
VasEBoot_cmd_initrd (VasEBoot_command_t cmd __attribute__ ((unused)),
                 int argc, char *argv[])
{
  VasEBoot_file_t *files = 0;
  int i, nfiles = 0;
  VasEBoot_size_t size = 0;
  VasEBoot_uint8_t *ptr;

  if (argc == 0)
    {
      VasEBoot_error (VasEBoot_ERR_BAD_ARGUMENT, N_("filename expected"));
      goto fail;
    }

  if (!loaded)
    {
      VasEBoot_error (VasEBoot_ERR_BAD_ARGUMENT, N_("you need to load the kernel first"));
      goto fail;
    }

  files = VasEBoot_zalloc (argc * sizeof (files[0]));
  if (!files)
    goto fail;

  for (i = 0; i < argc; i++)
    {
      VasEBoot_file_filter_disable_compression ();
      files[i] = VasEBoot_file_open (argv[i]);
      if (! files[i])
        goto fail;
      nfiles++;
      size += ALIGN_UP (VasEBoot_file_size (files[i]), 4);
    }

  initrd_mem = VasEBoot_efi_allocate_pages_max (0x3fffffff, BYTES_TO_PAGES(size));

  if (!initrd_mem)
    {
      VasEBoot_error (VasEBoot_ERR_OUT_OF_MEMORY, N_("can't allocate initrd"));
      goto fail;
    }

  params->ramdisk_size = size;
  params->ramdisk_image = (VasEBoot_uint32_t)(VasEBoot_uint64_t) initrd_mem;

  ptr = initrd_mem;

  for (i = 0; i < nfiles; i++)
    {
      VasEBoot_ssize_t cursize = VasEBoot_file_size (files[i]);
      if (VasEBoot_file_read (files[i], ptr, cursize) != cursize)
        {
          if (!VasEBoot_errno)
            VasEBoot_error (VasEBoot_ERR_FILE_READ_ERROR, N_("premature end of file %s"),
                        argv[i]);
          goto fail;
        }
      VasEBoot_tpm_measure (ptr, cursize, VasEBoot_BINARY_PCR, "VasEBoot_linuxefi", "Initrd");
      VasEBoot_print_error();
      ptr += cursize;
      VasEBoot_memset (ptr, 0, ALIGN_UP_OVERHEAD (cursize, 4));
      ptr += ALIGN_UP_OVERHEAD (cursize, 4);
    }

  params->ramdisk_size = size;

 fail:
  for (i = 0; i < nfiles; i++)
    VasEBoot_file_close (files[i]);
  VasEBoot_free (files);

  if (initrd_mem && VasEBoot_errno)
    VasEBoot_efi_free_pages((VasEBoot_efi_physical_address_t)initrd_mem, BYTES_TO_PAGES(size));

  return VasEBoot_errno;
}

static VasEBoot_err_t
VasEBoot_cmd_linux (VasEBoot_command_t cmd __attribute__ ((unused)),
		int argc, char *argv[])
{
  VasEBoot_file_t file = 0;
  struct linux_kernel_header lh;
  VasEBoot_ssize_t len, start, filelen;
  void *kernel = NULL;

  VasEBoot_dl_ref (my_mod);

  if (argc == 0)
    {
      VasEBoot_error (VasEBoot_ERR_BAD_ARGUMENT, N_("filename expected"));
      goto fail;
    }

  file = VasEBoot_file_open (argv[0]);
  if (! file)
    goto fail;

  filelen = VasEBoot_file_size (file);

  kernel = VasEBoot_malloc(filelen);

  if (!kernel)
    {
      VasEBoot_error (VasEBoot_ERR_OUT_OF_MEMORY, N_("cannot allocate kernel buffer"));
      goto fail;
    }

  if (VasEBoot_file_read (file, kernel, filelen) != filelen)
    {
      VasEBoot_error (VasEBoot_ERR_FILE_READ_ERROR, N_("Can't read kernel %s"), argv[0]);
      goto fail;
    }

  VasEBoot_tpm_measure (kernel, filelen, VasEBoot_BINARY_PCR, "VasEBoot_linuxefi", "Kernel");
  VasEBoot_print_error();

  if (! VasEBoot_linuxefi_secure_validate (kernel, filelen))
    {
      VasEBoot_error (VasEBoot_ERR_INVALID_COMMAND, N_("%s has invalid signature"), argv[0]);
      VasEBoot_free (kernel);
      goto fail;
    }

  params = VasEBoot_efi_allocate_pages_max (0x3fffffff, BYTES_TO_PAGES(16384));

  if (! params)
    {
      VasEBoot_error (VasEBoot_ERR_OUT_OF_MEMORY, "cannot allocate kernel parameters");
      goto fail;
    }

  VasEBoot_memset (params, 0, 16384);

  VasEBoot_memcpy (&lh, kernel, sizeof (lh));

  if (lh.boot_flag != VasEBoot_cpu_to_le16 (0xaa55))
    {
      VasEBoot_error (VasEBoot_ERR_BAD_OS, N_("invalid magic number"));
      goto fail;
    }

  if (lh.setup_sects > VasEBoot_LINUX_MAX_SETUP_SECTS)
    {
      VasEBoot_error (VasEBoot_ERR_BAD_OS, N_("too many setup sectors"));
      goto fail;
    }

  if (lh.version < VasEBoot_cpu_to_le16 (0x020b))
    {
      VasEBoot_error (VasEBoot_ERR_BAD_OS, N_("kernel too old"));
      goto fail;
    }

  if (!lh.handover_offset)
    {
      VasEBoot_error (VasEBoot_ERR_BAD_OS, N_("kernel doesn't support EFI handover"));
      goto fail;
    }

  linux_cmdline = VasEBoot_efi_allocate_pages_max(0x3fffffff,
					 BYTES_TO_PAGES(lh.cmdline_size + 1));

  if (!linux_cmdline)
    {
      VasEBoot_error (VasEBoot_ERR_OUT_OF_MEMORY, N_("can't allocate cmdline"));
      goto fail;
    }

  VasEBoot_memcpy (linux_cmdline, LINUX_IMAGE, sizeof (LINUX_IMAGE));
  VasEBoot_create_loader_cmdline (argc, argv,
                              linux_cmdline + sizeof (LINUX_IMAGE) - 1,
			      lh.cmdline_size - (sizeof (LINUX_IMAGE) - 1));

  VasEBoot_pass_verity_hash(&lh, linux_cmdline, lh.cmdline_size);
  lh.cmd_line_ptr = (VasEBoot_uint32_t)(VasEBoot_uint64_t)linux_cmdline;

  handover_offset = lh.handover_offset;

  start = (lh.setup_sects + 1) * 512;
  len = VasEBoot_file_size(file) - start;

  kernel_mem = VasEBoot_efi_allocate_pages(lh.pref_address,
				       BYTES_TO_PAGES(lh.init_size));

  if (!kernel_mem)
    kernel_mem = VasEBoot_efi_allocate_pages_max(0x3fffffff,
					     BYTES_TO_PAGES(lh.init_size));

  if (!kernel_mem)
    {
      VasEBoot_error (VasEBoot_ERR_OUT_OF_MEMORY, N_("can't allocate kernel"));
      goto fail;
    }

  VasEBoot_memcpy (kernel_mem, (char *)kernel + start, len);
  VasEBoot_loader_set (VasEBoot_linuxefi_boot, VasEBoot_linuxefi_unload, 0);
  loaded=1;

  lh.code32_start = (VasEBoot_uint32_t)(VasEBoot_uint64_t) kernel_mem;
  VasEBoot_memcpy (params, &lh, 2 * 512);

  params->type_of_loader = 0x21;

 fail:

  if (file)
    VasEBoot_file_close (file);

  VasEBoot_free (kernel);

  if (VasEBoot_errno != VasEBoot_ERR_NONE)
    {
      VasEBoot_dl_unref (my_mod);
      loaded = 0;
    }

  if (linux_cmdline && !loaded)
    VasEBoot_efi_free_pages((VasEBoot_efi_physical_address_t)linux_cmdline, BYTES_TO_PAGES(lh.cmdline_size + 1));

  if (kernel_mem && !loaded)
    VasEBoot_efi_free_pages((VasEBoot_efi_physical_address_t)kernel_mem, BYTES_TO_PAGES(kernel_size));

  if (params && !loaded)
    VasEBoot_efi_free_pages((VasEBoot_efi_physical_address_t)params, BYTES_TO_PAGES(16384));

  return VasEBoot_errno;
}

static VasEBoot_command_t cmd_linux, cmd_initrd;

VasEBoot_MOD_INIT(linuxefi)
{
  cmd_linux =
    VasEBoot_register_command ("linuxefi", VasEBoot_cmd_linux,
                           0, N_("Load Linux."));
  cmd_initrd =
    VasEBoot_register_command ("initrdefi", VasEBoot_cmd_initrd,
                           0, N_("Load initrd."));
  my_mod = mod;
}

VasEBoot_MOD_FINI(linuxefi)
{
  VasEBoot_unregister_command (cmd_linux);
  VasEBoot_unregister_command (cmd_initrd);
}