Merge pull request iovisor#1792 from pchaigno/refactor-ext-ptr-assign…

…ments

Refactor external pointer assignments

src/cc/frontends/clang/b_frontend_action.cc

@@ -214,11 +214,56 @@ ProbeVisitor::ProbeVisitor(ASTContext &C, Rewriter &rewriter,
  set<Decl *> &m, bool track_helpers) :
  C(C), rewriter_(rewriter), m_(m), track_helpers_(track_helpers) {}
 
+bool ProbeVisitor::assignsExtPtr(Expr *E, int *nbAddrOf) {
+ if (IsContextMemberExpr(E)) {
+ *nbAddrOf = 0;
+ return true;
+ }
+
+ ProbeChecker checker = ProbeChecker(E, ptregs_, track_helpers_,
+ true);
+ if (checker.is_transitive()) {
+ // The negative of the number of dereferences is the number of addrof. In
+ // an assignment, if we went through n addrof before getting the external
+ // pointer, then we'll need n dereferences on the left-hand side variable
+ // to get to the external pointer.
+ *nbAddrOf = -checker.get_nb_derefs();
+ return true;
+ }
+
+ if (E->getStmtClass() == Stmt::CallExprClass) {
+ CallExpr *Call = dyn_cast<CallExpr>(E);
+ 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 false;
+
+ if (memb_name == "lookup" || memb_name == "lookup_or_init") {
+ if (m_.find(Ref->getDecl()) != m_.end()) {
+ // Retrieved an ext. pointer from a map, mark LHS as ext. pointer.
+ // Pointers from maps always need a single dereference to get the
+ // actual value. The value may be an external pointer but cannot
+ // be a pointer to an external pointer as the verifier prohibits
+ // storing known pointers (to map values, context, the stack, or
+ // the packet) in maps.
+ *nbAddrOf = 1;
+ return true;
+ }
+ }
+ }
+ }
+ }
+ }
+ return false;
+}
 bool ProbeVisitor::VisitVarDecl(VarDecl *D) {
  if (Expr *E = D->getInit()) {
- ProbeChecker checker = ProbeChecker(E, ptregs_, track_helpers_, true);
- if (checker.is_transitive() || IsContextMemberExpr(E)) {
- tuple<Decl *, int> pt = make_tuple(D, checker.get_nb_derefs());
+ int nbAddrOf;
+ if (assignsExtPtr(E, &nbAddrOf)) {
+ // The negative of the number of addrof is the number of dereferences.
+ tuple<Decl *, int> pt = make_tuple(D, -nbAddrOf);
  set_ptreg(pt);
  }
  }
@@ -249,42 +294,11 @@ bool ProbeVisitor::VisitCallExpr(CallExpr *Call) {
 bool ProbeVisitor::VisitBinaryOperator(BinaryOperator *E) {
  if (!E->isAssignmentOp())
  return true;
- // copy probe attribute from RHS to LHS if present
- ProbeChecker checker = ProbeChecker(E->getRHS(), ptregs_, track_helpers_,
- true);
- if (checker.is_transitive()) {
- // The negative of the number of dereferences is the number of addrof. In
- // an assignment, if we went through n addrof before getting the external
- // pointer, then we'll need n dereferences on the left-hand side variable
- // to get to the external pointer.
- ProbeSetter setter(&ptregs_, -checker.get_nb_derefs());
- setter.TraverseStmt(E->getLHS());
- } else if (E->getRHS()->getStmtClass() == Stmt::CallExprClass) {
- CallExpr *Call = dyn_cast<CallExpr>(E->getRHS());
- 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;
 
- if (memb_name == "lookup" || memb_name == "lookup_or_init") {
- if (m_.find(Ref->getDecl()) != m_.end()) {
- // Retrieved an ext. pointer from a map, mark LHS as ext. pointer.
- // Pointers from maps always need a single dereference to get the
- // actual value. The value may be an external pointer but cannot
- // be a pointer to an external pointer as the verifier prohibits
- // storing known pointers (to map values, context, the stack, or
- // the packet) in maps.
- ProbeSetter setter(&ptregs_, 1);
- setter.TraverseStmt(E->getLHS());
- }
- }
- }
- }
- }
- } else if (IsContextMemberExpr(E->getRHS())) {
- ProbeSetter setter(&ptregs_);
+ // copy probe attribute from RHS to LHS if present
+ int nbAddrOf;
+ if (assignsExtPtr(E->getRHS(), &nbAddrOf)) {
+ ProbeSetter setter(&ptregs_, nbAddrOf);
  setter.TraverseStmt(E->getLHS());
  }
  return true;
@@ -355,7 +369,6 @@ bool ProbeVisitor::IsContextMemberExpr(Expr *E) {
  bool found = false;
  MemberExpr *M;
  for (M = Memb; M; M = dyn_cast<MemberExpr>(M->getBase())) {
- memb_visited_.insert(M);
  rhs_start = M->getLocEnd();
  base = M->getBase();
  member = M->getMemberLoc();

src/cc/frontends/clang/b_frontend_action.h

@@ -99,6 +99,7 @@ class ProbeVisitor : public clang::RecursiveASTVisitor<ProbeVisitor> {
  void set_ctx(clang::Decl *D) { ctx_ = D; }
  std::set<std::tuple<clang::Decl *, int>> get_ptregs() { return ptregs_; }
  private:
+ bool assignsExtPtr(clang::Expr *E, int *nbAddrOf);
  bool IsContextMemberExpr(clang::Expr *E);
  clang::SourceRange expansionRange(clang::SourceRange range);
  template <unsigned N>

tests/python/test_clang.py

@@ -538,7 +538,7 @@ def test_update_macro_arg(self):
  t = b["act"]
  self.assertEqual(len(t), 32);
 
- def test_ext_ptr_maps(self):
+ def test_ext_ptr_maps1(self):
  bpf_text = """
 #include <uapi/linux/ptrace.h>
 #include <net/sock.h>
@@ -568,6 +568,35 @@ def test_ext_ptr_maps(self):
  b.load_func("trace_entry", BPF.KPROBE)
  b.load_func("trace_exit", BPF.KPROBE)
 
+ def test_ext_ptr_maps2(self):
+ bpf_text = """
+#include <uapi/linux/ptrace.h>
+#include <net/sock.h>
+#include <bcc/proto.h>
+
+BPF_HASH(currsock, u32, struct sock *);
+
+int trace_entry(struct pt_regs *ctx, struct sock *sk,
+ struct sockaddr *uaddr, int addr_len) {
+ u32 pid = bpf_get_current_pid_tgid();
+ currsock.update(&pid, &sk);
+ return 0;
+};
+
+int trace_exit(struct pt_regs *ctx) {
+ u32 pid = bpf_get_current_pid_tgid();
+ struct sock **skpp = currsock.lookup(&pid);
+ if (skpp) {
+ struct sock *skp = *skpp;
+ return skp->__sk_common.skc_dport;
+ }
+ return 0;
+}
+ """
+ b = BPF(text=bpf_text)
+ b.load_func("trace_entry", BPF.KPROBE)
+ b.load_func("trace_exit", BPF.KPROBE)
+
  def test_ext_ptr_maps_reverse(self):
  bpf_text = """
 #include <uapi/linux/ptrace.h>