/*
 * Copyright 2002  John W. De Boskey <jwd@FreeBSD.Org>
 * All Rights Reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY JOHN W. DE BOSKEY ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL JOHN W. DE BOSKEY BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT
 * OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
 * USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
 * DAMAGE.
 *
 */


#include <elf.h>
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <link.h>
#include <unistd.h>
#include <machine/reg.h>
#include <signal.h>
#include <stdio.h>
#include <string.h>
#include <sys/syslimits.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <sys/mman.h>
#include <sys/ptrace.h>
#include <sys/stat.h>

#include "ptrace.h"
#include "elfutil.h"

pid_t    pid;               /* child process we're tracing       */

struct link_map linkmap;    /* linkmap struct from load_hook     */
struct r_debug  rdebug;     /* r_debug struct from load_hook     */

/*
 * Ok, provide the most simplistic usage function.
 */
void
usage(me)
char *me;
{
   fprintf(stderr,"Usage: %s executable <executable options>\n",me);
   exit(1);
}

char *
sigstr(sig)
int sig;
{
  switch (sig) {
    case SIGHUP    : return "SIGHUP";    break;
    case SIGINT    : return "SIGINT";    break;
    case SIGQUIT   : return "SIGQUIT";   break;
    case SIGILL    : return "SIGILL";    break;
    case SIGTRAP   : return "SIGTRAP";   break;
    case SIGABRT   : return "SIGABRT";   break;
    case SIGEMT    : return "SIGEMT";    break;
    case SIGFPE    : return "SIGFPE";    break;
    case SIGKILL   : return "SIGKILL";   break;
    case SIGBUS    : return "SIGBUS";    break;
    case SIGSEGV   : return "SIGSEGV";   break;
    case SIGSYS    : return "SIGSYS";    break;
    case SIGPIPE   : return "SIGPIPE";   break;
    case SIGALRM   : return "SIGALRM";   break;
    case SIGTERM   : return "SIGTERM";   break;
    case SIGURG    : return "SIGURG";    break;
    case SIGSTOP   : return "SIGSTOP";   break;
    case SIGTSTP   : return "SIGTSTP";   break;
    case SIGCONT   : return "SIGCONT";   break;
    case SIGCHLD   : return "SIGCHLD";   break;
    case SIGTTIN   : return "SIGTTIN";   break;
    case SIGTTOU   : return "SIGTTOU";   break;
    case SIGIO     : return "SIGIO";     break;
    case SIGXCPU   : return "SIGXCPU";   break;
    case SIGXFSZ   : return "SIGXFSZ";   break;
    case SIGVTALRM : return "SIGVTALRM"; break;
    case SIGPROF   : return "SIGPROF";   break;
    case SIGWINCH  : return "SIGWINCH";  break;
    case SIGINFO   : return "SIGINFO";   break;
    case SIGUSR1   : return "SIGUSR1";   break;
    case SIGUSR2   : return "SIGUSR2";   break;
    default        : return "UNKNOWN";   break;
  }
}



/*
 * Here is where we start the child process within which we
 * want to trace the loads of shared libs. Note that the
 * 'magic' is the ptrace() call between the fork() and exec().
 */
pid_t
start_child(argv)
char *argv[];
{
   int   rc;

   pid = fork();
   if (pid < 0) {
      fprintf(stderr, "can't fork: %s\n", strerror(errno) ); 
      exit(1);
   }

   if (pid == 0) { /* child */
      rc = ptrace(PT_TRACE_ME, getpid(), NULL, 0);
      if (rc < 0) {
         fprintf(stderr,"ptrace(PT_TRACE_ME)=%d: %s\n",errno,strerror(errno));
         _exit(1);
      }

      rc = execv(argv[0], argv);
      if (rc < 0) {
         fprintf(stderr,"execv(%s)=%d: %s\n",argv[0],errno,strerror(errno));  
         _exit(1);
      }

      /* not reached */
   }

   return(pid);
}


/*
 * Given the information about a watch, fill out the dbreg
 * structure. Currently, (the hardware) only supports up to
 * 4 hardware breaks.
 */
int
hw_watch(watchnum,watchaddr,size,access,d)
int watchnum;
unsigned int watchaddr;
int size;
unsigned int access;
struct dbreg *d;
{
   unsigned int dr0, dr1, dr2, dr3, dr7;
   unsigned int mask1, mask2, smask;
   int shift;

