/* Finalize operations on the assembler context, free all resources.
   Copyright (C) 2002, 2003, 2005 Red Hat, Inc.
   Written by Ulrich Drepper <drepper@redhat.com>, 2002.

   This program is Open Source software; you can redistribute it and/or
   modify it under the terms of the Open Software License version 1.0 as
   published by the Open Source Initiative.

   You should have received a copy of the Open Software License along
   with this program; if not, you may obtain a copy of the Open Software
   License version 1.0 from http://www.opensource.org/licenses/osl.php or
   by writing the Open Source Initiative c/o Lawrence Rosen, Esq.,
   3001 King Ranch Road, Ukiah, CA 95482.   */

#ifdef HAVE_CONFIG_H
# include <config.h>
#endif

#include <assert.h>
#include <error.h>
#include <libintl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/stat.h>

#include <libasmP.h>
#include <libelf.h>
#include <system.h>


static int
text_end (AsmCtx_t *ctx __attribute__ ((unused)))
{
  // XXX Does anything have to be done?
  return 0;
}


static int
binary_end (AsmCtx_t *ctx)
{
  void *symtab = NULL;
  struct Ebl_Strent *symscn_strent = NULL;
  struct Ebl_Strent *strscn_strent = NULL;
  struct Ebl_Strent *xndxscn_strent = NULL;
  Elf_Scn *shstrscn;
  struct Ebl_Strent *shstrscn_strent;
  size_t shstrscnndx;
  size_t symscnndx = 0;
  size_t strscnndx = 0;
  size_t xndxscnndx = 0;
  Elf_Data *data;
  Elf_Data *shstrtabdata;
  Elf_Data *strtabdata = NULL;
  Elf_Data *xndxdata = NULL;
  GElf_Shdr shdr_mem;
  GElf_Shdr *shdr;
  GElf_Ehdr ehdr_mem;
  GElf_Ehdr *ehdr;
  AsmScn_t *asmscn;
  int result = 0;

  /* Iterate over the created sections and compute the offsets of the
     various subsections and fill in the content.  */
  for (asmscn = ctx->section_list; asmscn != NULL; asmscn = asmscn->allnext)
    {
#if 0
      Elf_Scn *scn = elf_getscn (ctx->out.elf, asmscn->data.main.scnndx);
#else
      Elf_Scn *scn = asmscn->data.main.scn;
#endif
      off_t offset = 0;
      AsmScn_t *asmsubscn = asmscn;

      do
	{
	  struct AsmData *content = asmsubscn->content;
	  bool first = true;

	  offset = ((offset + asmsubscn->max_align - 1)
		    & ~(asmsubscn->max_align - 1));

	  /* Update the offset for this subsection.  This field now
	     stores the offset of the first by in this subsection.  */
	  asmsubscn->offset = offset;

	  /* Note that the content list is circular.  */
	  if (content != NULL)
	    do
	      {
		Elf_Data *newdata = elf_newdata (scn);

		if (newdata == NULL)
		  {
		    __libasm_seterrno (ASM_E_LIBELF);
		    return -1;
		  }

		newdata->d_buf = content->data;
		newdata->d_type = ELF_T_BYTE;
		newdata->d_size = content->len;
		newdata->d_off = offset;
		newdata->d_align = first ? asmsubscn->max_align : 1;

		offset += content->len;
	      }
	    while ((content = content->next) != asmsubscn->content);
	}
      while ((asmsubscn = asmsubscn->subnext) != NULL);
    }


  /* Create the symbol table if necessary.  */
  if (ctx->nsymbol_tab > 0)
    {
      /* Create the symbol table and string table section names.  */
      symscn_strent = ebl_strtabadd (ctx->section_strtab, ".symtab", 8);
      strscn_strent = ebl_strtabadd (ctx->section_strtab, ".strtab", 8);

      /* Create the symbol string table section.  */
      Elf_Scn *strscn = elf_newscn (ctx->out.elf);
      strtabdata = elf_newdata (strscn);
      shdr = gelf_getshdr (strscn, &shdr_mem);
      if (strtabdata == NULL || shdr == NULL)
	{
	  __libasm_seterrno (ASM_E_LIBELF);
	  return -1;
	}
      strscnndx = elf_ndxscn (strscn);

      ebl_strtabfinalize (ctx->symbol_strtab, strtabdata);

      shdr->sh_type = SHT_STRTAB;
      assert (shdr->sh_entsize == 0);

      (void) gelf_update_shdr (strscn, shdr);

      /* Create the symbol table section.  */
      Elf_Scn *symscn = elf_newscn (ctx->out.elf);
      data = elf_newdata (symscn);
      shdr = gelf_getshdr (symscn, &shdr_mem);
      if (data == NULL || shdr == NULL)
	{
	  __libasm_seterrno (ASM_E_LIBELF);
	  return -1;
	}
      symscnndx = elf_ndxscn (symscn);

      /* We know how many symbols there will be in the symbol table.  */
      data->d_size = gelf_fsize (ctx->out.elf, ELF_T_SYM,
				 ctx->nsymbol_tab + 1, EV_CURRENT);
      symtab = malloc (data->d_size);
      if (symtab == NULL)
	return -1;
      data->d_buf = symtab;
      data->d_type = ELF_T_SYM;
      data->d_off = 0;

      /* Clear the first entry.  */
      GElf_Sym syment;
      memset (&syment, '\0', sizeof (syment));
      (void) gelf_update_sym (data, 0, &syment);

      /* Iterate over the symbol table.  */
      void *runp = NULL;
      int ptr_local = 1;	/* Start with index 1; zero remains unused.  */
      int ptr_nonlocal = ctx->nsymbol_tab;
      uint32_t *xshndx = NULL;
      AsmSym_t *sym;
      while ((sym = asm_symbol_tab_iterate (&ctx->symbol_tab, &runp)) != NULL)
	if (asm_emit_symbol_p (ebl_string (sym->strent)))
	  {
	    assert (ptr_local <= ptr_nonlocal);

	    syment.st_name = ebl_strtaboffset (sym->strent);
	    syment.st_info = GELF_ST_INFO (sym->binding, sym->type);
	    syment.st_other = 0;
	    syment.st_value = sym->scn->offset + sym->offset;
	    syment.st_size = sym->size;

	    /* Add local symbols at the beginning, the other from
	       the end.  */
	    int ptr = sym->binding == STB_LOCAL ? ptr_local++ : ptr_nonlocal--;

	    /* Determine the section index.  We have to handle the
	       overflow correctly.  */
	    Elf_Scn *scn = (sym->scn->subsection_id == 0
			    ? sym->scn->data.main.scn
			    : sym->scn->data.up->data.main.scn);

	    Elf32_Word ndx;
	    if (unlikely (scn == ASM_ABS_SCN))
	      ndx = SHN_ABS;
	    else if (unlikely (scn == ASM_COM_SCN))
	      ndx = SHN_COMMON;
	    else if (unlikely ((ndx = elf_ndxscn (scn)) >= SHN_LORESERVE))
	      {
		if (unlikely (xshndx == NULL))
		  {
		    /* The extended section index section does not yet
		       exist.  */
		    Elf_Scn *xndxscn;

		    xndxscn = elf_newscn (ctx->out.elf);
		    xndxdata = elf_newdata (xndxscn);
		    shdr = gelf_getshdr (xndxscn, &shdr_mem);
		    if (xndxdata == NULL || shdr == NULL)
		      {
			__libasm_seterrno (ASM_E_LIBELF);
			return -1;
		      }
		    xndxscnndx = elf_ndxscn (xndxscn);

		    shdr->sh_type = SHT_SYMTAB_SHNDX;
		    shdr->sh_entsize = sizeof (Elf32_Word);
		    shdr->sh_addralign = sizeof (Elf32_Word);
		    shdr->sh_link = symscnndx;

		    (void) gelf_update_shdr (xndxscn, shdr);

		    xndxscn_strent = ebl_strtabadd (ctx->section_strtab,
						    ".symtab_shndx", 14);

		    /* Note that using 'elf32_fsize' instead of
		       'gelf_fsize' here is correct.  */
		    xndxdata->d_size = elf32_fsize (ELF_T_WORD,
						    ctx->nsymbol_tab + 1,
						    EV_CURRENT);
		    xshndx = xndxdata->d_buf = calloc (1, xndxdata->d_size);
		    if (xshndx == NULL)
		      return -1;
		    /* Using ELF_T_WORD here relies on the fact that the
		       32- and 64-bit types are the same size.  */
		    xndxdata->d_type = ELF_T_WORD;
		    xndxdata->d_off = 0;
		  }

		/* Store the real section index in the extended setion
		   index table.  */
		assert ((size_t) ptr < ctx->nsymbol_tab + 1);
		xshndx[ptr] = ndx;

		/* And signal that this happened.  */
		ndx = SHN_XINDEX;
	      }
	    syment.st_shndx = ndx;

	    /* Remember where we put the symbol.  */
	    sym->symidx = ptr;

	    (void) gelf_update_sym (data, ptr, &syment);
	  }

      assert (ptr_local == ptr_nonlocal + 1);

      shdr->sh_type = SHT_SYMTAB;
      shdr->sh_link = strscnndx;
      shdr->sh_info = ptr_local;
      shdr->sh_entsize = gelf_fsize (ctx->out.elf, ELF_T_SYM, 1, EV_CURRENT);
      shdr->sh_addralign = gelf_fsize (ctx->out.elf, ELF_T_ADDR, 1,
				       EV_CURRENT);

      (void) gelf_update_shdr (symscn, shdr);
    }


  /* Create the section header string table section and fill in the
     references in the section headers.  */
  shstrscn = elf_newscn (ctx->out.elf);
  shstrtabdata = elf_newdata (shstrscn);
  shdr = gelf_getshdr (shstrscn, &shdr_mem);
  if (shstrscn == NULL || shstrtabdata == NULL || shdr == NULL)
    {
      __libasm_seterrno (ASM_E_LIBELF);
      return -1;
    }


  /* Add the name of the section header string table.  */
  shstrscn_strent = ebl_strtabadd (ctx->section_strtab, ".shstrtab", 10);

  ebl_strtabfinalize (ctx->section_strtab, shstrtabdata);

  shdr->sh_type = SHT_STRTAB;
  assert (shdr->sh_entsize == 0);
  shdr->sh_name = ebl_strtaboffset (shstrscn_strent);

  (void) gelf_update_shdr (shstrscn, shdr);


  /* Create the section groups.  */
  if (ctx->groups != NULL)
    {
      AsmScnGrp_t *runp = ctx->groups->next;

      do
	{
	  Elf_Scn *scn;
	  Elf32_Word *grpdata;

	  scn = runp->scn;
	  assert (scn != NULL);
	  shdr = gelf_getshdr (scn, &shdr_mem);
	  assert (shdr != NULL);

	  data = elf_newdata (scn);
	  if (data == NULL)
	    {
	      __libasm_seterrno (ASM_E_LIBELF);
	      return -1;
	    }

	  /* It is correct to use 'elf32_fsize' instead of 'gelf_fsize'
	     here.  */
	  data->d_size = elf32_fsize (ELF_T_WORD, runp->nmembers + 1,
				      EV_CURRENT);
	  grpdata = data->d_buf = malloc (data->d_size);
	  if (grpdata == NULL)
	    return -1;
	  data->d_type = ELF_T_WORD;
	  data->d_off = 0;
	  data->d_align = elf32_fsize (ELF_T_WORD, 1, EV_CURRENT);

	  /* The first word of the section is filled with the flag word.  */
	  *grpdata++ = runp->flags;

	  if (runp->members != NULL)
	    {
	      AsmScn_t *member = runp->members->data.main.next_in_group;

	      do
		{
		  /* Only sections, not subsections, can be registered
		     as member of a group.  The subsections get
		     automatically included.  */
		  assert (member->subsection_id == 0);

		  *grpdata++ = elf_ndxscn (member->data.main.scn);
		}
	      while ((member = member->data.main.next_in_group)
		     != runp->members->data.main.next_in_group);
	    }

	  /* Construct the section header.  */
	  shdr->sh_name = ebl_strtaboffset (runp->strent);
	  shdr->sh_type = SHT_GROUP;
	  shdr->sh_flags = 0;
	  shdr->sh_link = symscnndx;
	  /* If the user did not specify a signature we use the initial
	     empty symbol in the symbol table as the signature.  */
	  shdr->sh_info = (runp->signature != NULL
			   ? runp->signature->symidx : 0);

	  (void) gelf_update_shdr (scn, shdr);
	}
      while ((runp = runp->next) != ctx->groups->next);
    }


  /* Add the name to the symbol section.  */
  if (likely (symscnndx != 0))
    {
      Elf_Scn *scn = elf_getscn (ctx->out.elf, symscnndx);

      shdr = gelf_getshdr (scn, &shdr_mem);

      shdr->sh_name = ebl_strtaboffset (symscn_strent);

      (void) gelf_update_shdr (scn, shdr);


      /* Add the name to the string section.  */
      assert (strscnndx != 0);
      scn = elf_getscn (ctx->out.elf, strscnndx);

      shdr = gelf_getshdr (scn, &shdr_mem);

      shdr->sh_name = ebl_strtaboffset (strscn_strent);

      (void) gelf_update_shdr (scn, shdr);


      /* Add the name to the extended symbol index section.  */
      if (xndxscnndx != 0)
	{
	  scn = elf_getscn (ctx->out.elf, xndxscnndx);

	  shdr = gelf_getshdr (scn, &shdr_mem);

	  shdr->sh_name = ebl_strtaboffset (xndxscn_strent);

	  (void) gelf_update_shdr (scn, shdr);
	}
    }


  /* Iterate over the created sections and fill in the names.  */
  for (asmscn = ctx->section_list; asmscn != NULL; asmscn = asmscn->allnext)
    {
      shdr = gelf_getshdr (asmscn->data.main.scn, &shdr_mem);
      /* This better should not fail.  */
      assert (shdr != NULL);

      shdr->sh_name = ebl_strtaboffset (asmscn->data.main.strent);

      /* We now know the maximum alignment.  */
      shdr->sh_addralign = asmscn->max_align;

      (void) gelf_update_shdr (asmscn->data.main.scn, shdr);
    }

  /* Put the reference to the section header string table in the ELF
     header.  */
  ehdr = gelf_getehdr (ctx->out.elf, &ehdr_mem);
  assert (ehdr != NULL);

  shstrscnndx = elf_ndxscn (shstrscn);
  if (unlikely (shstrscnndx > SHN_HIRESERVE)
      || unlikely (shstrscnndx == SHN_XINDEX))
    {
      /* The index of the section header string sectio is too large.  */
      Elf_Scn *scn = elf_getscn (ctx->out.elf, 0);

      /* Get the header for the zeroth section.  */
      shdr = gelf_getshdr (scn, &shdr_mem);
      /* This better does not fail.  */
      assert (shdr != NULL);

      /* The sh_link field of the zeroth section header contains the value.  */
      shdr->sh_link = shstrscnndx;

      (void) gelf_update_shdr (scn, shdr);

      /* This is the sign for the overflow.  */
      ehdr->e_shstrndx = SHN_XINDEX;
    }
  else
    ehdr->e_shstrndx = elf_ndxscn (shstrscn);

  gelf_update_ehdr (ctx->out.elf, ehdr);

  /* Write out the ELF file.  */
  if (unlikely (elf_update (ctx->out.elf, ELF_C_WRITE_MMAP)) < 0)
    {
      __libasm_seterrno (ASM_E_LIBELF);
      result = -1;
    }

  /* We do not need the section header and symbol string tables anymore.  */
  free (shstrtabdata->d_buf);
  if (strtabdata != NULL)
    free (strtabdata->d_buf);
  /* We might have allocated the extended symbol table index.  */
  if (xndxdata != NULL)
    free (xndxdata->d_buf);

  /* Free section groups memory.  */
  AsmScnGrp_t *scngrp = ctx->groups;
  if (scngrp != NULL)
    do
      free (elf_getdata (scngrp->scn, NULL)->d_buf);
    while ((scngrp = scngrp->next) != ctx->groups);

  /* Finalize the ELF handling.  */
  if (unlikely (elf_end (ctx->out.elf)) != 0)
    {
      __libasm_seterrno (ASM_E_LIBELF);
      result = -1;
    }

  /* Free the temporary resources.  */
  free (symtab);

  return result;
}


