Clean some brittle parts of clang rewriter

* Rewrites of text inside of a macro (even if just the arguments) is not
  support by clang. Convert macro definitions to pure rewrites instead.
* For packet field access, no longer require 'skb' named
  argument...instead, learn it from the function parameter list.
* Add a complex test case...supposedly this should have failed issue
  #10, but the C version does not exhibit the same failure as the B
  version.
Signed-off-by: Brenden Blanco <bblanco@plumgrid.com>

src/CMakeLists.txt

@@ -4,8 +4,6 @@
 set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS}")
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fno-rtti")
 
-#include_directories(${CMAKE_CURRENT_SOURCE_DIR}/compat/include)
-
 include_directories(${CMAKE_CURRENT_SOURCE_DIR})
 include_directories(${CMAKE_CURRENT_BINARY_DIR})
 

src/cc/b_frontend_action.cc

@@ -32,6 +32,7 @@ using std::map;
 using std::string;
 using std::to_string;
 using std::unique_ptr;
+using std::vector;
 using namespace clang;
 
 // Encode the struct layout as a json description
@@ -80,12 +81,34 @@ bool BTypeVisitor::VisitFunctionDecl(FunctionDecl *D) {
   // put each non-static non-inline function decl in its own section, to be
   // extracted by the MemoryManager
   if (D->isExternallyVisible() && D->hasBody()) {
-    string attr = string("__attribute__((section(\".") + D->getName().str() + "\")))";
+    string attr = string("__attribute__((section(\".") + D->getName().str() + "\")))\n";
     rewriter_.InsertText(D->getLocStart(), attr);
+    // remember the arg names of the current function...first one is the ctx
+    fn_args_.clear();
+    for (auto arg : D->params()) {
+      if (arg->getName() == "") {
+        C.getDiagnostics().Report(arg->getLocEnd(), diag::err_expected)
+            << "named arguments in BPF program definition";
+        return false;
+      }
+      fn_args_.push_back(arg->getName());
+    }
   }
   return true;
 }
 
+// Reverse the order of call traversal so that parameters inside of
+// function calls will get rewritten before the call itself, otherwise
+// text mangling will result.
+bool BTypeVisitor::TraverseCallExpr(CallExpr *Call) {
+  for (auto child : Call->children())
+    if (!TraverseStmt(child))
+      return false;
+  if (!WalkUpFromCallExpr(Call))
+    return false;
+  return true;
+}
+
 // convert calls of the type:
 //  table.foo(&key)
 // to:
@@ -152,10 +175,56 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
 
           txt = prefix + args + suffix;
         }
+        if (!rewriter_.isRewritable(Call->getLocStart())) {
+          C.getDiagnostics().Report(Call->getLocStart(), diag::err_expected)
+              << "use of map function not in a macro";
+          return false;
+        }
         rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), txt);
         return true;
       }
     }
+  } else if (Call->getCalleeDecl()) {
+    NamedDecl *Decl = dyn_cast<NamedDecl>(Call->getCalleeDecl());
+    if (!Decl) return true;
+    if (AsmLabelAttr *A = Decl->getAttr<AsmLabelAttr>()) {
+      // Functions with the tag asm("llvm.bpf.extra") are implemented in the
+      // rewriter rather than as a macro since they may also include nested
+      // rewrites, and clang::Rewriter does not support rewrites in macros,
+      // unless one preprocesses the entire source file.
+      if (A->getLabel() == "llvm.bpf.extra") {
+        if (!rewriter_.isRewritable(Call->getLocStart())) {
+          C.getDiagnostics().Report(Call->getLocStart(), diag::err_expected)
+              << "use of extra builtin not in a macro";
+          return false;
+        }
+
+        vector<string> args;
+        for (auto arg : Call->arguments())
+          args.push_back(rewriter_.getRewrittenText(SourceRange(arg->getLocStart(), arg->getLocEnd())));
+
+        string text;
+        if (Decl->getName() == "incr_cksum_l3") {
+          text = "bpf_l3_csum_replace_(" + fn_args_[0] + ", (u64)";
+          text += args[0] + ", " + args[1] + ", " + args[2] + ", sizeof(" + args[2] + "))";
+        } else if (Decl->getName() == "incr_cksum_l4") {
+          text = "bpf_l4_csum_replace_(" + fn_args_[0] + ", (u64)";
+          text += args[0] + ", " + args[1] + ", " + args[2];
+          text += ", ((" + args[3] + " & 0x1) << 4) | sizeof(" + args[2] + "))";
+        } else if (Decl->getName() == "bpf_trace_printk") {
+          //  #define bpf_trace_printk(fmt, args...)
+          //    ({ char _fmt[] = fmt; bpf_trace_printk_(_fmt, sizeof(_fmt), args...); })
+          text = "({ char _fmt[] = " + args[0] + "; bpf_trace_printk_(_fmt, sizeof(_fmt), ";
+          for (auto arg = args.begin() + 1; arg != args.end(); ++arg) {
+            text += *arg;
+            if (arg + 1 != args.end())
+              text += ", ";
+          }
+          text += "); })";
+        }
+        rewriter_.ReplaceText(SourceRange(Call->getLocStart(), Call->getLocEnd()), text);
+      }
+    }
   }
   return true;
 }
