#include <stdio.h>
#include <fcntl.h>
#include "BPatch.h"
#include "BPatch_process.h"
#include "BPatch_function.h"
#include "BPatch_Vector.h"
#include "BPatch_thread.h"

/*
 * retee.C
 *
 * This program (mutator) provides an example of several facets of
 * Dyninst's behavior, and is a good basis for many Dyninst
 * mutators. We want to intercept all output from a target application
 * (the mutatee), duplicating output to a file as well as the
 * original destination (e.g., stdout).
 *
 * This mutator operates in several phases. In brief:
 * 1) Attach to the running process and get a handle (BPatch_process
 * object)
 * 2) Get a handle for the parsed image of the mutatee for function
 * lookup (BPatch_image object)
 * 3) Open a file for output
 * 3a) Look up the "open" function
 * 3b) Build a code snippet to call open with the file name.
 * 3c) Run that code snippet via a oneTimeCode, saving the returned
 * file descriptor
 * 4) Write the returned file descriptor into a memory variable for
 * mutatee-side use
 * 5) Build a snippet that copies output to the file
 * 5a) Locate the "write" library call
 * 5b) Access its parameters
 * 5c) Build a snippet calling write(fd, parameters)
 * 5d) Insert the snippet at write
 * 6) Add a hook to exit to ensure that we close the file (using
 * a callback at exit and another oneTimeCode)
 */

void usage() {
  fprintf(stderr, "Usage: retee <process pid> <filename>\n");
  fprintf(stderr, " note: <filename> is relative to the application process.\n");
}

// We need to use a callback, and so the things that callback requires
// are made global - this includes the file descriptor snippet (see below)
BPatch_variableExpr *fdVar = NULL;
// Before we add instrumentation, we need to open the file for
// writing. We can do this with a oneTimeCode - a piece of code run at
// a particular time, rather than at a particular location.
int openFileForWrite(BPatch_process *app, BPatch_image *appImage, char *fileName) {
  // The code to be generated is:
  // fd = open(argv[2], O_WRONLY|O_CREAT, 0666);
  // (1) Find the open function
  BPatch_Vector<BPatch_function *>openFuncs;
  appImage->findFunction("open", openFuncs);
  if (openFuncs.size() == 0) {
    fprintf(stderr, "ERROR: Unable to find function for open()\n");
    return -1;
  }
  // (2) Allocate a vector of snippets for the parameters to open
  BPatch_Vector<BPatch_snippet *> openArgs;
  // (3) Create a string constant expression from argv[3]
  BPatch_constExpr fileNameExpr(fileName);
  // (4) Create two more constant expressions _WRONLY|O_CREAT and 0666
  BPatch_constExpr fileFlagsExpr(O_WRONLY|O_CREAT);
  BPatch_constExpr fileModeExpr(0666);
  // (5) Push 3 & 4 onto the list from step 2, push first to last parameter.
  openArgs.push_back(&fileNameExpr);
  openArgs.push_back(&fileFlagsExpr);
  openArgs.push_back(&fileModeExpr);
  // (6) create a procedure call using function found at 1 and
  // parameters from step 5.
  BPatch_funcCallExpr openCall(*openFuncs[0], openArgs);
  // (7) The oneTimeCode returns whatever the return result from
  // the BPatch_snippet is. In this case, the return result of
  // open -> the file descriptor.
  void *openFD = app->oneTimeCode( openCall );
  // We need to do this to turn the 64-bit upcast of an integer
  // (return type from open) into a 32-bit value. 
  return (int) (long) openFD;
}

// We have used a oneTimeCode to open the file descriptor. However,
// this returns the file descriptor to the mutator - the mutatee has
// no idea what the descriptor is. We need to allocate a variable in
// the mutatee to hold this value for future use and copy the
// (mutator-side) value into the mutatee variable.
// Note: there are alternatives to this technique. We could have
// allocated the variable before the oneTimeCode and augmented the
// snippet to do the assignment. We could also write the file
// descriptor as a constant into any inserted instrumentation.
BPatch_variableExpr *writeFileDescIntoMutatee(BPatch_process *app,
					      BPatch_image *appImage,
					      int fileDescriptor) {
  // (1) Allocate a variable in the mutatee of size (and type) int
  BPatch_variableExpr *fdVar = app->malloc(*appImage->findType("int"));
  if (fdVar == NULL) return NULL;
  // (2) Write the value into the variable
  // Like memcpy, writeValue takes a pointer
  // The third parameter is for functionality called "saveTheWorld",
  // which we don't worry about here (and so is false)
  bool ret = fdVar->writeValue((void *) &fileDescriptor, sizeof(int),
			       false);
  if (ret == false) return NULL;
  return fdVar;
}

