/* elfgrep: Searches for object files inside other object files.
 * Copyright (C) 2002 Dion Mendel
 *
 * This program 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 2
 * of the License, or (at your option) any later version.
 *
 * This program 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 this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330,
 * Boston, MA 02111-1307, USA.
 */

#include <assert.h>
#include <errno.h>
#include <stdio.h>
#include <stdarg.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>


/*--------------------------------- Macros -----------------------------------*/

#define OBJDUMP "/usr/bin/objdump"    /* full path to objdump executable */

#define MAXLINE 4096                  /* maximum buffer size */
#define WILDCARD ((char) 0xf)         /* wildcard to replace relocation data */
#define ARRAY_INCREMENT 200           /* increase array by this many elements */

#ifndef FALSE
# define FALSE 0
#endif

#ifndef TRUE
# define TRUE !FALSE
#endif

static char *program_name = "elfgrep";

/*----------------------------- Error handling -------------------------------*/

/* Fatal error related to a system call.
 * Print a message and terminate. */
static void
err_sys (const char *fmt, ...)
{
   char buf[MAXLINE];
   va_list ap;

   va_start(ap, fmt);

   vsprintf(buf, fmt, ap);
   sprintf(buf+strlen(buf), ": %s", strerror(errno));
   strcat(buf, "\n");
   fflush(stdout);         /* in case stdout and stderr are the same */
   fputs(buf, stderr);
   fflush(NULL);           /* flushes all stdio output streams */

   va_end(ap);
   exit(EXIT_FAILURE);
}

/* Fatal error unrelated to a system call.
 * Print a message and terminate. */
static void
err_quit (const char *fmt, ...)
{
   char buf[MAXLINE];
   va_list ap;

   va_start(ap, fmt);

   vsprintf(buf, fmt, ap);
   strcat(buf, "\n");
   fflush(stdout);         /* in case stdout and stderr are the same */
   fputs(buf, stderr);
   fflush(NULL);           /* flushes all stdio output streams */

   va_end(ap);
   exit(EXIT_FAILURE);
}

/*---------------------------- Memory allocation -----------------------------*/

/* A safe realloc that will allocate
 * the memory or exit the program. */
static void *
xrealloc (void *ptr, size_t size)
{
   void *ret = NULL;

   if (size == 0) {
      free(ptr);
   }
   else if ( (ret = (void *) realloc(ptr, size)) == NULL) {
      fputs("Error: out of memory\n", stderr);
      exit(EXIT_FAILURE);
   }

   return ret;
}

/* A safe malloc that will allocate
 * the memory or exit the program. */
static void *
xmalloc (size_t size)
{
   void *ptr;

   /* ensure do not try to malloc 0 bytes */
   if (size == 0)
      size = 1;

   if ( (ptr = (void *) malloc(size)) == NULL) {
      fputs("Error: out of memory\n", stderr);
      exit(EXIT_FAILURE);
   }

   return ptr;
}

/* Make a copy of a string in a newly
 * allocated block of memory. */
static char *
xstrdup (char *str)
{
   void *ptr;

   ptr = xmalloc(strlen (str) + 1);
   strcpy(ptr, str);
   return ptr;
}

/*---------------------------- Command execution -----------------------------*/

typedef struct {
   size_t allocated;                /* allocated slots */
   size_t used;                     /* used slots */
   char **lines;                    /* lines of results */
} ExecResult;

#define exec_result_length(e) ((e)->used)
#define exec_result_index(e, i) ((e)->lines[(i)])

/*
 * Releases resources allocated to the result of a call to exec_cmd.
 */
static void
release_exec_result (ExecResult *er)
{
   int i;
   assert(er != NULL);

   for (i = 0; i < exec_result_length(er); i++) {
      free(exec_result_index(er, i));
   }
   free(er->lines);
   free(er);
}

/*
 * Similar to the system() function except returns the contents of the stdout
 * of the executed command.
 * The optional arguments are a list of NULL terminated arguments to pass
 * to exec.
 * TODO: handle stderr and return value of child
 */
