blog.cpp

// blog.cpp --- Simple example demonstrating GTIRB usage
//
// Search a GTIRB instance printing the calls preceeding `system' calls.
//
// Compiled with (with system installs):
//     g++ --std=c++17 -lgtirb blog.cpp
//
// or with (with explicit includes):
//     g++ -I/gtirb/build/include/ -I/gtirb/build/protobuf-src/include/ \
//         --std=c++17 -lgtirb blog.cpp
//
#include <gtirb/gtirb.hpp>
#include <fstream>
#include <iomanip>
#include <set>
#include <vector>
using namespace gtirb;

// Create a context to manage memory for gtirb objects
Context context;

// Print Addrs in hex format
std::ostream& operator<<(std::ostream& Os, Addr A) {
  auto Flags = Os.flags();
  Os << "0x" << std::hex << std::setw(8) << std::setfill('0') << uint64_t(A);
  Os.flags(Flags);
  return Os;
}

template <class T> T* nodeFromUUID(Context& C, UUID id) {
  return dyn_cast_or_null<T>(Node::getByUUID(C, id));
}

// Return the string name of the plt function at addr.
std::optional<Symbol> forwardSymbol(const std::map<UUID, UUID>* forwarding,
                                    Symbol* symbol) {
  auto found = forwarding->find(symbol->getUUID());
  if (found != forwarding->end()) {
    Symbol* destSymbol = nodeFromUUID<Symbol>(context, found->second);
    if (destSymbol)
      return *destSymbol;
    else
      return std::nullopt;
  } else {
    return std::nullopt;
  }
}

// Walk backwards in a graph printing all symbols
class BackwardsScan {
public:
  using Vertex = CFG::vertex_descriptor;

  BackwardsScan(Module& M, std::map<UUID, UUID>& fwd, CFG& G, int depth)
      : Graph(G), M(M), fwd(fwd), Depth(depth) {}

  void visit(Vertex V, int length) {
    if ((length >= Depth))
      return;

    length += 1;

    // Process this vertex (if not previously visited).
    if (Visited.find(V) == Visited.end()) {
      auto block = dyn_cast<Block>(Graph[V]);
      auto BStart = block->getAddress();
      std::cout << "Predecessor " << BStart << " references ";
      auto BEnd = BStart + block->getSize();
      auto E = M.findSymbolicExpression(BStart, BEnd);
      std::find_if(E.begin(), E.end(), [this](auto& arg) {
        return std::visit(
            [this](auto&& arg) {
              using T = std::decay_t<decltype(arg)>;
              if constexpr (std::is_same_v<T, SymAddrConst>) {
                auto fwded = forwardSymbol(&fwd, arg.Sym);
                if (fwded)
                  std::cout << fwded->getName();
                else
                  std::cout << arg.Sym->getName();
                return true;
              } else {
                return false;
              }
            },
            arg);
      });
      std::cout << std::endl;
    }

    // Continue along all incoming edges.
    Visited.insert(V);
    for (auto [edge, end] = in_edges(V, Graph); edge != end; ++edge) {
      visit(source(*edge, Graph), length);
    }
  }

  CFG& Graph;
  Module& M;
  std::map<UUID, UUID> fwd;
  int Depth;
  std::set<Vertex> Visited;
};

int main(int argc, char** argv) {
  if (argc < 3) {
    std::cerr << "Usage " << argv[0] << " GTIRB-FILE PATH-LENGTH" << std::endl;
    exit(1);
  }
  std::ifstream in(argv[1]);
  int depth = std::stoi(argv[2]);

  // Get the module
  auto I = IR::load(context, in);
  auto& M = *I->begin();
  auto* fwd = I->modules().begin()->getAuxData<std::map<UUID, UUID>>(
      "symbolForwarding");
  if (!fwd) {
    std::cerr << "No symbolForwarding table." << std::endl;
    exit(1);
  }

  // Find the block calling system.  This constitutes the back door.
  auto& Cfg = M.getCFG();
  auto Blocks = blocks(Cfg);
  auto SystemBlockIt =
      std::find_if(Blocks.begin(), Blocks.end(), [&M, fwd](auto& B) {
        auto BStart = B.getAddress();
        auto BEnd = BStart + B.getSize();
        auto E = M.findSymbolicExpression(BStart, BEnd);
        return E.end() != std::find_if(E.begin(), E.end(), [fwd](auto& arg) {
                 return std::visit(
                     [fwd](auto&& arg) {
                       using T = std::decay_t<decltype(arg)>;
                       if constexpr (std::is_same_v<T, SymAddrConst>) {
                         auto fwded = forwardSymbol(fwd, arg.Sym);
                         auto s = arg.Sym->getName();
                         if (fwded)
                           s = fwded->getName();
                         return s == "system";
                       } else {
                         return false;
                       }
                     },
                     arg);
               });
      });
  if (SystemBlockIt == Blocks.end()) {
    std::cerr << "No block referencing \"system\" was found." << std::endl;
    exit(1);
  }

  auto BackDoorBlock = &*SystemBlockIt;
  std::cout << "Basic Block calling system found at "
            << BackDoorBlock->getAddress() << std::endl;

  // Walk backwards in the CFG printing called functions.
  BackwardsScan(M, *fwd, Cfg, depth).visit(*getVertex(BackDoorBlock, Cfg), 0);
}