Merge pull request #632 from markdrayton/probe-strings

Probe registration fixes

src/python/bcc/__init__.py

@@ -334,11 +334,12 @@ class BPF(object):
         fn.sock = sock
 
     def _get_kprobe_functions(self, event_re):
-        blacklist = set([line.rstrip().split()[1] for line in
-                open("%s/../kprobes/blacklist" % TRACEFS)])
+        with open("%s/../kprobes/blacklist" % TRACEFS) as blacklist_file:
+            blacklist = set([line.rstrip().split()[1] for line in
+                    blacklist_file])
         fns = []
-        with open("%s/available_filter_functions" % TRACEFS) as f:
-            for line in f:
+        with open("%s/available_filter_functions" % TRACEFS) as avail_file:
+            for line in avail_file:
                 fn = line.rstrip().split()[0]
                 if re.match(event_re, fn) and fn not in blacklist:
                     fns.append(fn)
@@ -363,7 +364,6 @@ class BPF(object):
     def attach_kprobe(self, event="", fn_name="", event_re="",
             pid=-1, cpu=0, group_fd=-1):
 
-        assert isinstance(event, str), "event must be a string"
         # allow the caller to glob multiple functions together
         if event_re:
             for line in self._get_kprobe_functions(event_re):
@@ -374,6 +374,7 @@ class BPF(object):
                     pass
             return
 
+        event = str(event)
         self._check_probe_quota(1)
         fn = self.load_func(fn_name, BPF.KPROBE)
         ev_name = "p_" + event.replace("+", "_").replace(".", "_")
@@ -388,7 +389,7 @@ class BPF(object):
         return self
 
     def detach_kprobe(self, event):
-        assert isinstance(event, str), "event must be a string"
+        event = str(event)
         ev_name = "p_" + event.replace("+", "_").replace(".", "_")
         if ev_name not in self.open_kprobes:
             raise Exception("Kprobe %s is not attached" % event)
@@ -402,7 +403,6 @@ class BPF(object):
     def attach_kretprobe(self, event="", fn_name="", event_re="",
             pid=-1, cpu=0, group_fd=-1):
 
-        assert isinstance(event, str), "event must be a string"
         # allow the caller to glob multiple functions together
         if event_re:
             for line in self._get_kprobe_functions(event_re):
@@ -413,6 +413,7 @@ class BPF(object):
                     pass
             return
 
+        event = str(event)
         self._check_probe_quota(1)
         fn = self.load_func(fn_name, BPF.KPROBE)
         ev_name = "r_" + event.replace("+", "_").replace(".", "_")
@@ -427,7 +428,7 @@ class BPF(object):
         return self
 
     def detach_kretprobe(self, event):
-        assert isinstance(event, str), "event must be a string"
+        event = str(event)
         ev_name = "r_" + event.replace("+", "_").replace(".", "_")
         if ev_name not in self.open_kprobes:
             raise Exception("Kretprobe %s is not attached" % event)
@@ -526,6 +527,7 @@ class BPF(object):
                  BPF(text).attach_uprobe("/usr/bin/python", "main")
         """
 
+        name = str(name)
         (path, addr) = BPF._check_path_symbol(name, sym, addr)
 
         self._check_probe_quota(1)
@@ -548,6 +550,7 @@ class BPF(object):
         or binary 'name'.
         """
 
+        name = str(name)
         (path, addr) = BPF._check_path_symbol(name, sym, addr)
         ev_name = "p_%s_0x%x" % (self._probe_repl.sub("_", path), addr)
         if ev_name not in self.open_uprobes:
@@ -569,6 +572,7 @@ class BPF(object):
         meaning of additional parameters.
         """
 
+        name = str(name)
         (path, addr) = BPF._check_path_symbol(name, sym, addr)
 
         self._check_probe_quota(1)
@@ -591,6 +595,7 @@ class BPF(object):
         or binary 'name'.
         """
 
+        name = str(name)
         (path, addr) = BPF._check_path_symbol(name, sym, addr)
         ev_name = "r_%s_0x%x" % (self._probe_repl.sub("_", path), addr)
         if ev_name not in self.open_uprobes:
@@ -604,7 +609,7 @@ class BPF(object):
 
     def _trace_autoload(self):
         for i in range(0, lib.bpf_num_functions(self.module)):
-            func_name = lib.bpf_function_name(self.module, i)
+            func_name = str(lib.bpf_function_name(self.module, i).decode())
             if func_name.startswith("kprobe__"):
                 fn = self.load_func(func_name, BPF.KPROBE)
                 self.attach_kprobe(event=fn.name[8:], fn_name=fn.name)
@@ -768,14 +773,14 @@ class BPF(object):
             exit()
 
     def cleanup(self):
-        for k, v in self.open_kprobes.items():
+        for k, v in list(self.open_kprobes.items()):
             lib.perf_reader_free(v)
             # non-string keys here include the perf_events reader
             if isinstance(k, str):
                 desc = "-:kprobes/%s" % k
                 lib.bpf_detach_kprobe(desc.encode("ascii"))
             self._del_kprobe(k)
-        for k, v in self.open_uprobes.items():
+        for k, v in list(self.open_uprobes.items()):
             lib.perf_reader_free(v)
             desc = "-:uprobes/%s" % k
             lib.bpf_detach_uprobe(desc.encode("ascii"))