static ExecResult *
exec_cmd (char *exec, ...)
{
   ExecResult *result = NULL;
   char *argv[100];
   va_list va;
   char *arg;
   int i;
   int status;
   int fd[2];
   FILE *fptr;
   char buffer[MAXLINE];
   pid_t pid;

   /* prepare argv vector by adding exec command plus optional arguments */
   memset(argv, 0, 100 * sizeof(char *));
   i = 0;
   argv[i++] = exec;  /* argv[0] */

   va_start(va, exec);
   do {
      /* copy all arguments, including the NULL */
      arg = va_arg(va, char *);
      argv[i++] = arg;
   } while (arg != NULL && i < 100);
   va_end(va);

   if (pipe(fd) < 0)
      err_sys("pipe error");

   if ( (pid = fork()) < 0)
      err_sys("fork error");
   else if (pid == 0) {                        /* child */
      close(fd[0]);    /* close read end */
/*      close(STDERR_FILENO);
*/
      if (fd[1] != STDOUT_FILENO) {
         if (dup2(fd[1], STDOUT_FILENO) != STDOUT_FILENO)
            err_sys("dup2 error to stdout");
         close(fd[1]);       /* don't need this after dup2 */
      }

      if (execvp(exec, argv) < 0)
         err_sys("execvp error for %s", exec);
   }
   else {                                      /* parent */
      close(fd[1]);    /* close write end */

      /* create results structure */
      result = xmalloc(sizeof(*result));
      result->used = 0;
      result->allocated = 0;
      result->lines = NULL;

      fptr = fdopen(fd[0], "r");
      while ( (fgets(buffer, MAXLINE, fptr)) != NULL) {
         /* remove '\n' from end of buffer */
         if (buffer[strlen(buffer) - 1] == '\n')
            buffer[strlen(buffer) - 1] = '\0';

         /* store results */
         if (result->used >= result->allocated) {
            /* increase lines array */
            result->allocated += ARRAY_INCREMENT;
            result->lines = xrealloc(result->lines,
                                     result->allocated * sizeof(char *));
         }
         result->lines[result->used] = xstrdup(buffer);
         result->used++;
      }
      if (ferror(fptr))
         err_sys("fgets error");
      fclose(fptr);      /* also closes fd[0] */

      if (waitpid(pid, &status, 0) < 0)
         err_sys("waitpid error");
      if (status != 0)
         err_quit("%s exited with non zero return value", exec);
   }

   return result;
}

/*---------------------------- Objdump processing ----------------------------*/

typedef struct {
   char *name;                      /* name of the object file */
   int start_addr;                  /* starting address */
   int offset;                      /* offset of section in file */
   int size;                        /* size of section in bytes */
   int has_reloc;                   /* does object file have relocatable data */
   char *signature;                 /* section signature */
} ObjInfo;

/*
 * Releases resources allocated to the result of a call to read_obj_info.
 */
static void
release_obj_info (ObjInfo *obj)
{
   assert(obj != NULL);

   free(obj->signature);
   free(obj);
}

/*
 * Reads in a signature for the specified section of the given object file.
 */
static ObjInfo *
read_obj_info (char *filename, char *section)
{
   ExecResult *res;
   ObjInfo *obj;
   int i;
   char *line, *ptr;
   FILE *fptr;
   int offset;
   char buf[12];

   assert(filename != NULL);
   assert(section != NULL);
   assert(section[0] == '.');

   /* create and initialise obj info structure */
   obj = xmalloc(sizeof(*obj));
   obj->name = filename;
   obj->size = 0;
   obj->offset = 0;
   obj->start_addr = 0;
   obj->has_reloc = FALSE;
   obj->signature = NULL;

   /* determine the start_address, plus offset and size of the data for the */
   /* given section */
   if ( (res = exec_cmd(OBJDUMP, "-f", "-h", filename, NULL)) != NULL) {
      for (i = 0; i < exec_result_length(res); i++) {
         line = exec_result_index(res, i);

         /* does the file have relocatable data? */
         if (strstr(line, "HAS_RELOC") == line) {
            obj->has_reloc = TRUE;
         }
         /* is section size and offset defined on this line? */
         if ( (ptr = strstr(line, section)) != NULL) {
            /* read size */
            if (sscanf(line+18, "%08x", &obj->size) != 1) {
               err_quit("could not parse %s output", OBJDUMP);
               break;
            }
            /* read start address */
            if (sscanf(line+28, "%08x", &obj->start_addr) != 1) {
               err_quit("could not parse %s output", OBJDUMP);
               break;
            }
            /* read offset */
            if (sscanf(line+48, "%08x", &obj->offset) != 1) {
               err_quit("could not parse %s output", OBJDUMP);
               break;
            }
         }
      }

      release_exec_result(res);
   }

   /* read in data for the given section */
   if ( (fptr = fopen(filename, "rb")) != NULL) {
      fseek(fptr, obj->offset, SEEK_SET);
      obj->signature = xmalloc(obj->size * sizeof(*obj->signature));
      if (fread(obj->signature, 1, obj->size, fptr) != obj->size) {
         err_quit("fread for file `%s' was short", filename);
      }
      fclose(fptr);
   }
   else {
      err_sys("fopen failed");
   }

   if (obj->has_reloc) {
      res = exec_cmd(OBJDUMP, "-j", section, "-r", filename, NULL);
      if (res != NULL) {
         /* skip first 5 lines of output */
         for (i = 5; i < exec_result_length(res); i++) {
            line = exec_result_index(res, i);

            /* skip empty lines */
            if (line[0] ==  '\0') {
               continue;
            }

            /* get offset of relocation data */
            if (sscanf(line, "%08x %10s", &offset, buf) != 2) {
               err_quit("could not parse %s output", OBJDUMP);
            }

            /* ensure type is known */
            if ((strcmp(buf, "R_386_32") != 0) &&
                (strcmp(buf, "R_386_PC32") != 0))
            {
               err_quit("unknown relocation type `%s'", buf);
            }

            /* replace relocation data with known wildcard */
            obj->signature[offset + 0] = WILDCARD;
            obj->signature[offset + 1] = WILDCARD;
            obj->signature[offset + 2] = WILDCARD;
            obj->signature[offset + 3] = WILDCARD;
         }

         release_exec_result(res);
      }
   }

   return obj;
}