   switch ( access ) {
      case DR7_EXEC:
      case DR7_WRITEONLY:
      case DR7_READWRITE: break;
      default : fprintf(stderr,"hw_watch(): invalid access value %d\n",access);
      return -1;
   }

   /*
    * We can watch a 1, 2, or 4 byte sized location
    */
   switch (size) {
     case 1: smask = 0x00; break;
     case 2: smask = 0x01; break;
     case 4: smask = 0x03; break;
     default: fprintf(stderr,
                 "hw_watch(): illegal size value %d, must be 1, 2, or 4\n",
                  size );
              return -2;
   }

   mask1 = (smask << 2) | access;

   dr7 = d->dr7;
   dr0 = d->dr0;
   dr1 = d->dr1;
   dr2 = d->dr2;
   dr3 = d->dr3;
 
   mask2 = 0x02;  /* just enable the appropriate GE bit */

   switch (watchnum) {
      case 0 : dr0 = watchaddr; shift = 0; break;
      case 1 : dr1 = watchaddr; shift = 2; break;
      case 2 : dr2 = watchaddr; shift = 4; break;
      case 3 : dr3 = watchaddr; shift = 6; break;
      default : fprintf(stderr,"hw_watch(): invalid watch number %d\n",
                               watchnum);
              return -1;
              break;
   }

   dr7 = dr7 & ~((0x3 << shift) | (0x0f << (watchnum*4+16)));
   dr7 |= (mask2 << shift) | (mask1 << (watchnum*4+16));

   d->dr7 = dr7;
   d->dr0 = dr0;
   d->dr1 = dr1;
   d->dr2 = dr2;
   d->dr3 = dr3;

   return 0;
}



void
enable_hook(pid,interp)
pid_t pid;
__interp *interp;
{
   int rc;
   struct dbreg d;
   unsigned int pc;

   /*
    * Get the current program counter and relocate the hookpt the
    * first time through this code sequence.
    */
   if (!interp->relocated) {
      pc = GET_PC(pid);
      interp->hookpt = pc - interp->entrypt + interp->hookpt;
      interp->relocated = 1;
   }
   
   memset ( &d, 0, sizeof(d) );
   hw_watch(0,interp->hookpt,1,DR7_EXEC,&d);

   /*
    * We blow off any previous debug register settings here
    */
   rc = ptrace ( PT_SETDBREGS, pid, (caddr_t)&d, 0 );
   if (rc < 0) {
      fprintf(stderr,"Can't set debug registers for process %d: %s\n",
                      pid, strerror(errno) );
      exit(1);
   }

}

void
disable_hw_all(pid)
pid_t pid;
{
   struct dbreg d;
   int rc;

   memset ( &d, 0, sizeof(d) );
   rc = ptrace ( PT_SETDBREGS, pid, (caddr_t)&d, 0 );
   if (rc < 0) {
      fprintf(stderr,"Can't set debug registers for process %d: %s\n",
                      pid, strerror(errno) );
      exit(1);
   }

   return;
}