int
asm_end (ctx)
     AsmCtx_t *ctx;
{
  int result;

  if (ctx == NULL)
    /* Something went wrong earlier.  */
    return -1;

  result = unlikely (ctx->textp) ? text_end (ctx) : binary_end (ctx);
  if (result != 0)
    return result;

  /* Make the new file globally readable and user/group-writable.  */
  if (fchmod (ctx->fd, S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH) != 0)
    {
      __libasm_seterrno (ASM_E_CANNOT_CHMOD);
      return -1;
    }

  /* Rename output file.  */
  if (rename (ctx->tmp_fname, ctx->fname) != 0)
    {
      __libasm_seterrno (ASM_E_CANNOT_RENAME);
      return -1;
    }

  /* Free the resources.  */
  __libasm_finictx (ctx);

  return 0;
}


static void
free_section (AsmScn_t *scnp)
{
  void *oldp;

  if (scnp->subnext != NULL)
    free_section (scnp->subnext);

  struct AsmData *data = scnp->content;
  if (data != NULL)
    do
      {
	oldp = data;
	data = data->next;
	free (oldp);
      }
    while (oldp != scnp->content);

  free (scnp);
}


void
__libasm_finictx (ctx)
     AsmCtx_t *ctx;
{
  /* Iterate through section table and free individual entries.  */
  AsmScn_t *scn = ctx->section_list;
  while (scn != NULL)
    {
      AsmScn_t *oldp = scn;
      scn = scn->allnext;
      free_section (oldp);
    }

  /* Free the resources of the symbol table.  */
  void *runp = NULL;
  AsmSym_t *sym;
  while ((sym = asm_symbol_tab_iterate (&ctx->symbol_tab, &runp)) != NULL)
    free (sym);
  asm_symbol_tab_free (&ctx->symbol_tab);


  /* Free section groups.  */
  AsmScnGrp_t *scngrp = ctx->groups;
  if (scngrp != NULL)
    do
      {
	AsmScnGrp_t *oldp = scngrp;

	scngrp = scngrp->next;
	free (oldp);
      }
    while (scngrp != ctx->groups);


  if (unlikely (ctx->textp))
    {
      /* Close the stream.  */
      fclose (ctx->out.file);
    }
  else
    {
      /* Close the output file.  */
      /* XXX We should test for errors here but what would we do if we'd
	 find any.  */
      (void) close (ctx->fd);

      /* And the string tables.  */
      ebl_strtabfree (ctx->section_strtab);
      ebl_strtabfree (ctx->symbol_strtab);
    }

  /* Initialize the lock.  */
  rwlock_fini (ctx->lock);

  /* Finally free the data structure.   */
  free (ctx);
}