/*
    DXE2GEN - dynamically-loadable modules builder for DJGPP

    Copyright (C) 2000 Andrew Zabolotny <bit@eltech.ru>
    Partly based on work by Charles Sandmann and DJ Delorie.

    Usage of this library is not restricted in any way.
    The full license text can be found in the file dxe.txt.
*/

/*
    Program exit codes:
       0: o.k.
      -1: wrong command line
      -2: i/o error
      -3: unresolved symbols
      otherwise the exit code of GNU ld is returned
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <process.h>
#include <dir.h>
#include <errno.h>
#include <coff.h>
#include <dxe2.h>

#define TEMP_O_FILE	"$$dxe$$.o"
#define TEMP_S_FILE	"$$dxe$$.s"

#define VALID_RELOC(r) ((r.r_type != RELOC_REL32) && (r.r_symndx != -1UL))

static char *progname;
static char *version = "0.1.0";

// Command-line options
struct
{
  bool unresolved;		// allow unresolved symbols in output
  bool verbose;			// verbose output
  bool showexp;			// show exported symbols
  bool showunres;		// show unresolved symbols
  char *output;			// output file name
  int objcount;			// number of object files given on command line
  char *implib;			// name of import library
  char *objfile;		// the first object file on command line
  char *dxefile;		// the name of dxe file on command line
  char *description;		// a description of the module
  unsigned num_prefix;		// number of exported prefixes
  unsigned max_prefix;		// maximal number of exported prefixes
  char **export_prefix;		// exported symbol prefixes
} opt =
{
  false,
  false,
  false,
  false,
  NULL,
  0,
  NULL,
  NULL,
  NULL,
  "",
  0,
  0,
  NULL
};

void exit_cleanup ()
{
  remove (TEMP_O_FILE);
  remove (TEMP_S_FILE);
}

void display_version ()
{
  printf ("dxe2gen version %s\n", version);
}

void display_help ()
{
  display_version ();
  printf ("Usage: dxe2gen [-o output.dxe] [options] [object-files] [ld-options]\n");
  printf ("Create a dynamically-loadable executable module for DJGPP\n\n");
  printf ("-o output.dxe\tDefine the name of output DXE file\n");
  printf ("-D description\tSet module description string\n");
  printf ("-E prefix\tExport only symbols that start with <prefix> (cumulative)\n");
  printf ("-I import.a\tCreate an import library for given DXE file\n");
  printf ("-U\t\tAllow unresolved symbols in DXE file\n");
  printf ("-V\t\tVerbose output (minimal output by default)\n");
  printf ("--show-exp\tShow symbols exported by the DXE module\n");
  printf ("--show-unres\tShow unresolved symbols in the DXE module\n");
  printf ("[ld-options]\tAny other options are passed unchanged to ld\n\n");
  printf ("If the DXE file contains unresolved symbols, you should provide\n");
  printf ("the appropiate environment at load time, so that they can be\n");
  printf ("resolved dynamically at runtime.\n");
  exit (-1);
}

void process_args (int argc, char *argv[], char *new_argv[])
{
  new_argv [0] = "ld";
  new_argv [1] = "-X";
  new_argv [2] = "-S";
  new_argv [3] = "-r";
  new_argv [4] = "-o";
  new_argv [5] = TEMP_O_FILE;
  new_argv [6] = "-L";
  static char libdir [50];
  sprintf (libdir, "%s/lib", getenv ("DJDIR"));
  new_argv [7] = libdir;
  new_argv [8] = "-T";
  new_argv [9] = "dxe2.ld";
  int new_argc = 10;

  for (int i = 1; i < argc; i++)
  {
    if (!strcmp (argv [i], "-h") || !strcmp (argv [i], "--help"))
      display_help ();
    else if (!strcmp (argv [i], "-U"))
      opt.unresolved = true;
    else if (!strcmp (argv [i], "-o"))
      opt.output = argv [++i];
    else if (!strcmp (argv [i], "-V"))
      opt.verbose = true;
    else if (!strcmp (argv [i], "-D"))
      opt.description = argv [++i];
    else if (!strcmp (argv [i], "-I"))
      opt.implib = argv [++i];
    else if (!strcmp (argv [i], "-E"))
    {
      if (opt.num_prefix >= opt.max_prefix)
      {
        opt.max_prefix += 16;
        opt.export_prefix = (char **)realloc (opt.export_prefix,
          opt.max_prefix * sizeof (char *));
      }
      opt.export_prefix [opt.num_prefix++] = argv [++i];
    }
    else if (!strcmp (argv [i], "--show-exp"))
      opt.showexp = true;
    else if (!strcmp (argv [i], "--show-unres"))
      opt.showunres = true;
    else
    {
      new_argv [new_argc++] = argv [i];
      char *dot = strrchr (argv [i], '.');
      if (dot)
      {
        if (!strcasecmp (dot, ".o") || !strcasecmp (dot, ".a"))
        {
          opt.objcount++;
          if (!opt.objfile)
            opt.objfile = argv [i];
        }
        else if (!strcasecmp (dot, ".dxe"))
          opt.dxefile = argv [i];
      }
    }
  }
  new_argv [new_argc] = NULL;
}

FILE *run_ld (char *argv [], FILHDR &fh)
{
  if (opt.objcount == 1)
  {
    // See if the object file has just one section
    FILE *f = fopen (opt.objfile, "rb");
    if (!f)
    {
      fprintf (stderr, "%s: cannot open file `%s'\n", progname, opt.objfile);
      exit (-2);
    }

    fread (&fh, 1, FILHSZ, f);
    if (fh.f_nscns == 1)
      return f;
    fclose (f);
  }

  // Despite what DJGPP docs says, spawn does NOT search along the PATH
  int rc = spawnvp (P_WAIT, "ld.exe", argv);
  if (rc)
  {
    if (rc == -1) perror ("ld.exe");
    exit (rc);
  }

  FILE *f = fopen (TEMP_O_FILE, "rb");
  if (!f)
  {
    fprintf (stderr, "%s: cannot open linker output file `%s'\n", progname, TEMP_O_FILE);
    exit (-2);
  }
  else
    atexit (exit_cleanup);

  fread (&fh, 1, FILHSZ, f);
  if (fh.f_nscns != 1)
  {
    fclose (f);
    fprintf (stderr, "%s: linker output file has more than one section\n", TEMP_O_FILE);
    exit (-2);
  }

  return f;
}

int write_dxe (FILE *inf, FILE *outf, FILHDR &fh)
{
  // Get the section header
  SCNHDR sc;
  fseek (inf, fh.f_opthdr, SEEK_CUR);
  fread (&sc, 1, SCNHSZ, inf);

  dxe2_header dh;
  dh.magic = DXE2_MAGIC;
  dh.sec_size = sc.s_size;
  dh.n_relocs = sc.s_nreloc;
  dh.n_exp_syms = 0;
  dh.n_unres_syms = 0;

  // Read section data
  char *data = new char [sc.s_size];
  fseek (inf, sc.s_scnptr, 0);
  fread (data, 1, sc.s_size, inf);

  // Read all symbols
  SYMENT *sym = new SYMENT [fh.f_nsyms];
  fseek (inf, fh.f_symptr, SEEK_SET);
  fread (sym, fh.f_nsyms, SYMESZ, inf);

  // Read symbol name table
  size_t stsz;
  fread (&stsz, 1, sizeof (stsz), inf);
  char *strings = new char [stsz];
  fread (strings + 4, 1, stsz - 4, inf);
  strings [0] = 0;

  // Read the relocation table
  RELOC *relocs = new RELOC [sc.s_nreloc];
  fseek (inf, sc.s_relptr, SEEK_SET);
  fread (relocs, sc.s_nreloc, RELSZ, inf);

  // Close input file
  fclose (inf);

  // Exported symbols table
  char *expsym_table = NULL;
  size_t expsym_size = 0, expsym_maxsize = 0;
  // Unresolved symbols table
  char *unres_table = NULL;
  size_t unres_size = 0, unres_maxsize = 0;

  unsigned i, j, errcount = 0;

  for (i = 0; i < fh.f_nsyms; i += 1 + sym [i].e_numaux)
  {
    char tmp [9], *name;
    if (sym [i].e.e.e_zeroes)
    {
      memcpy (tmp, sym [i].e.e_name, E_SYMNMLEN);
      tmp [E_SYMNMLEN] = 0;
      name = tmp;
    }
    else
      name = strings + sym[i].e.e.e_offset;
    int namelen = strlen (name) + 1;

    // If the symbol begins with '_', skip it
    bool underline = (name [0] == '_');
    if (underline)
      name++, namelen--;

    // sanity check
    if (namelen <= 1)
      continue;

#if 0
    printf("[%3d] 0x%08lx 0x%08x 0x%04x %d %s\n",
           i,
           sym [i].e_value,
           sym [i].e_scnum,
           sym [i].e_sclass,
           sym [i].e_numaux,
           name
           );
#endif

    if (sym [i].e_scnum == 0)
    {
      // unresolved symbol
      dh.n_unres_syms++;

      // count the amount of relocations pointing to this symbol
      int n_abs_relocs = 0, n_rel_relocs = 0;
      for (j = 0; j < sc.s_nreloc; j++)
        if (relocs [j].r_symndx == i)
          if (relocs [j].r_type == RELOC_REL32)
            n_rel_relocs++;
          else
            n_abs_relocs++;

      // If there are no references to this symbol, skip it
      if (n_rel_relocs == 0 && n_abs_relocs == 0)
        continue;

      if (!opt.unresolved)
      {
        fprintf (stderr, "%s: unresolved symbol `%s'\n", progname, name);
        errcount++; continue;
      }

      if (n_rel_relocs > 0xffff || n_abs_relocs > 0xffff)
      {
        fprintf (stderr, "%s: FATAL ERROR: too many relocations for unresolved"
          " symbol `%s'\n", progname, name);
        fclose (outf);
        remove (opt.output);
        exit (-4);
      }

      size_t newsize = unres_size + 2 * sizeof (short) + namelen +
        (n_rel_relocs + n_abs_relocs) * sizeof (long);
      if (newsize > unres_maxsize)
        unres_table = (char *)realloc (unres_table, unres_maxsize += 1024);
      memcpy (unres_table + unres_size + 2 * sizeof (short), name, namelen);

      // Store number of references to this unresolved symbol
      short *count = (short *)(unres_table + unres_size);
      count [0] = n_rel_relocs;
      count [1] = n_abs_relocs;

      long *rel_relocs = (long *)(unres_table + unres_size +
        2 * sizeof (short) + namelen);
      long *abs_relocs = (long *)(unres_table + newsize);

      unres_size = newsize;

      for (j = 0; j < sc.s_nreloc; j++)
        if (relocs [j].r_symndx == i)
        {
          // mark the relocation as processed
          relocs [j].r_symndx = -1UL;

          if (relocs [j].r_type == RELOC_REL32)
            *rel_relocs++ = relocs [j].r_vaddr;
          else
            *--abs_relocs = relocs [j].r_vaddr;
        }
      if (opt.verbose)
        printf ("unresolved: `%s' (%d references)\n", name, n_rel_relocs + n_abs_relocs);
    }
    else if ((sym [i].e_sclass == C_EXT)
          && underline
          && strncmp (name, "_GLOBAL_$", 9))
    {
      if (opt.num_prefix)
      {
        bool ok = false;
        for (j = 0; j < opt.num_prefix; j++)
          if (memcmp (opt.export_prefix [j], name, strlen (opt.export_prefix [j])) == 0)
          { ok = true; break; }
        if (!ok) continue;
      }

      // exported symbol
      dh.n_exp_syms++;

      size_t newsize = expsym_size + sizeof (long) + namelen;
      if (newsize > expsym_maxsize)
        expsym_table = (char *)realloc (expsym_table, expsym_maxsize += 1024);

      *(long *)(expsym_table + expsym_size) = sym [i].e_value;
      memcpy (expsym_table + expsym_size + sizeof (long), name, namelen);
      expsym_size = newsize;
      if (opt.verbose)
        printf ("export: `%s'\n", name);
    }
  }

  if (errcount)
  {
    fclose (outf);
    remove (opt.output);
    exit (-3);
  }

  // Compute the amount of valid relocations
  for (i = 0; i < sc.s_nreloc; i++)
  {
#if 0
    printf ("[%3d] %08lX %03ld %04X\n",
      i,
      relocs [i].r_vaddr,
      relocs [i].r_symndx,
      relocs [i].r_type);
#endif
    if (!VALID_RELOC (relocs [i]))
      dh.n_relocs--;
  }

  // A small array for padding with zeros
  char fill [16];
  memset (fill, 0, 16);

  // Compute the symbol tables offset
  size_t desc_len = strlen (opt.description) + 1;
  if (desc_len <= 1) desc_len = 0;
  dh.sym_f_offset = (sizeof (dh) + desc_len + 15) & ~15UL;

  // Compute the offset of .text+.data+.bss sections
  size_t hdrsize = dh.sym_f_offset + expsym_size + unres_size +
    dh.n_relocs * sizeof (long);
  dh.sec_f_offset = (hdrsize + 15) & ~15UL;

  dh.sec_f_size = dh.sec_size;
  while (dh.sec_f_size && !data [dh.sec_f_size - 1])
    dh.sec_f_size--;

  // Output the DXE2 header
  fwrite (&dh, 1, sizeof (dh), outf);

  // If we have a description string, put it here
  if (desc_len)
    fwrite (opt.description, 1, desc_len, outf);
  fwrite (fill, 1, (dh.sec_f_offset - desc_len) & 15, outf);

  // Output the exported symbols table
  fwrite (expsym_table, 1, expsym_size, outf);
  free (expsym_table);

  // Output the unresolved symbols table
  fwrite (unres_table, 1, unres_size, outf);
  free (unres_table);

  // Output the relocations
  for (i = 0; i < sc.s_nreloc; i++)
    if (VALID_RELOC (relocs [i]))
      fwrite (&relocs [i].r_vaddr, 1, sizeof (long), outf);

  // Finally, write the actual code+data+bss section
  fwrite (fill, 1, dh.sec_f_offset - hdrsize, outf);
  fwrite (data, 1, dh.sec_f_size, outf);
  delete [] data;

  fclose (outf);

  delete [] strings;
  delete [] relocs;
  delete [] sym;

  return 0;
}

static FILE *open_dxe_file (const char *fname, dxe2_header &dh)
{
  FILE *f = fopen (fname, "rb");
  if (!f)
  {
    fprintf (stderr, "%s: cannot read DXE module `%s'\n", progname, fname);
    exit (-2);
  }

  fread (&dh, 1, sizeof (dh), f);
  if (dh.magic != DXE2_MAGIC)
  {
    fclose (f);
    fprintf (stderr, "%s: the file `%s' is not an DXE module\n", progname, fname);
    exit (-2);
  }

  return f;
}

int make_implib ()
{
  int i;
  char *scan;
  dxe2_header dh;
  FILE *f = open_dxe_file (opt.dxefile, dh);

  fseek (f, dh.sym_f_offset, SEEK_SET);
  size_t symtsize = dh.sec_f_offset - dh.sym_f_offset;
  char *symtab = new char [symtsize];
  fread (symtab, 1, symtsize, f);
  fclose (f);

  // Compute the base name of the library
  char basename [8];
  scan = strchr (opt.dxefile, 0);
  while ((scan > opt.dxefile) && (scan [-1] != '/') && (scan [-1] != '\\')
      && (scan [-1] != ':'))
    scan--;
  for (i = 0; *scan && *scan != '.'; i++, scan++)
    basename [i] = *scan;
  basename [i] = 0;
  strupr (basename);

  FILE *implib = fopen (TEMP_S_FILE, "w");
  if (!implib)
  {
    delete [] symtab;
    fprintf (stderr, "%s: cannot open file `%s' for writing\n", progname, TEMP_S_FILE);
    exit (-2);
  }

  fprintf (implib,
       ".text\n"
       "Lmodh:	.long	0\n"
       "Lmodn:	.ascii	\"%s.DXE\\0\"\n"
       "	.globl	_load_%s\n"
       "_load_%s:\n"
       "	pushl	$Lsyms\n"		// symbol names
       "	pushl	$Lstubs\n"		// stubs
       "	pushl	$Lmodh\n"		// module handle
       "	pushl	$Lmodn\n"		// module name
       "	call	_dlstatbind\n"		// statically bind
       "	addl	$16,%%esp\n"
       "	ret\n"
       "	.globl	_unload_%s\n"
       "_unload_%s:\n"
       "	pushl	$Lload\n"		// loader address
       "	pushl	$Lsyms\n"		// symbol names
       "	pushl	$Lstubs\n"		// stubs
       "	pushl	$Lmodh\n"		// module handle
       "	pushl	$Lmodn\n"		// module name
       "	call	_dlstatunbind\n"	// unbind module
       "	addl	$20,%%esp\n"
       "	ret\n"
       "Lload:	pushal\n"
       "	call	_load_%s\n"
       "	popal\n"
       "	subl	$5,(%%esp)\n"
       "	ret\n"
       , basename, basename, basename, basename, basename, basename);

  // Emit the names of all imported functions
  fprintf (implib, "Lsyms:\n");
  for (i = 0, scan = symtab; i < dh.n_exp_syms; i++)
  {
    scan += sizeof (long);
    // Do NOT export djgpp_first_ctor, djgpp_last_ctor and so on
    // since you will get duplicate symbols if you link with more
    // than one import library
    if (memcmp (scan, "djgpp", 5))
      fprintf (implib, "\t.ascii\t\"%s\\0\"\n", scan);
    scan = strchr (scan, 0) + 1;
  }
  fprintf (implib, "\t.byte\t0\n");

  // And now emit the stubs
  for (i = 0, scan = symtab; i < dh.n_exp_syms; i++)
  {
    scan += sizeof (long);
    if (memcmp (scan, "djgpp", 5))
    {
      fprintf (implib, "\t.align\t2,0xcc\n");
      if (i == 0)
        fprintf (implib, "Lstubs:\n");
      fprintf (implib, "\t.globl\t_%s\n_%s:\n\tcall\tLload\n", scan, scan);
    }
    scan = strchr (scan, 0) + 1;
  }

  fclose (implib);
  delete [] symtab;

  // We already have what to clean up
  atexit (exit_cleanup);

  // Allright, now run the assembler on the resulting file
  int rc = spawnlp (P_WAIT, "as.exe", "as", "-o", TEMP_O_FILE, TEMP_S_FILE, NULL);
  if (rc)
  {
    if (rc == -1) perror ("as.exe");
    exit (rc);
  }

  // Okey-dokey, let's stuff the object file into the archive
  rc = spawnlp (P_WAIT, "ar.exe", "ar", "crs", opt.implib, TEMP_O_FILE, NULL);
  if (rc)
  {
    if (rc == -1) perror ("ar.exe");
    exit (rc);
  }

  return 0;
}

static int show_symbols (const char *fname)
{
  dxe2_header dh;
  FILE *f = open_dxe_file (fname, dh);

  fseek (f, dh.sym_f_offset, SEEK_SET);
  size_t symtsize = dh.sec_f_offset - dh.sym_f_offset;
  char *symtab = new char [symtsize];
  fread (symtab, 1, symtsize, f);
  fclose (f);

  int i;
  char *scan;
  for (i = 0, scan = symtab; i < dh.n_exp_syms; i++)
  {
    scan += sizeof (long);
    if (opt.showexp)
      puts (scan);
    scan = strchr (scan, 0) + 1;
  }

  for (i = 0; i < dh.n_unres_syms; i++)
  {
    unsigned short n1 = ((unsigned short *)scan) [0];
    unsigned short n2 = ((unsigned short *)scan) [1];
    scan += sizeof (short) * 2;
    if (opt.showunres)
      puts (scan);
    scan = strchr (scan, 0) + 1 + (n1 + n2) * sizeof (long);
  }

  return 0;
}

int main(int argc, char *argv[])
{
  progname = argv [0];

  // Prepare the command line for ld
  char *new_argv[argc + 11];
  process_args (argc, argv, new_argv);

  if (opt.showexp || opt.showunres)
  {
    for (int i = 1; i < argc; i++)
    {
      if (argv [i][0] == '-')
        continue;
      char *dot = strchr (argv [i], '.');
      if (dot && !strcasecmp (dot, ".dxe"))
      {
        int rc = show_symbols (argv [i]);
        if (rc) return rc;
      }
    }
    return 0;
  }

  if (!opt.output && (!opt.implib || opt.objcount))
  {
    fprintf (stderr, "%s: no output file name given (-h for help)\n", progname);
    exit (-1);
  }

  if (opt.output && !opt.objcount)
  {
    fprintf (stderr, "%s: no object file(s) given (-h for help)\n", progname);
    exit (-1);
  }

  if (opt.implib)
  {
    if (!opt.dxefile)
      opt.dxefile = opt.output;
    if (!opt.dxefile)
    {
      fprintf (stderr, "%s: no DXE module name given (-h for help)\n", progname);
      exit (-1);
    }
  }

  if (opt.verbose)
  {
    // print the command line for ld
    for (int i = 0; new_argv [i]; i++)
      printf ("%s ", new_argv [i]);
    printf ("\n");
  }

  if (opt.objcount)
  {
    // Run linker
    int rc;
    FILHDR fh;
    FILE *inf = run_ld (new_argv, fh);

    // Now `inf' is an opened single-section COFF module
    // Create the output file
    FILE *outf = fopen (opt.output, "wb");
    if (!outf)
    {
      fclose (inf);
      fprintf (stderr, "%s: cannot open file `%s' for writing\n", progname, opt.output);
      exit (-2);
    }

    // Allright, now write the DXE file and quit
    rc = write_dxe (inf, outf, fh);
    if (rc) return rc;
  }

  if (opt.implib)
  {
    int rc = make_implib ();
    if (rc) return rc;
  }

  return 0;
}