<commit_before>#include <linux/bpf.h>

#include <clang/AST/ASTConsumer.h>
#include <clang/AST/ASTContext.h>
#include <clang/AST/RecordLayout.h>
#include <clang/Frontend/CompilerInstance.h>
#include <clang/Rewrite/Core/Rewriter.h>

#include "b_frontend_action.h"

extern "C"
int bpf_create_map(enum bpf_map_type map_type, int key_size, int value_size,
                   int max_entries);

namespace ebpf {
using std::map;
using std::string;
using std::to_string;
using std::unique_ptr;
using namespace clang;

BTypeVisitor::BTypeVisitor(ASTContext &C, Rewriter &rewriter, map<string, BPFTable> &tables)
    : C(C), rewriter_(rewriter), out_(llvm::errs()), tables_(tables) {
}

bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
  return true;
}

// convert calls of the type:
//  table.foo(&key)
// to:
//  bpf_table_foo_elem(bpf_pseudo_fd(table), &key [,&leaf])
bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
  // make sure node is a reference to a bpf table, which is assured by the
  // presence of the section("maps/<typename>") GNU __attribute__
  if (MemberExpr *Memb = dyn_cast<MemberExpr>(Call->getCallee()->IgnoreImplicit())) {
    StringRef memb_name = Memb->getMemberDecl()->getName();
    if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Memb->getBase())) {
      if (SectionAttr *A = Ref->getDecl()->getAttr<SectionAttr>()) {
        if (!A->getName().startswith("maps"))
          return true;

        SourceRange argRange(Call->getArg(0)->getLocStart(),
                             Call->getArg(Call->getNumArgs()-1)->getLocEnd());
        string args = rewriter_.getRewrittenText(argRange);

        // find the table fd, which was opened at declaration time
        auto table_it = tables_.find(Ref->getDecl()->getName());
        if (table_it == tables_.end()) {
          C.getDiagnostics().Report(Ref->getLocEnd(), diag::err_expected)
              << "initialized handle for bpf_table";
          return false;
        }
        string fd = to_string(table_it->second.fd);
        string prefix, suffix;
        string map_update_policy = "BPF_ANY";
        string txt;
        if (memb_name == "lookup_or_init") {
          string map_update_policy = "BPF_NOEXIST";
          string name = Ref->getDecl()->getName();
          string arg0 = rewriter_.getRewrittenText(SourceRange(Call->getArg(0)->getLocStart(),
                                                               Call->getArg(0)->getLocEnd()));
          string arg1 = rewriter_.getRewrittenText(SourceRange(Call->getArg(1)->getLocStart(),
                                                               Call->getArg(1)->getLocEnd()));
          string lookup = "bpf_map_lookup_elem_(bpf_pseudo_fd(1, " + fd + ")";
          string update = "bpf_map_update_elem_(bpf_pseudo_fd(1, " + fd + ")";
          txt  = "({typeof(" + name + ".leaf) *leaf = " + lookup + ", " + arg0 + "); ";
          txt += "if (!leaf) {";
          txt += " " + update + ", " + arg0 + ", " + arg1 + ", " + map_update_policy + ");";
          txt += " leaf = " + lookup + ", " + arg0 + ");";
          txt += " if (!leaf) return -1;";
          txt += "}";
          txt += "leaf;})";
        } else {
          if (memb_name == "lookup") {
            prefix = "bpf_map_lookup_elem";
            suffix = ")";
          } else if (memb_name == "update") {
            prefix = "bpf_map_update_elem";
            suffix = ", " + map_update_policy + ")";
          } else if (memb_name == "delete") {
            prefix = "bpf_map_delete_elem";
            suffix = ")";
          } else if (memb_name == "call") {
            prefix = "bpf_tail_call_";
            suffix = ")";
          } else {
            llvm::errs() << "error: unknown bpf_table operation " << memb_name << "\n";
            return false;
          }
          prefix += "((void *)bpf_pseudo_fd(1, " + fd + "), ";

          txt = prefix + args + suffix;
        }
        rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), txt);
        return true;
      }
    }
  }
  return true;
}

