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>

man/man8/inject.8

@@ -3,7 +3,7 @@
 inject \- injects appropriate error into function if input call chain and
 predicates are satisfied. Uses Linux eBPF/bcc.
 .SH SYNOPSIS
-.B inject -h [-I header] [-P probability] [-v] mode spec
+.B inject -h [-I header] [-P probability] [-v] [-C count] mode spec
 .SH DESCRIPTION
 inject injects errors into specified kernel functionality when a given call
 chain and associated predicates are satsified.
@@ -29,6 +29,9 @@ Necessary headers to be included.
 .TP
 \-P probability
 Optional probability of failure, default 1.
+.TP
+\-C count
+Number of errors to inject before stopping, default never stops.
 .SH EXAMPLES
 Please see inject_example.txt
 .SH SOURCE

tools/inject.py

@@ -44,9 +44,10 @@ class Probe:
     }
 
     @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 @@ class Probe:
         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 @@ class Probe:
         /*
          * 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 @@ class Probe:
         /*
          * 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 @@ EXAMPLES:
                 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 @@ EXAMPLES:
     # 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 @@ struct pid_struct {
 
         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"