/*
 * Finds whether the signature of needle exists anywhere in haystack.  The
 * offsets of any matches are stored int the given matches array.
 * Returns the number of matches found.
 */
static int
find_matches (ObjInfo *needle, ObjInfo *haystack, int *matches, int max_matches)
{
   int looping, match;                        /* boolean */
   int i, pos, needle_size, needle_data;
   int num_matches;

   num_matches = 0;
   pos = 0;
   looping = TRUE;
   needle_size = needle->size;

   /* early exit if nothing to search for */
   if (needle_size == 0)
      return 0;

   while (looping) {

      if ((haystack->size - pos) >= needle_size) {
         /* assume match at this position until proven otherwise */
         match = TRUE;
         for (i = 0; i < needle_size; i++) {
            needle_data = needle->signature[i];

            if (needle_data != haystack->signature[pos + i] &&
                  needle_data != WILDCARD)
            {
               /* no match */
               match = FALSE;
               break;
            }
         }

         if (match && (num_matches < max_matches)) {
            /* store match */
            matches[num_matches++] = (pos + haystack->start_addr);
         }
      }
      else {   /* (haystack->size - pos) < needle_size) */
         looping = FALSE;
      }

      /* increment search */
      pos++;
   }

   return num_matches;
}

/* Prints brief usage message and exits. */
void
usage ()
{
   printf("%s - Search for the existance of OBJFILE inside FILE\n",
          program_name);
   printf("Usage: %s [OPTION] <OBJFILE> <FILE>\n", program_name);
   printf("\
Options:

Select one of:
  -t                         compare .text sections
  -d                         compare .data sections
  -r                         compare .rodata sections
");
   exit(EXIT_FAILURE);
}

int
main (int argc, char *argv[])
{
   ObjInfo *haystack, *needle;
   int matches[20];
   int i, num;
   char *section;

   if (argc != 4) {
      usage();
   }

   /* process single option */
   if (argv[1][0] == '-') {
      switch (argv[1][1]) {
         case 'd': section = ".data"; break;
         case 'r': section = ".rodata"; break;
         case 't': section = ".text"; break;
         default: usage();
      }
   }
   else {
      usage();
   }

   /* read objdump info */
   needle = read_obj_info(argv[2], section);
   haystack = read_obj_info(argv[3], section);

   /* search for matches and print the result */
   num = find_matches(needle, haystack, matches, 20);
   if (num == 0) {
      printf("%s - no matches\n", needle->name);
   }
   else if (num == 1) {
      printf("%s - match at 0x%08x (0x%08x bytes)\n", needle->name,
             matches[0], needle->size);
   }
   else {
      for (i = 0; i < num; i++) {
        printf("%s - match %d at 0x%08x (0x%08x bytes)\n", needle->name, i,
               matches[i], needle->size);
      }
   }

   /* release resources */
   release_obj_info(haystack);
   release_obj_info(needle);

   return EXIT_SUCCESS;
}