// look for table subscript references, and turn them into auto table entries:
// table.data[key]
//  becomes:
// struct Key key = {123};
// struct Leaf *leaf = table.get(&key);
// if (!leaf) {
//   struct Leaf zleaf = {0};
//   table.put(&key, &zleaf, BPF_NOEXIST);
//   leaf = table.get(&key);
//   if (!leaf) return -1;
// }
bool BTypeVisitor::VisitArraySubscriptExpr(ArraySubscriptExpr *Arr) {
  Expr *LHS = Arr->getLHS()->IgnoreImplicit();
  Expr *RHS = Arr->getRHS()->IgnoreImplicit();
  if (MemberExpr *Memb = dyn_cast<MemberExpr>(LHS)) {
    if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Memb->getBase())) {
      if (SectionAttr *A = Ref->getDecl()->getAttr<SectionAttr>()) {
        if (A->getName().startswith("maps")) {
          auto table_it = tables_.find(Ref->getDecl()->getName());
          if (table_it == tables_.end()) {
            C.getDiagnostics().Report(Ref->getLocEnd(), diag::err_expected)
                << "initialized handle for bpf_table";
            return false;
          }
          string fd = to_string(table_it->second.fd);
          string map_update_policy = "BPF_NOEXIST";
          string name = Ref->getDecl()->getName();
          SourceRange argRange(RHS->getLocStart(), RHS->getLocEnd());
          string args = rewriter_.getRewrittenText(argRange);
          string lookup = "bpf_map_lookup_elem_(bpf_pseudo_fd(1, " + fd + ")";
          string update = "bpf_map_update_elem_(bpf_pseudo_fd(1, " + fd + ")";
          string txt = "(*({typeof(" + name + ".leaf) *leaf = " + lookup + ", " + args + "); ";
          txt         += "if (!leaf) {";
          txt         += " typeof(" + name + ".leaf) zleaf = {0};";
          txt         += " " + update + ", " + args + ", &zleaf, " + map_update_policy + ");";
          txt         += " leaf = " + lookup + ", " + args + ");";
          txt         += " if (!leaf) return -1;";
          txt         += "}";
          txt         += "leaf;}))";
          rewriter_.ReplaceText(SourceRange(Arr->getLocStart(), Arr->getLocEnd()), txt);
        }
      }
    }
  }
  return true;
}

bool BTypeVisitor::VisitMemberExpr(MemberExpr *E) {
  Expr *Base = E->getBase()->IgnoreImplicit();
  if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Base)) {
    if (DeprecatedAttr *A = Ref->getDecl()->getAttr<DeprecatedAttr>()) {
      if (A->getMessage() == "packet") {
        if (FieldDecl *F = dyn_cast<FieldDecl>(E->getMemberDecl())) {<commit_msg>Return 0 in lookup_or_init error path<commit_after>#include <linux/bpf.h>

#include <clang/AST/ASTConsumer.h>
#include <clang/AST/ASTContext.h>
#include <clang/AST/RecordLayout.h>
#include <clang/Frontend/CompilerInstance.h>
#include <clang/Rewrite/Core/Rewriter.h>

#include "b_frontend_action.h"

extern "C"
int bpf_create_map(enum bpf_map_type map_type, int key_size, int value_size,
                   int max_entries);

namespace ebpf {
using std::map;
using std::string;
using std::to_string;
using std::unique_ptr;
using namespace clang;

BTypeVisitor::BTypeVisitor(ASTContext &C, Rewriter &rewriter, map<string, BPFTable> &tables)
    : C(C), rewriter_(rewriter), out_(llvm::errs()), tables_(tables) {
}

bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
  return true;
}

// convert calls of the type:
//  table.foo(&key)
// to:
//  bpf_table_foo_elem(bpf_pseudo_fd(table), &key [,&leaf])
bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
  // make sure node is a reference to a bpf table, which is assured by the
  // presence of the section("maps/<typename>") GNU __attribute__
  if (MemberExpr *Memb = dyn_cast<MemberExpr>(Call->getCallee()->IgnoreImplicit())) {
    StringRef memb_name = Memb->getMemberDecl()->getName();
    if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Memb->getBase())) {
      if (SectionAttr *A = Ref->getDecl()->getAttr<SectionAttr>()) {
        if (!A->getName().startswith("maps"))
          return true;

        SourceRange argRange(Call->getArg(0)->getLocStart(),
                             Call->getArg(Call->getNumArgs()-1)->getLocEnd());
        string args = rewriter_.getRewrittenText(argRange);

        // find the table fd, which was opened at declaration time
        auto table_it = tables_.find(Ref->getDecl()->getName());
        if (table_it == tables_.end()) {
          C.getDiagnostics().Report(Ref->getLocEnd(), diag::err_expected)
              << "initialized handle for bpf_table";
          return false;
        }
        string fd = to_string(table_it->second.fd);
        string prefix, suffix;
        string map_update_policy = "BPF_ANY";
        string txt;
        if (memb_name == "lookup_or_init") {
          string map_update_policy = "BPF_NOEXIST";
          string name = Ref->getDecl()->getName();
          string arg0 = rewriter_.getRewrittenText(SourceRange(Call->getArg(0)->getLocStart(),
                                                               Call->getArg(0)->getLocEnd()));
          string arg1 = rewriter_.getRewrittenText(SourceRange(Call->getArg(1)->getLocStart(),
                                                               Call->getArg(1)->getLocEnd()));
          string lookup = "bpf_map_lookup_elem_(bpf_pseudo_fd(1, " + fd + ")";
          string update = "bpf_map_update_elem_(bpf_pseudo_fd(1, " + fd + ")";
          txt  = "({typeof(" + name + ".leaf) *leaf = " + lookup + ", " + arg0 + "); ";
          txt += "if (!leaf) {";
          txt += " " + update + ", " + arg0 + ", " + arg1 + ", " + map_update_policy + ");";
          txt += " leaf = " + lookup + ", " + arg0 + ");";
          txt += " if (!leaf) return 0;";
          txt += "}";
          txt += "leaf;})";
        } else {
          if (memb_name == "lookup") {
            prefix = "bpf_map_lookup_elem";
            suffix = ")";
          } else if (memb_name == "update") {
            prefix = "bpf_map_update_elem";
            suffix = ", " + map_update_policy + ")";
          } else if (memb_name == "delete") {
            prefix = "bpf_map_delete_elem";
            suffix = ")";
          } else if (memb_name == "call") {
            prefix = "bpf_tail_call_";
            suffix = ")";
          } else {
            llvm::errs() << "error: unknown bpf_table operation " << memb_name << "\n";
            return false;
          }
          prefix += "((void *)bpf_pseudo_fd(1, " + fd + "), ";

          txt = prefix + args + suffix;
        }
        rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), txt);
        return true;
      }
    }
  }
  return true;
}

