text: /*
    Copyright 2015 Google LLC 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
  
    Unless required by applicable law or agreed to in writing, software
    distributed under the License is distributed on an "AS IS" BASIS,
    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    See the License for the specific language governing permissions and
    limitations under the License.
  */
  
  /*
     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.
  
     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 "afl-ngram.h"
  
  #include <stdio.h>
  #include <stdlib.h>
  #include <unistd.h>
  
  #include "llvm/ADT/Statistic.h"
  #include "llvm/IR/IRBuilder.h"
  #include "llvm/IR/LegacyPassManager.h"
  #include "llvm/IR/Module.h"
  #include "llvm/Support/Debug.h"
  #include "llvm/Support/MathExtras.h"
  #include "llvm/Transforms/IPO/PassManagerBuilder.h"
  
  using namespace llvm;
  
  namespace {
  
  class AFLCoverage : public ModulePass {
  
  public:
    static char ID;
    AFLCoverage() : ModulePass(ID) {}
  
    bool runOnModule(Module &M) override;
  
    // StringRef getPassName() const override {
    //  return "American Fuzzy Lop Instrumentation";
    // }
  };
  
  } // namespace
  
  char AFLCoverage::ID = 0;
  
  bool AFLCoverage::runOnModule(Module &M) {
  
    LLVMContext &C = M.getContext();
  
    IntegerType *Int8Ty = IntegerType::getInt8Ty(C);
    IntegerType *Int32Ty = IntegerType::getInt32Ty(C);
  
    /* Show a banner */
  
    char be_quiet = 0;
  
    if (isatty(2) && !getenv("AFL_QUIET")) {
  
      SAYF(cCYA "afl-llvm-pass " cBRI VERSION cRST
                " by <lszekeres@google.com, adrian.herrera@anu.edu.au>\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)");
    }
  
    /* 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");
  
    /* Decide previous location vector size (must be a power of two) */
  
    char *ngram_size_str = getenv("AFL_NGRAM_SIZE");
    unsigned int ngram_size = 4;
  
    if (ngram_size_str) {
      if (sscanf(ngram_size_str, "%u", &ngram_size) != 1 || ngram_size < 2 ||
          ngram_size > MAX_NGRAM_SIZE) {
        FATAL(
            "Bad value of AFL_NGRAM_SIZE (must be between 2 and MAX_NGRAM_SIZE)");
      }
    }
  
    unsigned PrevLocSize = ngram_size - 1;
    uint64_t PrevLocVecSize = PowerOf2Ceil(PrevLocSize);
    Type *PrevLocTy = VectorType::get(Int32Ty, PrevLocVecSize);
  
    GlobalVariable *AFLPrevLoc = new GlobalVariable(
        M, PrevLocTy, /* isConstant */ false, GlobalValue::ExternalLinkage,
        /* Initializer */ nullptr, "__afl_prev_loc",
        /* InsertBefore */ nullptr, GlobalVariable::GeneralDynamicTLSModel,
        /* AddressSpace */ 0, /* IsExternallyInitialized */ false);
  
    /* Create the vector shuffle mask for updating the previous block history.
       Note that the first element of the vector will store cur_loc, so just set
       it to undef to allow the optimizer to do its thing. */
  
    SmallVector<Constant *, 32> PrevLocShuffle = {UndefValue::get(Int32Ty)};
  
    for (unsigned I = 0; I < PrevLocSize - 1; ++I) {
      PrevLocShuffle.push_back(ConstantInt::get(Int32Ty, I));
    }
  
    for (unsigned I = PrevLocSize; I < PrevLocVecSize; ++I) {
      PrevLocShuffle.push_back(ConstantInt::get(Int32Ty, PrevLocSize));
    }
  
    Constant *PrevLocShuffleMask = ConstantVector::get(PrevLocShuffle);
  
    /* 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 (AFL_R(100) >= inst_ratio)
          continue;
  
        /* Make up cur_loc */
  
        unsigned int cur_loc = AFL_R(MAP_SIZE);
  
        ConstantInt *CurLoc = ConstantInt::get(Int32Ty, cur_loc);
  
        /* Load prev_loc_trans */
  
        LoadInst *PrevLocVec = IRB.CreateLoad(AFLPrevLoc);
        PrevLocVec->setMetadata(M.getMDKindID("nosanitize"),
                                MDNode::get(C, None));
  
        /* "For efficiency, we propose to hash the tuple as a key into the
           hit_count map as (prev_block_trans << 1) ^ curr_block_trans, where
           prev_block_trans = (block_trans_1 ^ ... ^ block_trans_(n-1)" */
  
        Value *PrevLocTrans = IRB.CreateXorReduce(PrevLocVec);
  
        /* Load SHM pointer */
  
        LoadInst *MapPtr = IRB.CreateLoad(AFLMapPtr);
        MapPtr->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
        Value *MapPtrIdx = IRB.CreateGEP(
            MapPtr, IRB.CreateZExt(IRB.CreateXor(PrevLocTrans, CurLoc), Int32Ty));
  
        /* Update bitmap */
  
        LoadInst *Counter = IRB.CreateLoad(MapPtrIdx);
        Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
        Value *Incr = IRB.CreateAdd(Counter, ConstantInt::get(Int8Ty, 1));
        IRB.CreateStore(Incr, MapPtrIdx)
            ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
  
        /* Update prev_loc history vector (by placing cur_loc at the head of the
           vector and shuffle the other elements back by one) */
  
        Value *ShuffledPrevLoc = IRB.CreateShuffleVector(
            PrevLocVec, UndefValue::get(PrevLocTy), PrevLocShuffleMask);
        Value *UpdatedPrevLoc = IRB.CreateInsertElement(
            ShuffledPrevLoc, IRB.CreateLShr(CurLoc, (uint64_t)1), (uint64_t)0);
        StoreInst *Store = IRB.CreateStore(UpdatedPrevLoc, 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_ModuleOptimizerEarly,
                      registerAFLPass);
  
  static RegisterStandardPasses
      RegisterAFLPass0(PassManagerBuilder::EP_EnabledOnOptLevel0,
                       registerAFLPass);
  