AFLplusplus/llvm_mode/afl-llvm-pass.so.cc

/*
   american fuzzy lop - LLVM-mode instrumentation pass
   ---------------------------------------------------

   Written by Laszlo Szekeres <lszekeres@google.com> and
              Michal Zalewski <lcamtuf@google.com>

   LLVM integration design comes from Laszlo Szekeres. C bits copied-and-pasted
   from afl-as.c are Michal's fault.

   Copyright 2015, 2016 Google Inc. All rights reserved.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at:

     http://www.apache.org/licenses/LICENSE-2.0

   This library is plugged into LLVM when invoking clang through afl-clang-fast.
   It tells the compiler to add code roughly equivalent to the bits discussed
   in ../afl-as.h.

 */

#define AFL_LLVM_PASS

#include "../config.h"
#include "../debug.h"

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

#include <list>
#include <string>
#include <fstream>

#include "llvm/IR/DebugInfo.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
#include "llvm/Support/Debug.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include "llvm/IR/CFG.h"

using namespace llvm;

namespace {

  class AFLCoverage : public ModulePass {

    public:

      static char ID;
      AFLCoverage() : ModulePass(ID) {
        char* instWhiteListFilename = getenv("AFL_LLVM_WHITELIST");
        if (instWhiteListFilename) {
          std::string line;
          std::ifstream fileStream;
          fileStream.open(instWhiteListFilename);
          if (!fileStream)
            report_fatal_error("Unable to open AFL_LLVM_WHITELIST");
          getline(fileStream, line);
          while (fileStream) {
            myWhitelist.push_back(line);
            getline(fileStream, line);
          }
        }
      }

      bool runOnModule(Module &M) override;

      // StringRef getPassName() const override {
      //  return "American Fuzzy Lop Instrumentation";
      // }

    protected:

      std::list<std::string> myWhitelist;

  };

}


char AFLCoverage::ID = 0;