// look for table subscript references, and turn them into auto table entries:
// table.data[key]
//  becomes:
// struct Key key = {123};
// struct Leaf *leaf = table.get(&key);
// if (!leaf) {
//   struct Leaf zleaf = {0};
//   table.put(&key, &zleaf, BPF_NOEXIST);
//   leaf = table.get(&key);
//   if (!leaf) return -1;
// }
bool BTypeVisitor::VisitArraySubscriptExpr(ArraySubscriptExpr *Arr) {
  Expr *LHS = Arr->getLHS()->IgnoreImplicit();
  Expr *RHS = Arr->getRHS()->IgnoreImplicit();
  if (MemberExpr *Memb = dyn_cast<MemberExpr>(LHS)) {
    if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Memb->getBase())) {
      if (SectionAttr *A = Ref->getDecl()->getAttr<SectionAttr>()) {
        if (A->getName().startswith("maps")) {
          auto table_it = tables_.find(Ref->getDecl()->getName());
          if (table_it == tables_.end()) {
            C.getDiagnostics().Report(Ref->getLocEnd(), diag::err_expected)
                << "initialized handle for bpf_table";
            return false;
          }
          string fd = to_string(table_it->second.fd);
          string map_update_policy = "BPF_NOEXIST";
          string name = Ref->getDecl()->getName();
          SourceRange argRange(RHS->getLocStart(), RHS->getLocEnd());
          string args = rewriter_.getRewrittenText(argRange);
          string lookup = "bpf_map_lookup_elem_(bpf_pseudo_fd(1, " + fd + ")";
          string update = "bpf_map_update_elem_(bpf_pseudo_fd(1, " + fd + ")";
          string txt = "(*({typeof(" + name + ".leaf) *leaf = " + lookup + ", " + args + "); ";
          txt         += "if (!leaf) {";
          txt         += " typeof(" + name + ".leaf) zleaf = {0};";
          txt         += " " + update + ", " + args + ", &zleaf, " + map_update_policy + ");";
          txt         += " leaf = " + lookup + ", " + args + ");";
          txt         += " if (!leaf) return 0;";
          txt         += "}";
          txt         += "leaf;}))";
          rewriter_.ReplaceText(SourceRange(Arr->getLocStart(), Arr->getLocEnd()), txt);
        }
      }
    }
  }
  return true;
}

bool BTypeVisitor::VisitMemberExpr(MemberExpr *E) {
  Expr *Base = E->getBase()->IgnoreImplicit();
  if (DeclRefExpr *Ref = dyn_cast<DeclRefExpr>(Base)) {
    if (DeprecatedAttr *A = Ref->getDecl()->getAttr<DeprecatedAttr>()) {
      if (A->getMessage() == "packet") {
        if (FieldDecl *F = dyn_cast<FieldDecl>(E->getMemberDecl())) {
{'fraction_non_alphanumeric': 0.1674117465700632, 'fraction_numerical': 0.00393093880067828, 'mean_word_length': 2.4108832807570977, 'pattern_counts': {'":': 0, '<': 57, '<?xml version=': 0, '>': 143, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '896284', 'n_tokens_mistral': 4131, 'n_tokens_neox': 4073, 'n_words': 829}