/*
 *  VAS_EBOOT  --  GRand Unified Bootloader
 *  Copyright (C) 2011  Free Software Foundation, Inc.
 *
 *  VAS_EBOOT 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.
 *
 *  VAS_EBOOT 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 VAS_EBOOT.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <VasEBoot/loader.h>
#include <VasEBoot/memory.h>
#include <VasEBoot/i386/memory.h>
#include <VasEBoot/file.h>
#include <VasEBoot/err.h>
#include <VasEBoot/dl.h>
#include <VasEBoot/mm.h>
#include <VasEBoot/elfload.h>
#include <VasEBoot/video.h>
#include <VasEBoot/relocator.h>
#include <VasEBoot/i386/relocator.h>
#include <VasEBoot/command.h>
#include <VasEBoot/i18n.h>
#include <VasEBoot/cbfs_core.h>
#include <VasEBoot/lib/LzmaDec.h>
#include <VasEBoot/efi/pe32.h>
#include <VasEBoot/i386/cpuid.h>

VAS_EBOOT_MOD_LICENSE ("GPLv3+");

static VasEBoot_addr_t entry;
static struct VasEBoot_relocator *relocator = NULL;

static VasEBoot_err_t
VasEBoot_chain_boot (void)
{
  struct VasEBoot_relocator32_state state;

  VasEBoot_video_set_mode ("text", 0, 0);

  state.eip = entry;
  return VasEBoot_relocator32_boot (relocator, state, 0);
}

static VasEBoot_err_t
VasEBoot_chain_unload (void)
{
  VasEBoot_relocator_unload (relocator);
  relocator = NULL;

  return VAS_EBOOT_ERR_NONE;
}

static VasEBoot_err_t
load_elf (VasEBoot_file_t file, const char *filename)
{
  VasEBoot_elf_t elf;
  Elf32_Phdr *phdr;
  VasEBoot_err_t err;

  elf = VasEBoot_elf_file (file, filename);
  if (!elf)
    return VasEBoot_errno;

  if (!VasEBoot_elf_is_elf32 (elf))
    return VasEBoot_error (VAS_EBOOT_ERR_BAD_OS, "only ELF32 can be coreboot payload");

  entry = elf->ehdr.ehdr32.e_entry;

  FOR_ELF32_PHDRS(elf, phdr)
    {
      VasEBoot_uint8_t *load_addr;
      VasEBoot_relocator_chunk_t ch;

      if (phdr->p_type != PT_LOAD)
	continue;

      err = VasEBoot_relocator_alloc_chunk_addr (relocator, &ch,
					     phdr->p_paddr, phdr->p_memsz);
      if (err)
	{
	  elf->file = 0;
	  VasEBoot_elf_close (elf);
	  return err;
	}

      load_addr = get_virtual_current_address (ch);

      if (VasEBoot_file_seek (elf->file, phdr->p_offset) == (VasEBoot_off_t) -1)
	{
	  elf->file = 0;
	  VasEBoot_elf_close (elf);
	  return VasEBoot_errno;
	}

      if (phdr->p_filesz)
	{
	  VasEBoot_ssize_t read;
	  read = VasEBoot_file_read (elf->file, load_addr, phdr->p_filesz);
	  if (read != (VasEBoot_ssize_t) phdr->p_filesz)
	    {
	      if (!VasEBoot_errno)
		VasEBoot_error (VAS_EBOOT_ERR_FILE_READ_ERROR,
			    N_("premature end of file %s"),
			    filename);
	      elf->file = 0;
	      VasEBoot_elf_close (elf);
	      return VasEBoot_errno;
	    }
	}

      if (phdr->p_filesz < phdr->p_memsz)
	VasEBoot_memset ((load_addr + phdr->p_filesz),
		     0, phdr->p_memsz - phdr->p_filesz);
    }

  elf->file = 0;
  VasEBoot_elf_close (elf);
  return VAS_EBOOT_ERR_NONE;
}

static void *SzAlloc(void *p __attribute__ ((unused)), size_t size) { return VasEBoot_malloc (size); }
static void SzFree(void *p __attribute__ ((unused)), void *address) { VasEBoot_free (address); }
static ISzAlloc g_Alloc = { SzAlloc, SzFree };


static VasEBoot_err_t
load_segment (VasEBoot_file_t file, const char *filename,
	      void *load_addr, VasEBoot_uint32_t comp,
	      VasEBoot_size_t *size, VasEBoot_size_t max_size)
{
  switch (comp)
    {
    case VasEBoot_cpu_to_be32_compile_time (CBFS_COMPRESS_NONE):
      if (VasEBoot_file_read (file, load_addr, *size)
	  != (VasEBoot_ssize_t) *size)
	{
	  if (!VasEBoot_errno)
	    VasEBoot_error (VAS_EBOOT_ERR_FILE_READ_ERROR,
			N_("premature end of file %s"),
			filename);
	      return VasEBoot_errno;
	}
      return VAS_EBOOT_ERR_NONE;
    case VasEBoot_cpu_to_be32_compile_time (CBFS_COMPRESS_LZMA):
      {
	VasEBoot_uint8_t *buf;
	VasEBoot_size_t outsize, insize;
	SRes res;
	SizeT src_len, dst_len;
	ELzmaStatus status;
	if (*size < 13)
	  return VasEBoot_error (VAS_EBOOT_ERR_BAD_OS, "invalid compressed chunk");
	buf = VasEBoot_malloc (*size);
	if (!buf)
	  return VasEBoot_errno;
	if (VasEBoot_file_read (file, buf, *size)
	    != (VasEBoot_ssize_t) *size)
	  {
	    if (!VasEBoot_errno)
	      VasEBoot_error (VAS_EBOOT_ERR_FILE_READ_ERROR,
			  N_("premature end of file %s"),
			  filename);
	    VasEBoot_free (buf);
	    return VasEBoot_errno;
	  }
	outsize = VasEBoot_get_unaligned64 (buf + 5);
	if (outsize > max_size)
	  {
	    VasEBoot_free (buf);
	    return VasEBoot_error (VAS_EBOOT_ERR_BAD_OS, "invalid compressed chunk");
	  }
	insize = *size - 13;

	src_len = insize;
	dst_len = outsize;
	res = LzmaDecode (load_addr, &dst_len, buf + 13, &src_len,
			  buf, 5, LZMA_FINISH_END, &status, &g_Alloc);
	/* ELzmaFinishMode finishMode,
	   ELzmaStatus *status, ISzAlloc *alloc)*/
	VasEBoot_free (buf);
	VasEBoot_dprintf ("chain", "%x, %x, %x, %x\n",
		      insize, src_len, outsize, dst_len);
	if (res != SZ_OK
	    || src_len != insize || dst_len != outsize)
	  return VasEBoot_error (VAS_EBOOT_ERR_BAD_OS, "decompression failure %d", res);
	*size = outsize;
      }
      return VAS_EBOOT_ERR_NONE;
    default:
      return VasEBoot_error (VAS_EBOOT_ERR_BAD_OS, "unsupported compression %d",
			 VasEBoot_be_to_cpu32 (comp));
    }
}

