Merge pull request #308 from iovisor/cflags

Add ability to set custom cflags when loading programs
diff --git a/src/cc/bpf_common.cc b/src/cc/bpf_common.cc
index 029da49..f68372c 100644
--- a/src/cc/bpf_common.cc
+++ b/src/cc/bpf_common.cc
@@ -26,18 +26,18 @@
   return mod;
 }
 
-void * bpf_module_create_c(const char *filename, unsigned flags) {
+void * bpf_module_create_c(const char *filename, unsigned flags, const char *cflags[], int ncflags) {
   auto mod = new ebpf::BPFModule(flags);
-  if (mod->load_c(filename) != 0) {
+  if (mod->load_c(filename, cflags, ncflags) != 0) {
     delete mod;
     return nullptr;
   }
   return mod;
 }
 
-void * bpf_module_create_c_from_string(const char *text, unsigned flags) {
+void * bpf_module_create_c_from_string(const char *text, unsigned flags, const char *cflags[], int ncflags) {
   auto mod = new ebpf::BPFModule(flags);
-  if (mod->load_string(text) != 0) {
+  if (mod->load_string(text, cflags, ncflags) != 0) {
     delete mod;
     return nullptr;
   }
diff --git a/src/cc/bpf_common.h b/src/cc/bpf_common.h
index e066c58..9d02d9b 100644
--- a/src/cc/bpf_common.h
+++ b/src/cc/bpf_common.h
@@ -25,8 +25,8 @@
 #endif
 
 void * bpf_module_create_b(const char *filename, const char *proto_filename, unsigned flags);
-void * bpf_module_create_c(const char *filename, unsigned flags);
-void * bpf_module_create_c_from_string(const char *text, unsigned flags);
+void * bpf_module_create_c(const char *filename, unsigned flags, const char *cflags[], int ncflags);
+void * bpf_module_create_c_from_string(const char *text, unsigned flags, const char *cflags[], int ncflags);
 void bpf_module_destroy(void *program);
 char * bpf_module_license(void *program);
 unsigned bpf_module_kern_version(void *program);
diff --git a/src/cc/bpf_module.cc b/src/cc/bpf_module.cc
index bd55a0e..767adee 100644
--- a/src/cc/bpf_module.cc
+++ b/src/cc/bpf_module.cc
@@ -299,9 +299,9 @@
 }
 
 // load an entire c file as a module
-int BPFModule::load_cfile(const string &file, bool in_memory) {
+int BPFModule::load_cfile(const string &file, bool in_memory, const char *cflags[], int ncflags) {
   clang_loader_ = make_unique<ClangLoader>(&*ctx_, flags_);
-  if (clang_loader_->parse(&mod_, &tables_, file, in_memory))
+  if (clang_loader_->parse(&mod_, &tables_, file, in_memory, cflags, ncflags))
     return -1;
   return 0;
 }
@@ -313,7 +313,7 @@
 // build an ExecutionEngine.
 int BPFModule::load_includes(const string &tmpfile) {
   clang_loader_ = make_unique<ClangLoader>(&*ctx_, flags_);
-  if (clang_loader_->parse(&mod_, &tables_, tmpfile, false))
+  if (clang_loader_->parse(&mod_, &tables_, tmpfile, false, nullptr, 0))
     return -1;
   return 0;
 }
@@ -668,7 +668,7 @@
 }
 
 // load a C file
-int BPFModule::load_c(const string &filename) {
+int BPFModule::load_c(const string &filename, const char *cflags[], int ncflags) {
   if (!sections_.empty()) {
     fprintf(stderr, "Program already initialized\n");
     return -1;
@@ -677,7 +677,7 @@
     fprintf(stderr, "Invalid filename\n");
     return -1;
   }
-  if (int rc = load_cfile(filename, false))
+  if (int rc = load_cfile(filename, false, cflags, ncflags))
     return rc;
   if (int rc = annotate())
     return rc;
@@ -687,12 +687,12 @@
 }
 
 // load a C text string
-int BPFModule::load_string(const string &text) {
+int BPFModule::load_string(const string &text, const char *cflags[], int ncflags) {
   if (!sections_.empty()) {
     fprintf(stderr, "Program already initialized\n");
     return -1;
   }
-  if (int rc = load_cfile(text, true))
+  if (int rc = load_cfile(text, true, cflags, ncflags))
     return rc;
   if (int rc = annotate())
     return rc;
diff --git a/src/cc/bpf_module.h b/src/cc/bpf_module.h
index 541341d..f79a938 100644
--- a/src/cc/bpf_module.h
+++ b/src/cc/bpf_module.h
@@ -48,15 +48,15 @@
   void dump_ir(llvm::Module &mod);
   int load_file_module(std::unique_ptr<llvm::Module> *mod, const std::string &file, bool in_memory);
   int load_includes(const std::string &tmpfile);
-  int load_cfile(const std::string &file, bool in_memory);
+  int load_cfile(const std::string &file, bool in_memory, const char *cflags[], int ncflags);
   int kbuild_flags(const char *uname_release, std::vector<std::string> *cflags);
   int run_pass_manager(llvm::Module &mod);
  public:
   BPFModule(unsigned flags);
   ~BPFModule();
   int load_b(const std::string &filename, const std::string &proto_filename);
-  int load_c(const std::string &filename);
-  int load_string(const std::string &text);
+  int load_c(const std::string &filename, const char *cflags[], int ncflags);
+  int load_string(const std::string &text, const char *cflags[], int ncflags);
   size_t num_functions() const;
   uint8_t * function_start(size_t id) const;
   uint8_t * function_start(const std::string &name) const;
diff --git a/src/cc/frontends/clang/loader.cc b/src/cc/frontends/clang/loader.cc
index 45fda7d..7b2de83 100644
--- a/src/cc/frontends/clang/loader.cc
+++ b/src/cc/frontends/clang/loader.cc
@@ -65,7 +65,7 @@
 ClangLoader::~ClangLoader() {}
 
 int ClangLoader::parse(unique_ptr<llvm::Module> *mod, unique_ptr<vector<TableDesc>> *tables,
-                       const string &file, bool in_memory) {
+                       const string &file, bool in_memory, const char *cflags[], int ncflags) {
   using namespace clang;
 
   struct utsname un;
@@ -103,6 +103,10 @@
   kflags.push_back(BCC_INSTALL_PREFIX "/share/bcc/include");
   for (auto it = kflags.begin(); it != kflags.end(); ++it)
     flags_cstr.push_back(it->c_str());
+  if (cflags) {
+    for (auto i = 0; i < ncflags; ++i)
+      flags_cstr.push_back(cflags[i]);
+  }
 
   // set up the error reporting class
   IntrusiveRefCntPtr<DiagnosticOptions> diag_opts(new DiagnosticOptions());
diff --git a/src/cc/frontends/clang/loader.h b/src/cc/frontends/clang/loader.h
index 2f45b64..22c64e1 100644
--- a/src/cc/frontends/clang/loader.h
+++ b/src/cc/frontends/clang/loader.h
@@ -39,7 +39,7 @@
   explicit ClangLoader(llvm::LLVMContext *ctx, unsigned flags);
   ~ClangLoader();
   int parse(std::unique_ptr<llvm::Module> *mod, std::unique_ptr<std::vector<TableDesc>> *tables,
-            const std::string &file, bool in_memory);
+            const std::string &file, bool in_memory, const char *cflags[], int ncflags);
  private:
   llvm::LLVMContext *ctx_;
   unsigned flags_;
diff --git a/src/python/bcc/__init__.py b/src/python/bcc/__init__.py
index b346146..01fd08a 100644
--- a/src/python/bcc/__init__.py
+++ b/src/python/bcc/__init__.py
@@ -30,9 +30,11 @@
 lib.bpf_module_create_b.restype = ct.c_void_p
 lib.bpf_module_create_b.argtypes = [ct.c_char_p, ct.c_char_p, ct.c_uint]
 lib.bpf_module_create_c.restype = ct.c_void_p
-lib.bpf_module_create_c.argtypes = [ct.c_char_p, ct.c_uint]
+lib.bpf_module_create_c.argtypes = [ct.c_char_p, ct.c_uint,
+        ct.POINTER(ct.c_char_p), ct.c_int]
 lib.bpf_module_create_c_from_string.restype = ct.c_void_p
-lib.bpf_module_create_c_from_string.argtypes = [ct.c_char_p, ct.c_uint]
+lib.bpf_module_create_c_from_string.argtypes = [ct.c_char_p, ct.c_uint,
+        ct.POINTER(ct.c_char_p), ct.c_int]
 lib.bpf_module_destroy.restype = None
 lib.bpf_module_destroy.argtypes = [ct.c_void_p]
 lib.bpf_module_license.restype = ct.c_char_p
@@ -395,7 +397,7 @@
                     raise Exception("Could not find file %s" % filename)
         return filename
 
-    def __init__(self, src_file="", hdr_file="", text=None, cb=None, debug=0):
+    def __init__(self, src_file="", hdr_file="", text=None, cb=None, debug=0, cflags=[]):
         """Create a a new BPF module with the given source code.
 
         Note:
@@ -416,8 +418,11 @@
         self.debug = debug
         self.funcs = {}
         self.tables = {}
+        cflags_array = (ct.c_char_p * len(cflags))()
+        for i, s in enumerate(cflags): cflags_array[i] = s.encode("ascii")
         if text:
-            self.module = lib.bpf_module_create_c_from_string(text.encode("ascii"), self.debug)
+            self.module = lib.bpf_module_create_c_from_string(text.encode("ascii"),
+                    self.debug, cflags_array, len(cflags_array))
         else:
             src_file = BPF._find_file(src_file)
             hdr_file = BPF._find_file(hdr_file)
@@ -426,7 +431,7 @@
                         hdr_file.encode("ascii"), self.debug)
             else:
                 self.module = lib.bpf_module_create_c(src_file.encode("ascii"),
-                        self.debug)
+                        self.debug, cflags_array, len(cflags_array))
 
         if self.module == None:
             raise Exception("Failed to compile BPF module %s" % src_file)
diff --git a/tests/cc/test_clang.py b/tests/cc/test_clang.py
index 1b46648..d701b65 100755
--- a/tests/cc/test_clang.py
+++ b/tests/cc/test_clang.py
@@ -287,5 +287,13 @@
         import ctypes
         self.assertEqual(ctypes.sizeof(b["t3"].Leaf), 8)
 
+    def test_cflags(self):
+        text = """
+#ifndef MYFLAG
+#error "MYFLAG not set as expected"
+#endif
+"""
+        b = BPF(text=text, cflags=["-DMYFLAG"])
+
 if __name__ == "__main__":
     main()