// We now have an open file descriptor in the mutatee. We want to
// instrument write to intercept and copy the output. That happens
// here.
bool interceptAndCloneWrite(BPatch_process *app,
			    BPatch_image *appImage,
			    BPatch_variableExpr *fdVar) {
  // (1) Locate the write call
  BPatch_Vector<BPatch_function *> writeFuncs;
  appImage->findFunction("write",
			 writeFuncs);
  if(writeFuncs.size() == 0) {
    fprintf(stderr, "ERROR: Unable to find function for write()\n");
    return false;
  }
  // (2) Build the call to (our) write. Arguments are:
  // ours: fdVar (file descriptor)
  // parameter: buffer
  // parameter: buffer size
  // Declare a vector to hold these.
  BPatch_Vector<BPatch_snippet *> writeArgs;
  // Push on the file descriptor
  writeArgs.push_back(fdVar);
  // Well, we need the buffer... but that's a parameter to the
  // function we're implementing. That's not a problem - we can grab
  // it out with a BPatch_paramExpr.
  BPatch_paramExpr buffer(1); // Second (0, 1, 2) argument
  BPatch_paramExpr bufferSize(2);
  writeArgs.push_back(&buffer);
  writeArgs.push_back(&bufferSize);
  // And build the write call
  BPatch_funcCallExpr writeCall(*writeFuncs[0], writeArgs);
  // (3) Identify the BPatch_point for the entry of write. We're
  // instrumenting the function with itself; normally the findPoint
  // call would operate off a different function than the snippet.
  BPatch_Vector<BPatch_point *> *points;
  points = writeFuncs[0]->findPoint(BPatch_entry);
  if ((*points).size() == 0) {
    return false;
  }
  // (4) Insert the snippet at the start of write
  return app->insertSnippet(writeCall, *points);
  // Note: we have just instrumented write() with a call to
  // write(). This would ordinarily be a _bad thing_, as there is
  // nothing to stop infinite recursion - write -> instrumentation
  // -> write -> instrumentation....
  // However, Dyninst uses a feature called a "tramp guard" to
  // prevent this, and it's on by default.
}

// This function is called as an exit callback (that is, called
// immediately before the process exits when we can still affect it)
// and thus must match the exit callback signature:
//
// typedef void (*BPatchExitCallback) (BPatch_thread *, BPatch_exitType)
//
// Note that the callback gives us a thread, and we want a process - but
// each thread has an up pointer.
void closeFile(BPatch_thread *thread, BPatch_exitType) {
  fprintf(stderr, "Exit callback called for process...\n");
  // (1) Get the BPatch_process and BPatch_images
  BPatch_process *app = thread->getProcess();
  BPatch_image *appImage = app->getImage();
  // The code to be generated is:
  // close(fd);
  // (2) Find close
  BPatch_Vector<BPatch_function *> closeFuncs;
  appImage->findFunction("close", closeFuncs);
  if (closeFuncs.size() == 0) {
    fprintf(stderr, "ERROR: Unable to find function for close()\n");
    return;
  }
  // (3) Allocate a vector of snippets for the parameters to open
  BPatch_Vector<BPatch_snippet *> closeArgs;
  // (4) Add the fd snippet - fdVar is global since we can't
  // get it via the callback
  closeArgs.push_back(fdVar);
  // (5) create a procedure call using function found at 1 and
  // parameters from step 3.
  BPatch_funcCallExpr closeCall(*closeFuncs[0], closeArgs);
  // (6) Use a oneTimeCode to close the file
  app->oneTimeCode( closeCall );
  // (7) Tell the app to continue to finish it off.
  app->continueExecution();
  return;
}

BPatch bpatch;

// In main we perform the following operations.
// 1) Attach to the process and get BPatch_process and BPatch_image
// handles
// 2) Open a file descriptor
// 3) Instrument write
// 4) Continue the process and wait for it to terminate
int main(int argc, char *argv[]) {
  int pid;
  if (argc != 3) {
    usage();
    exit(1);
  }
  pid = atoi(argv[1]);
  // Attach to the program - we can attach with just a pid; the
  // program name is no longer necessary
  fprintf(stderr, "Attaching to process %d...\n", pid);
  BPatch_process *app = bpatch.processAttach(NULL, pid);
  if (!app) return -1;
  // Read the program's image and get an associated image object
  BPatch_image *appImage = app->getImage();
  BPatch_Vector<BPatch_function*> writeFuncs;
  fprintf(stderr, "Opening file %s for write...\n", argv[2]);
  int fileDescriptor = openFileForWrite(app, appImage, argv[2]);
  if (fileDescriptor == -1) {
    fprintf(stderr, "ERROR: opening file %s for write failed\n",
	    argv[2]);
    exit(1);
  }
  fprintf(stderr, "Writing returned file descriptor %d into"
	  "mutatee...\n", fileDescriptor);
  // This was defined globally as the exit callback needs it.
  fdVar = writeFileDescIntoMutatee(app, appImage, fileDescriptor);
  if (fdVar == NULL) {
    fprintf(stderr, "ERROR: failed to write mutatee-side variable\n");
    exit(1);
  }
  fprintf(stderr, "Instrumenting write...\n");
  bool ret = interceptAndCloneWrite(app, appImage, fdVar);
  if (!ret) {
    fprintf(stderr, "ERROR: failed to instrument mutatee\n");
    exit(1);
  }
  fprintf(stderr, "Adding exit callback...\n");
  bpatch.registerExitCallback(closeFile);
  // Continue the execution...
  fprintf(stderr, "Continuing execution and waiting for termination\n");
  app->continueExecution();
  while (!app->isTerminated())
    bpatch.waitForStatusChange();
  printf("Done.\n");
  return 0;
}