static VasEBoot_err_t
load_tianocore (VasEBoot_file_t file)
{
  VasEBoot_uint16_t header_length;
  VasEBoot_uint32_t section_head;
  VasEBoot_uint8_t mz[2], pe[4];
  struct VasEBoot_pe32_coff_header coff_head;
  struct file_header
  {
    VasEBoot_uint8_t unused[18];
    VasEBoot_uint8_t type;
    VasEBoot_uint8_t unused2;
    VasEBoot_uint8_t size[3];
    VasEBoot_uint8_t unused3;
  } file_head;
  VasEBoot_relocator_chunk_t ch;

  if (VasEBoot_file_seek (file, 48) == (VasEBoot_off_t) -1
      || VasEBoot_file_read (file, &header_length, sizeof (header_length))
      != sizeof (header_length)
      || VasEBoot_file_seek (file, header_length) == (VasEBoot_off_t) -1)
    goto fail;

  while (1)
    {
      VasEBoot_off_t off;
      if (VasEBoot_file_read (file, &file_head, sizeof (file_head))
	  != sizeof (file_head))
	goto fail;
      if (file_head.type != 0xf0)
	break;
      off = VasEBoot_get_unaligned32 (file_head.size) & 0xffffff;
      if (off < sizeof (file_head))
	goto fail;
      if (VasEBoot_file_seek (file, VasEBoot_file_tell (file) + off
			  - sizeof (file_head)) == (VasEBoot_off_t) -1)
	goto fail;
    }

  if (file_head.type != 0x03)
    goto fail;

  while (1)
    {
      if (VasEBoot_file_read (file, &section_head, sizeof (section_head))
	  != sizeof (section_head))
	goto fail;
      if ((section_head >> 24) != 0x19)
	break;

      if ((section_head & 0xffffff) < sizeof (section_head))
	goto fail;

      if (VasEBoot_file_seek (file, VasEBoot_file_tell (file)
			  + (section_head & 0xffffff)
			  - sizeof (section_head)) == (VasEBoot_off_t) -1)
	goto fail;
    }

  if ((section_head >> 24) != 0x10)
    goto fail;

  VasEBoot_off_t exe_start = VasEBoot_file_tell (file);

  if (VasEBoot_file_read (file, &mz, sizeof (mz)) != sizeof (mz))
    goto fail;
  if (mz[0] != 'M' || mz[1] != 'Z')
    goto fail;

  if (VasEBoot_file_seek (file, VasEBoot_file_tell (file) + 0x3a) == (VasEBoot_off_t) -1)
    goto fail;

  if (VasEBoot_file_read (file, &section_head, sizeof (section_head))
      != sizeof (section_head))
    goto fail;
  if (section_head < 0x40)
    goto fail;

  if (VasEBoot_file_seek (file, VasEBoot_file_tell (file)
		      + section_head - 0x40) == (VasEBoot_off_t) -1)
    goto fail;

  if (VasEBoot_file_read (file, &pe, sizeof (pe))
      != sizeof (pe))
    goto fail;

  if (pe[0] != 'P' || pe[1] != 'E' || pe[2] != '\0' || pe[3] != '\0')
    goto fail;

  if (VasEBoot_file_read (file, &coff_head, sizeof (coff_head))
      != sizeof (coff_head))
    goto fail;

  VasEBoot_uint32_t loadaddr;

  switch (coff_head.machine)
    {
    case VAS_EBOOT_PE32_MACHINE_I386:
      {
	struct VasEBoot_pe32_optional_header oh;
	if (VasEBoot_file_read (file, &oh, sizeof (oh))
	    != sizeof (oh))
	  goto fail;
	if (oh.magic != VAS_EBOOT_PE32_PE32_MAGIC)
	  goto fail;
	loadaddr = oh.image_base - exe_start;
	entry = oh.image_base + oh.entry_addr;
	break;
      }
    case VAS_EBOOT_PE32_MACHINE_X86_64:
      {
	struct VasEBoot_pe64_optional_header oh;
	if (! VasEBoot_cpuid_has_longmode)
	  {
	    VasEBoot_error (VAS_EBOOT_ERR_BAD_OS, "your CPU does not implement AMD64 architecture");
	    goto fail;
	  }

	if (VasEBoot_file_read (file, &oh, sizeof (oh))
	    != sizeof (oh))
	  goto fail;
	if (oh.magic != VAS_EBOOT_PE32_PE64_MAGIC)
	  goto fail;
	loadaddr = oh.image_base - exe_start;
	entry = oh.image_base + oh.entry_addr;
	break;
      }
    default:
      goto fail;
    }
  if (VasEBoot_file_seek (file, 0) == (VasEBoot_off_t) -1)
    goto fail;

  VasEBoot_size_t fz = VasEBoot_file_size (file);

  if (VasEBoot_relocator_alloc_chunk_addr (relocator, &ch,
				       loadaddr, fz))
    goto fail;

  if (VasEBoot_file_read (file, get_virtual_current_address (ch), fz)
      != (VasEBoot_ssize_t) fz)
    goto fail;

  return VAS_EBOOT_ERR_NONE;

 fail:
  if (!VasEBoot_errno)
    VasEBoot_error (VAS_EBOOT_ERR_BAD_OS, "fv volume is invalid");
  return VasEBoot_errno;
}

