inject: add option to limit number of overrides (#2044)

* inject: add option to limit number of overrides

* inject: update man pages with new count option

Signed-off-by: Nicolas Saenz Julienne <nsaenzjulienne@suse.de>
diff --git a/tools/inject.py b/tools/inject.py
index 031679b..741a0fb 100755
--- a/tools/inject.py
+++ b/tools/inject.py
@@ -44,9 +44,10 @@
     }
 
     @classmethod
-    def configure(cls, mode, probability):
+    def configure(cls, mode, probability, count):
         cls.mode = mode
         cls.probability = probability
+        cls.count = count
 
     def __init__(self, func, preds, length, entry):
         # length of call chain
@@ -207,6 +208,14 @@
         pred = self.preds[0][0]
         text = self._get_heading() + """
 {
+        u32 overriden = 0;
+        int zero = 0;
+        u32* val;
+
+        val = count.lookup(&zero);
+        if (val)
+            overriden = *val;
+
         /*
          * preparation for predicate, if necessary
          */
@@ -214,7 +223,8 @@
         /*
          * If this is the only call in the chain and predicate passes
          */
-        if (%s == 1 && %s) {
+        if (%s == 1 && %s && overriden < %s) {
+                count.increment(zero);
                 bpf_override_return(ctx, %s);
                 return 0;
         }
@@ -228,12 +238,15 @@
         /*
          * If all conds have been met and predicate passes
          */
-        if (p->conds_met == %s && %s)
+        if (p->conds_met == %s && %s && overriden < %s) {
+                count.increment(zero);
                 bpf_override_return(ctx, %s);
+        }
         return 0;
 }"""
-        return text % (self.prep, self.length, pred, self._get_err(),
-                    self.length - 1, pred, self._get_err())
+        return text % (self.prep, self.length, pred, Probe.count,
+                self._get_err(), self.length - 1, pred, Probe.count,
+                self._get_err())
 
     # presently parses and replaces STRCMP
     # STRCMP exists because string comparison is inconvenient and somewhat buggy
@@ -333,6 +346,8 @@
                 help="probability that this call chain will fail")
         parser.add_argument("-v", "--verbose", action="store_true",
                 help="print BPF program")
+        parser.add_argument("-c", "--count", action="store", default=-1,
+                help="Number of fails before bypassing the override")
         self.args = parser.parse_args()
 
         self.program = ""
@@ -344,7 +359,7 @@
     # create_probes and associated stuff
     def _create_probes(self):
         self._parse_spec()
-        Probe.configure(self.args.mode, self.args.probability)
+        Probe.configure(self.args.mode, self.args.probability, self.args.count)
         # self, func, preds, total, entry
 
         # create all the pair probes
@@ -482,6 +497,8 @@
 
         self.program += self._def_pid_struct()
         self.program += "BPF_HASH(m, u32, struct pid_struct);\n"
+        self.program += "BPF_ARRAY(count, u32, 1);\n"
+
         for p in self.probes:
             self.program += p.generate_program() + "\n"