bool AFLCoverage::runOnModule(Module &M) {

  LLVMContext &C = M.getContext();

  IntegerType *Int8Ty  = IntegerType::getInt8Ty(C);
  IntegerType *Int32Ty = IntegerType::getInt32Ty(C);
  unsigned int cur_loc = 0;

  /* Show a banner */

  char be_quiet = 0;

  if (isatty(2) && !getenv("AFL_QUIET")) {

    SAYF(cCYA "afl-llvm-pass" VERSION cRST " by <lszekeres@google.com>\n");

  } else be_quiet = 1;

  /* Decide instrumentation ratio */

  char* inst_ratio_str = getenv("AFL_INST_RATIO");
  unsigned int inst_ratio = 100;

  if (inst_ratio_str) {

    if (sscanf(inst_ratio_str, "%u", &inst_ratio) != 1 || !inst_ratio ||
        inst_ratio > 100)
      FATAL("Bad value of AFL_INST_RATIO (must be between 1 and 100)");

  }

  char* neverZero_counters_str = getenv("AFL_NZERO_COUNTS");

  /* Get globals for the SHM region and the previous location. Note that
     __afl_prev_loc is thread-local. */

  GlobalVariable *AFLMapPtr =
      new GlobalVariable(M, PointerType::get(Int8Ty, 0), false,
                         GlobalValue::ExternalLinkage, 0, "__afl_area_ptr");

  GlobalVariable *AFLPrevLoc = new GlobalVariable(
      M, Int32Ty, false, GlobalValue::ExternalLinkage, 0, "__afl_prev_loc",
      0, GlobalVariable::GeneralDynamicTLSModel, 0, false);

  /* Instrument all the things! */

  int inst_blocks = 0;

  for (auto &F : M)
    for (auto &BB : F) {

      BasicBlock::iterator IP = BB.getFirstInsertionPt();
      IRBuilder<> IRB(&(*IP));
      
      if (!myWhitelist.empty()) {
          bool instrumentBlock = false;

          /* Get the current location using debug information.
           * For now, just instrument the block if we are not able
           * to determine our location. */
          DebugLoc Loc = IP->getDebugLoc();
          if ( Loc ) {
              DILocation *cDILoc = dyn_cast<DILocation>(Loc.getAsMDNode());

              unsigned int instLine = cDILoc->getLine();
              StringRef instFilename = cDILoc->getFilename();

              if (instFilename.str().empty()) {
                  /* If the original location is empty, try using the inlined location */
                  DILocation *oDILoc = cDILoc->getInlinedAt();
                  if (oDILoc) {
                      instFilename = oDILoc->getFilename();
                      instLine = oDILoc->getLine();
                  }
              }

              /* Continue only if we know where we actually are */
              if (!instFilename.str().empty()) {
                  for (std::list<std::string>::iterator it = myWhitelist.begin(); it != myWhitelist.end(); ++it) {
                      /* We don't check for filename equality here because
                       * filenames might actually be full paths. Instead we
                       * check that the actual filename ends in the filename
                       * specified in the list. */
                      if (instFilename.str().length() >= it->length()) {
                          if (instFilename.str().compare(instFilename.str().length() - it->length(), it->length(), *it) == 0) {
                              instrumentBlock = true;
                              break;
                          }
                      }
                  }
              }
          }

          /* Either we couldn't figure out our location or the location is
           * not whitelisted, so we skip instrumentation. */
          if (!instrumentBlock) continue;
      }


      if (AFL_R(100) >= inst_ratio) continue;

      /* Make up cur_loc */

       //cur_loc++;
      cur_loc = AFL_R(MAP_SIZE);

      // only instrument if this basic block is the destination of a previous
      // basic block that has multiple successors
      // this gets rid of ~5-10% of instrumentations that are unnecessary
      // result: a little more speed and less map pollution
      int more_than_one = -1;
      //fprintf(stderr, "BB %u: ", cur_loc);
      for (BasicBlock *Pred : predecessors(&BB)) {
        int count = 0;
        if (more_than_one == -1)
          more_than_one = 0;
        //fprintf(stderr, " %p=>", Pred);
        for (BasicBlock *Succ : successors(Pred)) {
          //if (count > 0)
          //  fprintf(stderr, "|");
          if (Succ != NULL) count++;
          //fprintf(stderr, "%p", Succ);
        }
        if (count > 1)
          more_than_one = 1;
      }
      //fprintf(stderr, " == %d\n", more_than_one);
      if (more_than_one != 1)
        continue;

      ConstantInt *CurLoc = ConstantInt::get(Int32Ty, cur_loc);

      /* Load prev_loc */

      LoadInst *PrevLoc = IRB.CreateLoad(AFLPrevLoc);
      PrevLoc->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
      Value *PrevLocCasted = IRB.CreateZExt(PrevLoc, IRB.getInt32Ty());

      /* Load SHM pointer */

      LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr);
      MapPtr->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
      Value *MapPtrIdx = IRB.CreateGEP(MapPtr, IRB.CreateXor(PrevLocCasted, CurLoc));

      /* Update bitmap */

      LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
      Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));

      Value *Incr;
      if (neverZero_counters_str == NULL || neverZero_counters_str[0] != '4')
        Incr = IRB.CreateAdd(Counter, ConstantInt::get(Int8Ty, 1));

      if (neverZero_counters_str != NULL) {
          /* hexcoder: Realize a counter that skips zero during overflow.
           * Once this counter reaches its maximum value, it next increments to 1
           *
           * Instead of
           * Counter + 1 -> Counter
           * we inject now this
           * Counter + 1 -> {Counter, OverflowFlag}
           * Counter + OverflowFlag -> Counter
           */
           
          // Solution #1 - creates
          //mov    ecx,edx
          //add    cl,0x1
          //adc    dl,0x1
          if (neverZero_counters_str[0] == '1') {
            CallInst *AddOv = IRB.CreateBinaryIntrinsic(Intrinsic::uadd_with_overflow, Counter, ConstantInt::get(Int8Ty, 1));
            AddOv->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
            Value *SumWithOverflowBit = AddOv;
            Incr = IRB.CreateAdd(IRB.CreateExtractValue(SumWithOverflowBit, 0),  // sum 
                                 IRB.CreateZExt( // convert from one bit type to 8 bits type 
                                                IRB.CreateExtractValue(SumWithOverflowBit, 1), // overflow
                                                Int8Ty));

          // Solution #2 - creates the same code as #1
          } else if (neverZero_counters_str[0] == '2') {
            auto cf = IRB.CreateICmpULT(Incr, ConstantInt::get(Int8Ty, 1));
            Incr = IRB.CreateAdd(Incr, cf);
           
          // Solution #3 - creates
          //add    cl,0x1
          //cmp    cl,0x1
          //adc    cl,0x0
          } else if (neverZero_counters_str[0] == '3') {
            auto cf = IRB.CreateICmpEQ(Incr, ConstantInt::get(Int8Ty, 0));
            Incr = IRB.CreateAdd(Incr, cf);

          // Solution #4 - creates
          // cmp    dl, $0xff
          // sete   cl
          // add    dl,cl
          // add    dl,0x01
          } else if (neverZero_counters_str[0] == '4') {
             auto cf = IRB.CreateICmpEQ(Counter, ConstantInt::get(Int8Ty, 255));
             Value *HowMuch = IRB.CreateAdd(ConstantInt::get(Int8Ty, 1), cf);
             Incr = IRB.CreateAdd(Counter, HowMuch);

          } else if (neverZero_counters_str[0] == '5') {
            auto cf = IRB.CreateICmpEQ(Incr, ConstantInt::get(Int8Ty, 0));
            auto carry = IRB.CreateZExt(cf, Int8Ty);
            Incr = IRB.CreateAdd(Incr, carry);

          } else if (neverZero_counters_str[0] == '6') {
            auto cf = IRB.CreateICmpULT(Incr, ConstantInt::get(Int8Ty, 1));
            auto carry = IRB.CreateZExt(cf, Int8Ty);
            Incr = IRB.CreateAdd(Incr, carry);
           
          // no other implementations yet
          } else {
            fprintf(stderr, "Error: unknown value for AFL_NZERO_COUNTS: %s (valid is 1-4)\n", neverZero_counters_str);
            exit(-1);
          }
      }

      IRB.CreateStore(Incr, MapPtrIdx)->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));

      /* Set prev_loc to cur_loc >> 1 */

      StoreInst *Store = IRB.CreateStore(ConstantInt::get(Int32Ty, cur_loc >> 1), AFLPrevLoc);
      Store->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));

      inst_blocks++;

    }

  /* Say something nice. */

  if (!be_quiet) {

    if (!inst_blocks) WARNF("No instrumentation targets found.");
    else OKF("Instrumented %u locations (%s mode, ratio %u%%).",
             inst_blocks, getenv("AFL_HARDEN") ? "hardened" :
             ((getenv("AFL_USE_ASAN") || getenv("AFL_USE_MSAN")) ?
              "ASAN/MSAN" : "non-hardened"), inst_ratio);

  }

  return true;

}


static void registerAFLPass(const PassManagerBuilder &,
                            legacy::PassManagerBase &PM) {

  PM.add(new AFLCoverage());

}


static RegisterStandardPasses RegisterAFLPass(
    PassManagerBuilder::EP_OptimizerLast, registerAFLPass);

static RegisterStandardPasses RegisterAFLPass0(
    PassManagerBuilder::EP_EnabledOnOptLevel0, registerAFLPass);