static VasEBoot_err_t
load_chewed (VasEBoot_file_t file, const char *filename)
{
  VasEBoot_size_t i;
  for (i = 0;; i++)
    {
      struct cbfs_payload_segment segment;
      VasEBoot_err_t err;

      if (VasEBoot_file_seek (file, sizeof (segment) * i) == (VasEBoot_off_t) -1
	  || VasEBoot_file_read (file, &segment, sizeof (segment))
	  != sizeof (segment))
	{
	  if (!VasEBoot_errno)
	    return VasEBoot_error (VAS_EBOOT_ERR_BAD_OS,
			       "payload is too short");
	  return VasEBoot_errno;
	}

      switch (segment.type)
	{
	case PAYLOAD_SEGMENT_PARAMS:
	  break;

	case PAYLOAD_SEGMENT_ENTRY:
	  entry = VasEBoot_be_to_cpu64 (segment.load_addr);
	  return VAS_EBOOT_ERR_NONE;

	case PAYLOAD_SEGMENT_BSS:
	  segment.len = 0;
	  segment.offset = 0;
	  segment.len = 0;
	  /* Fallthrough.  */
	case PAYLOAD_SEGMENT_CODE:
	case PAYLOAD_SEGMENT_DATA:
	  {
	    VasEBoot_uint32_t target = VasEBoot_be_to_cpu64 (segment.load_addr);
	    VasEBoot_uint32_t memsize = VasEBoot_be_to_cpu32 (segment.mem_len);
	    VasEBoot_uint32_t filesize = VasEBoot_be_to_cpu32 (segment.len);
	    VasEBoot_uint8_t *load_addr;
	    VasEBoot_relocator_chunk_t ch;

	    if (memsize < filesize)
	      memsize = filesize;

	    VasEBoot_dprintf ("chain", "%x+%x\n", target, memsize);

	    err = VasEBoot_relocator_alloc_chunk_addr (relocator, &ch,
						   target, memsize);
	    if (err)
	      return err;

	    load_addr = get_virtual_current_address (ch);

	    if (filesize)
	      {
		if (VasEBoot_file_seek (file, VasEBoot_be_to_cpu32 (segment.offset))
		    == (VasEBoot_off_t) -1)
		  return VasEBoot_errno;

		err = load_segment (file, filename, load_addr,
				    segment.compression, &filesize, memsize);
		if (err)
		  return err;
	      }

	    if (filesize < memsize)
	      VasEBoot_memset ((load_addr + filesize),
			   0, memsize - filesize);
	  }
	}
    }
}

