lifting-bits-remill/tests/AArch64/Lift.cpp

/*
 * Copyright (c) 2017 Trail of Bits, Inc.
 *
 * 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.
 */

#include <gflags/gflags.h>
#include <glog/logging.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/GlobalValue.h>
#include <llvm/IR/Instructions.h>
#include <llvm/IR/LLVMContext.h>
#include <llvm/IR/Module.h>
#include <llvm/IR/Type.h>

#include <algorithm>
#include <cstdint>
#include <fstream>
#include <map>
#include <memory>
#include <sstream>
#include <string>

#include "remill/Arch/Arch.h"
#include "remill/Arch/Instruction.h"
#include "remill/Arch/Name.h"
#include "remill/BC/IntrinsicTable.h"
#include "remill/BC/Lifter.h"
#include "remill/BC/Util.h"
#include "remill/OS/OS.h"
#include "tests/AArch64/Test.h"

#ifdef __APPLE__
#  define SYMBOL_PREFIX "_"
#else
#  define SYMBOL_PREFIX ""
#endif

DEFINE_string(bc_out, "",
              "Name of the file in which to place the generated bitcode.");

DECLARE_string(arch);
DECLARE_string(os);

extern "C" {
int gNativeState [[gnu::used]] = 0;
int gLiftedState [[gnu::used]] = 0;
}  // extern

namespace {

// Decode a test and add it as a basic block to the module.
//
// TODO(pag): Eventually handle control-flow.
static void AddFunctionToModule(llvm::Module *module, const remill::Arch *arch,
                                const test::TestInfo &test) {
  DLOG(INFO) << "Adding block for: " << test.test_name;

  std::stringstream ss;
  ss << SYMBOL_PREFIX << test.test_name << "_lifted";

  auto func = remill::DeclareLiftedFunction(module, ss.str());
  remill::CloneBlockFunctionInto(func);

  func->setLinkage(llvm::GlobalValue::ExternalLinkage);
  func->setVisibility(llvm::GlobalValue::DefaultVisibility);

  remill::IntrinsicTable intrinsics(module);
  remill::InstructionLifter lifter(arch, &intrinsics);

  std::map<uint64_t, remill::Instruction> inst;
  std::map<uint64_t, llvm::BasicBlock *> blocks;

  // Function that will create basic blocks as needed.
  auto GetOrCreateBlock = [func, &blocks](uint64_t block_pc) {
    auto &block = blocks[block_pc];
    if (!block) {
      block = llvm::BasicBlock::Create(func->getContext(), "", func);
    }
    return block;
  };

  std::stringstream seen_insts;
  const char *sep = "";

  auto entry_block = GetOrCreateBlock(test.test_begin);
  llvm::BranchInst::Create(entry_block, &(func->front()));

  auto saw_isel = false;
  auto addr = test.test_begin;
  while (addr < test.test_end) {
    std::string inst_bytes;
    auto bytes = reinterpret_cast<const char *>(addr);
    inst_bytes.insert(inst_bytes.end(), bytes, bytes + 4);

    remill::Instruction inst;
    CHECK(arch->DecodeInstruction(addr, inst_bytes, inst))
        << "Can't decode test instruction " << inst.Serialize() << " in "
        << test.test_name;

    seen_insts << sep << inst.function;
    sep = ", ";

    LOG(INFO) << "Lifting " << inst.Serialize();

    auto block = GetOrCreateBlock(inst.pc);
    CHECK(remill::kLiftedInstruction == lifter.LiftIntoBlock(inst, block))
        << "Can't lift test instruction " << inst.Serialize() << " in "
        << test.test_name;

    saw_isel = saw_isel || inst.function == test.isel_name;
    addr += inst.NumBytes();

    // Connect together the basic blocks.
    switch (inst.category) {
      case remill::Instruction::kCategoryNormal:
      case remill::Instruction::kCategoryNoOp:
        llvm::BranchInst::Create(GetOrCreateBlock(inst.next_pc), block);
        break;

      case remill::Instruction::kCategoryDirectJump:
      case remill::Instruction::kCategoryDirectFunctionCall:
        llvm::BranchInst::Create(GetOrCreateBlock(inst.branch_taken_pc), block);
        break;

      case remill::Instruction::kCategoryConditionalBranch:
        llvm::BranchInst::Create(GetOrCreateBlock(inst.branch_taken_pc),
                                 GetOrCreateBlock(inst.branch_not_taken_pc),
                                 remill::LoadBranchTaken(block), block);
        break;

      default:
        remill::AddTerminatingTailCall(block, intrinsics.missing_block);
        break;
    }
  }

  CHECK(saw_isel) << "Test " << test.test_name
                  << " does not have an instruction that "
                  << "uses the semantics function " << test.isel_name
                  << ", saw " << seen_insts.str();

  // Terminate any stragglers.
  for (auto pc_to_block : blocks) {
    auto block = pc_to_block.second;
    if (!block->getTerminator()) {
      remill::AddTerminatingTailCall(block, intrinsics.missing_block);
    }
  }
}

}  // namespace

extern "C" int main(int argc, char *argv[]) {

  google::ParseCommandLineFlags(&argc, &argv, true);
  google::InitGoogleLogging(argv[0]);

  auto os = remill::GetOSName(REMILL_OS);
  auto context = new llvm::LLVMContext;

  auto arch = remill::Arch::Get(*context, os, remill::kArchAArch64LittleEndian);

  DLOG(INFO) << "Generating tests.";

  auto bc_file = remill::FindSemanticsBitcodeFile(FLAGS_arch);
  auto module = remill::LoadModuleFromFile(context, bc_file);
  remill::GetHostArch(*context)->PrepareModule(module.get());

  for (auto i = 0U;; ++i) {
    const auto &test = test::__aarch64_test_table_begin[i];
    if (&test >= &(test::__aarch64_test_table_end[0]))
      break;
    AddFunctionToModule(module.get(), arch, test);
  }

  DLOG(INFO) << "Serializing bitcode to " << FLAGS_bc_out;
  remill::StoreModuleToFile(module, FLAGS_bc_out);

  DLOG(INFO) << "Done.";
  return 0;
}