@@ -170,11 +239,16 @@ bool BTypeVisitor::VisitBinaryOperator(BinaryOperator *E) {
       if (DeprecatedAttr *A = Base->getDecl()->getAttr<DeprecatedAttr>()) {
         if (A->getMessage() == "packet") {
           if (FieldDecl *F = dyn_cast<FieldDecl>(Memb->getMemberDecl())) {
+            if (!rewriter_.isRewritable(E->getLocStart())) {
+              C.getDiagnostics().Report(E->getLocStart(), diag::err_expected)
+                  << "use of \"packet\" header type not in a macro";
+              return false;
+            }
             uint64_t ofs = C.getFieldOffset(F);
             uint64_t sz = F->isBitField() ? F->getBitWidthValue(C) : C.getTypeSize(F->getType());
             string base = rewriter_.getRewrittenText(SourceRange(Base->getLocStart(), Base->getLocEnd()));
             string rhs = rewriter_.getRewrittenText(SourceRange(RHS->getLocStart(), RHS->getLocEnd()));
-            string text = "bpf_dins_pkt(skb, (u64)" + base + "+" + to_string(ofs >> 3)
+            string text = "bpf_dins_pkt(" + fn_args_[0] + ", (u64)" + base + "+" + to_string(ofs >> 3)
                 + ", " + to_string(ofs & 0x7) + ", " + to_string(sz) + ", " + rhs + ")";
             rewriter_.ReplaceText(SourceRange(E->getLocStart(), E->getLocEnd()), text);
           }
@@ -196,11 +270,15 @@ bool BTypeVisitor::VisitImplicitCastExpr(ImplicitCastExpr *E) {
     if (DeprecatedAttr *A = Ref->getDecl()->getAttr<DeprecatedAttr>()) {
       if (A->getMessage() == "packet") {
         if (FieldDecl *F = dyn_cast<FieldDecl>(Memb->getMemberDecl())) {
+          if (!rewriter_.isRewritable(E->getLocStart())) {
+            C.getDiagnostics().Report(E->getLocStart(), diag::err_expected)
+                << "use of \"packet\" header type not in a macro";
+            return false;
+          }
           uint64_t ofs = C.getFieldOffset(F);
           uint64_t sz = F->isBitField() ? F->getBitWidthValue(C) : C.getTypeSize(F->getType());
-          string base = rewriter_.getRewrittenText(SourceRange(Base->getLocStart(), Base->getLocEnd()));
-          string text = "bpf_dext_pkt(skb, (u64)" + base + "+" + to_string(ofs >> 3)
-              + ", " + to_string(ofs & 0x7) + ", " + to_string(sz) + ")";
+          string text = "bpf_dext_pkt(" + fn_args_[0] + ", (u64)" + Ref->getDecl()->getName().str() + "+"
+              + to_string(ofs >> 3) + ", " + to_string(ofs & 0x7) + ", " + to_string(sz) + ")";
           rewriter_.ReplaceText(SourceRange(E->getLocStart(), E->getLocEnd()), text);
         }
       }

src/cc/b_frontend_action.h

@@ -68,6 +68,7 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
  public:
   explicit BTypeVisitor(clang::ASTContext &C, clang::Rewriter &rewriter,
                         std::map<std::string, BPFTable> &tables);
+  bool TraverseCallExpr(clang::CallExpr *Call);
   bool VisitFunctionDecl(clang::FunctionDecl *D);
   bool VisitCallExpr(clang::CallExpr *Call);
   bool VisitVarDecl(clang::VarDecl *Decl);
@@ -79,6 +80,7 @@ class BTypeVisitor : public clang::RecursiveASTVisitor<BTypeVisitor> {
   clang::Rewriter &rewriter_;  /// modifications to the source go into this class
   llvm::raw_ostream &out_;  /// for debugging
   std::map<std::string, BPFTable> &tables_;  /// store the open FDs
+  std::vector<std::string> fn_args_;
 };
 
 // A helper class to the frontend action, walks the decls

src/cc/export/helpers.h

@@ -72,8 +72,7 @@ static u64 (*bpf_ktime_get_ns)(void) =
 	(void *) BPF_FUNC_ktime_get_ns;
 static int (*bpf_trace_printk_)(const char *fmt, u64 fmt_size, ...) =
 	(void *) BPF_FUNC_trace_printk;
-#define bpf_trace_printk(_fmt, ...) \
-  ({ char fmt[] = _fmt; bpf_trace_printk_(fmt, sizeof(fmt), ##__VA_ARGS__); })
+int bpf_trace_printk(const char *fmt, ...) asm("llvm.bpf.extra");
 static u64 (*bpf_clone_redirect)(void *ctx, u64 ifindex, u64 flags) =
 	(void *) BPF_FUNC_clone_redirect;
 static u64 (*bpf_get_smp_processor_id)(void) =
@@ -286,10 +285,8 @@ int bpf_l4_csum_replace_(void *ctx, u64 off, u64 from, u64 to, u64 flags) {
   return bpf_l4_csum_replace(ctx, off, from, to, flags);
 }
 
-#define incr_cksum_l3(expr, oldval, newval) \
-  bpf_l3_csum_replace_(skb, (u64)expr, oldval, newval, sizeof(newval))
-#define incr_cksum_l4(expr, oldval, newval, is_pseudo) \
-  bpf_l4_csum_replace_(skb, (u64)expr, oldval, newval, ((is_pseudo & 0x1) << 4) | sizeof(newval))
+int incr_cksum_l3(void *off, u64 oldval, u64 newval) asm("llvm.bpf.extra");
+int incr_cksum_l4(void *off, u64 oldval, u64 newval, u64 flags) asm("llvm.bpf.extra");
 
 #define lock_xadd(ptr, val) ((void)__sync_fetch_and_add(ptr, val))
 

tests/cc/CMakeLists.txt

@@ -20,3 +20,5 @@ add_test(NAME py_test_brb WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_brb_c sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_brb.py test_brb.c)
 add_test(NAME py_test_brb2 WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
   COMMAND ${TEST_WRAPPER} py_brb2_c sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_brb2.py test_brb2.c)
+add_test(NAME py_test_clang WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  COMMAND ${TEST_WRAPPER} py_clang sudo ${CMAKE_CURRENT_SOURCE_DIR}/test_clang.py)

tests/cc/test_clang.py

+#!/usr/bin/env python
+# Copyright (c) PLUMgrid, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License")
+
+from bpf import BPF
+from unittest import main, TestCase
+
+class TestClang(TestCase):
+    def test_complex(self):
+        b = BPF(src_file="test_clang_complex.c", debug=0)
+        fn = b.load_func("handle_packet", BPF.SCHED_CLS)
+    def test_printk(self):
+        text = """
+#include <bcc/proto.h>
+int handle_packet(void *ctx) {
+  BEGIN(ethernet);
+  PROTO(ethernet) {
+    bpf_trace_printk("ethernet->dst = %llx, ethernet->src = %llx\\n",
+                     ethernet->dst, ethernet->src);
+  }
+EOP:
+  return 0;
+}
+"""
+        b = BPF(text=text, debug=0)
+        fn = b.load_func("handle_packet", BPF.SCHED_CLS)
+
+if __name__ == "__main__":
+    main()

tests/cc/test_clang_complex.c

+// Copyright (c) PLUMgrid, Inc.
+// Licensed under the Apache License, Version 2.0 (the "License")
+
+#include <bcc/proto.h>
+
+// hash
+struct FwdKey {
+  u32 dip:32;
+};
+struct FwdLeaf {
+  u32 fwd_idx:32;
+};
+BPF_TABLE("hash", struct FwdKey, struct FwdLeaf, fwd_map, 1);
+
+// array
+struct ConfigKey {
+  u32 index;
+};
+struct ConfigLeaf {
+  u32 bpfdev_ip;
+  u32 slave_ip;
+};
+BPF_TABLE("array", struct ConfigKey, struct ConfigLeaf, config_map, 1);
+
+// hash
+struct MacaddrKey {
+  u32 ip;
+};
+struct MacaddrLeaf {
+  u64 mac;
+};
+BPF_TABLE("hash", struct MacaddrKey, struct MacaddrLeaf, macaddr_map, 11);
+
+// hash
+struct SlaveKey {
+  u32 slave_ip;
+};
+struct SlaveLeaf {
+  u32 slave_ifindex;
+};
+BPF_TABLE("hash", struct SlaveKey, struct SlaveLeaf, slave_map, 10);
+
+int handle_packet(struct __sk_buff *skb) {
+  int ret = 0;
+
+  if (skb->pkt_type == 0) {
+    // tx
+    // make sure configured
+    u32 slave_ip;
+
+    struct ConfigKey cfg_key = {.index = 0};
+    struct ConfigLeaf *cfg_leaf = config_map.lookup(&cfg_key);
+    if (cfg_leaf) {
+      slave_ip = cfg_leaf->slave_ip;
+    } else {
+      return 0xffffffff;
+    }
+
+    // make sure slave configured
+    // tx, default to the single slave
+    struct SlaveKey slave_key = {.slave_ip = slave_ip};
+    struct SlaveLeaf *slave_leaf = slave_map.lookup(&slave_key);
+    if (slave_leaf) {
+      ret = slave_leaf->slave_ifindex;
+    } else {
+      return 0xffffffff;
+    }
+  } else {
+    // rx, default to stack
+    ret = 0;
+  }
+
+  BEGIN(ethernet);
+
+  PROTO(ethernet) {
+    switch (ethernet->type) {
+      case 0x0806: goto arp;
+      case 0x0800: goto ip;
+      case 0x8100: goto dot1q;
+    }
+    goto EOP;
+  }
+
+  PROTO(dot1q) {
+    switch (dot1q->type) {
+      case 0x0806: goto arp;
+      case 0x0800: goto ip;
+    }
+    goto EOP;
+  }
+
+  PROTO(arp) {
+    if (skb->pkt_type) {
+      if (arp->oper == 1) {
+        struct MacaddrKey mac_key = {.ip=arp->spa};
+        struct MacaddrLeaf mac_leaf = {.mac=arp->sha};
+        macaddr_map.update(&mac_key, &mac_leaf);
+      }
+    }
+    goto EOP;
+  }
+
+  PROTO(ip) {
+    switch (ip->nextp) {
+      case 6: goto tcp;
+      case 17: goto udp;
+    }
+    goto EOP;
+  }
+  PROTO(tcp) {
+    goto EOP;
+  }
+  PROTO(udp) {
+    if (udp->dport != 5000) {
+       goto EOP;
+    }
+    if (skb->pkt_type) {
+      // lookup and then forward
+      struct FwdKey fwd_key = {.dip=ip->dst};
+      struct FwdLeaf *fwd_val = fwd_map.lookup(&fwd_key);
+      if (fwd_val) {
+         return fwd_val->fwd_idx;
+      }
+    } else {
+      // rewrite the packet and send to a pre-configured index if needed
+      u32 new_ip;
+      u32 old_ip;
+      u64 src_mac;
+      u64 dst_mac;
+
+      struct ConfigKey cfg_key = {.index = 0};
+      struct ConfigLeaf *cfg_leaf = config_map.lookup(&cfg_key);
+      if (cfg_leaf) {
+        struct MacaddrKey mac_key = {.ip = cfg_leaf->bpfdev_ip};
+        struct MacaddrLeaf *mac_leaf;
+
+        mac_key.ip = cfg_leaf->bpfdev_ip;
+        mac_leaf = macaddr_map.lookup(&mac_key);
+        if (mac_leaf) {
+          src_mac = mac_leaf->mac;
+        } else {
+          goto EOP;
+        }
+
+        mac_key.ip = cfg_leaf->slave_ip;
+        mac_leaf = macaddr_map.lookup(&mac_key);
+        if (mac_leaf) {
+          dst_mac = mac_leaf->mac;
+        } else {
+          goto EOP;
+        }
+
+        // rewrite ethernet header
+        ethernet->dst = dst_mac;
+        ethernet->src = src_mac;
+
+        // ip & udp checksum
+        incr_cksum_l4(&udp->crc, ip->src, cfg_leaf->bpfdev_ip, 1);
+        incr_cksum_l4(&udp->crc, ip->dst, cfg_leaf->slave_ip, 1);
+
+        // rewrite ip src/dst fields
+        ip->src = cfg_leaf->bpfdev_ip;
+        ip->dst = cfg_leaf->slave_ip;
+      }
+    }
+    goto EOP;
+  }
+
+EOP:
+  return ret;
+}