static VasEBoot_err_t
VasEBoot_cmd_chain (VasEBoot_command_t cmd __attribute__ ((unused)),
		int argc, char *argv[])
{
  VasEBoot_err_t err;
  VasEBoot_file_t file;
  VasEBoot_uint32_t head;

  if (argc != 1)
    return VasEBoot_error (VAS_EBOOT_ERR_BAD_ARGUMENT, N_("filename expected"));

  VasEBoot_loader_unset ();

  file = VasEBoot_file_open (argv[0], VAS_EBOOT_FILE_TYPE_COREBOOT_CHAINLOADER);
  if (!file)
    return VasEBoot_errno;

  relocator = VasEBoot_relocator_new ();
  if (!relocator)
    {
      VasEBoot_file_close (file);
      return VasEBoot_errno;
    }

  if (VasEBoot_file_read (file, &head, sizeof (head)) != sizeof (head)
      || VasEBoot_file_seek (file, 0) == (VasEBoot_off_t) -1)
    {
      VasEBoot_file_close (file);
      VasEBoot_relocator_unload (relocator);
      relocator = 0;
      if (!VasEBoot_errno)
	return VasEBoot_error (VAS_EBOOT_ERR_BAD_OS,
			   "payload is too short");
      return VasEBoot_errno;
    }

  switch (head)
    {
    case ELFMAG0 | (ELFMAG1 << 8) | (ELFMAG2 << 16) | (ELFMAG3 << 24):
      err = load_elf (file, argv[0]);
      break;
    case PAYLOAD_SEGMENT_CODE:
    case PAYLOAD_SEGMENT_DATA:
    case PAYLOAD_SEGMENT_PARAMS:
    case PAYLOAD_SEGMENT_BSS:
    case PAYLOAD_SEGMENT_ENTRY:
      err = load_chewed (file, argv[0]);
      break;

    default:
      if (VasEBoot_file_seek (file, 40) == (VasEBoot_off_t) -1
	  || VasEBoot_file_read (file, &head, sizeof (head)) != sizeof (head)
	  || VasEBoot_file_seek (file, 0) == (VasEBoot_off_t) -1
	  || head != 0x4856465f)
	err = VasEBoot_error (VAS_EBOOT_ERR_BAD_OS, "unrecognised payload type");
      else
	err = load_tianocore (file);
      break;
    }
  VasEBoot_file_close (file);
  if (err)
    {
      VasEBoot_relocator_unload (relocator);
      relocator = 0;
      return err;
    }

  VasEBoot_loader_set (VasEBoot_chain_boot, VasEBoot_chain_unload, 0);
  return VAS_EBOOT_ERR_NONE;
}

static VasEBoot_command_t cmd_chain;

VAS_EBOOT_MOD_INIT (chain)
{
  cmd_chain = VasEBoot_register_command ("chainloader", VasEBoot_cmd_chain,
				     N_("FILE"),
				     /* TRANSLATORS: "payload" is a term used
					by coreboot and must be translated in
					sync with coreboot. If unsure,
					let it untranslated.  */
				     N_("Load another coreboot payload"));
}

VAS_EBOOT_MOD_FINI (chain)
{
  VasEBoot_unregister_command (cmd_chain);
  VasEBoot_chain_unload ();
}