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

* inject: add option to limit number of overrides * inject: update man pages with new count option Signed-off-by: Nicolas Saenz Julienne <[email protected]>
mwesolowski · Nov 15, 2018 · d51219a · d51219a
1 parent 8417f69
commit d51219a
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 7 deletions.
diff --git a/man/man8/inject.8 b/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

diff --git a/tools/inject.py b/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,14 +208,23 @@ def _generate_bottom(self):
  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
  */
  %s
  /*
  * 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 @@ def _generate_bottom(self):
  /*
  * 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 @@ def __init__(self):
  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 @@ def __init__(self):
  # 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 @@ def _generate_program(self):
 
  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"