void
process_loadhook(pid)
pid_t pid;
{
   unsigned int sp;
   unsigned int p1,p2;
   char     filename[PATH_MAX];
   Elf_Ehdr hdr;
   char    *type = "notset";

   /*
    * Get the two parms (pointers) off the stack
    */
   sp = GET_SP(pid);
   p1 = get_long(pid,(caddr_t)sp+4);
   p2 = get_long(pid,(caddr_t)sp+8);

   /*
    * Always read the r_debug struct if available.
    */
   if (p1) {
      memset(filename,0,sizeof(filename));
      get_string(pid,(caddr_t)p1,(caddr_t)&rdebug,sizeof(rdebug));
      get_bytes(pid,(caddr_t)p1,(caddr_t)&rdebug,sizeof(rdebug));
   }

   /*
    * Always read the filename and elf header from the traced process
    */
   if (p2) {
      get_bytes(pid,(caddr_t)p2,(caddr_t)&linkmap,sizeof(linkmap));
      get_string(pid,(caddr_t)linkmap.l_name,(caddr_t)filename,PATH_MAX);
      get_bytes(pid,(caddr_t)linkmap.l_addr,(caddr_t)&hdr,sizeof(hdr));
      /*
       * Since this is a load let's look for the ELF eye-catcher
       */
      if (hdr.e_ident[EI_MAG0] != ELFMAG0 ||
          hdr.e_ident[EI_MAG1] != ELFMAG1 ||
          hdr.e_ident[EI_MAG2] != ELFMAG2 ||
          hdr.e_ident[EI_MAG3] != ELFMAG3) {
         fprintf(stdout,"Load   shlib: No_ElfMark : %s\n",filename);
      }
      switch (hdr.e_type) {
         case ET_NONE  : type = "ET_NONE"; break;
         case ET_REL   : type = "ET_REL" ; break;
         case ET_EXEC  : type = "ET_EXEC"; break;
         case ET_DYN   : type = "ET_DYN" ; break;
         case ET_CORE  : type = "ET_CORE"; break;
         default       : type = "UKNOWN" ; break;
      }
   }

   /*
    * Check for the hello call from the dynamic loader.
    */
   if (!p1 && p2) {
      fprintf(stdout,"L %7s: @0x%08x: %s\n",type,(uint)linkmap.l_addr,filename);
   } else if (p1 && p2) {

      /*
       * Check for a shared object being deleted via dlclose()
       */
      if (rdebug.r_state == RT_DELETE) {
         fprintf(stdout,"D %7s: @%p: %s\n",type,linkmap.l_addr,filename);
      }
      else {
         fprintf(stdout,"L %7s: @%p: %s\n",type,linkmap.l_addr,filename);
      }
   } else if (p1 && !p2) {
      /*
       * The (p1 && !p2) case is boring and uninteresting so
       * we ignore it.
       *
       *  RT_ADD, 0
       */

   }
   
   return;
}

void
process_child(interp)
__interp *interp;
{
   int done = 0;
   int rc;
   int cstatus;
   int installed   = 0;
   int singlestep  = 0;

   while (!done) {
      rc = waitpid(pid, &cstatus, WUNTRACED);
      if (rc < 0) {
         fprintf(stderr,"waitpid(%d): %s\n", pid, strerror(errno));
         exit(1);
      }   

      if (WIFEXITED(cstatus)) {
         fprintf(stdout,"process %d exited, exit code=%d\n",
                          pid, WEXITSTATUS(cstatus));
         done = 1;
         
      }
      else if (WIFSIGNALED(cstatus)) {
         fprintf(stderr,"process %d terminated by signal %d%s\n",
                 pid, WTERMSIG(cstatus),
                 WCOREDUMP(cstatus) ? " (core dumped)" : "" );
         done = 1;
      }
      else if (WIFSTOPPED(cstatus)) {
         int sig;

         sig = WSTOPSIG(cstatus);

         switch ( sig ) {
            case SIGTRAP:
            case SIGSTOP:
                          if (installed && !singlestep) {
                             process_loadhook(pid);
                             disable_hw_all(pid);
                             set_singlestep(pid);
                             installed = 0;
                             singlestep = 1;
                          }
                          else if (singlestep) {
                             enable_hook(pid,interp);
                             installed = 1;
                             singlestep = 0;
                          }
                          else if (!installed) {
                             enable_hook(pid,interp);
                             installed = 1;
                             singlestep = 0;
                          }
                          else {
                          }
                          break;
            case SIGILL:
            case SIGBUS:
            case SIGSEGV: rc = ptrace(PT_DETACH, pid, NULL, 0);
                          if (rc < 0) {
                             fprintf(stderr, "ptrace(PT_DETACH): %s\n",
                                              strerror(errno));
                          }
                          done = 1;
                          break;
         }
         if (!done) {
            rc = ptrace(PT_CONTINUE, pid, (caddr_t)1, 0);
            if (rc < 0) {
               fprintf(stderr,"Can't make process %d continue: %s\n",
                              pid, strerror(errno) );
               exit(1);
            }
         }
      }

   }

   return;
}

int
main(argc,argv)
int argc;
char *argv[];
{
   int   fd;
   char *filename;
   char *filebuf;
   __interp _interp = { 0 };

   if (argc < 2)
      usage(argv[0]);
   filename = argv[1];

   open_elf_image(filename,&fd,&filebuf);

   _interp.name = get_elf_interp(filebuf);
   if (! _interp.name) {
      fprintf(stderr,"Could not locate [.interp] section in %s\n",filename);
      exit(1);
   }

   get_interp_info(&_interp);

   start_child(&argv[1]);

   process_child(&_interp);

   exit(1);
}