Add probability option to inject

tools/inject.py

@@ -22,12 +22,10 @@
 # Note: presently there are a few hacks to get around various rewriter/verifier
 # issues.
 #
-# Note: this tool requires(as of v4.16-rc5):
-# - commit f7174d08a5fc ("mm: make should_failslab always available for fault
-# injection")
+# Note: this tool requires:
 # - CONFIG_BPF_KPROBE_OVERRIDE
 #
-# USAGE: inject [-h] [-I header] [-v]
+# USAGE: inject [-h] [-I header] [-P probability] [-v] mode spec
 #
 # Copyright (c) 2018 Facebook, Inc.
 # Licensed under the Apache License, Version 2.0 (the "License")
@@ -45,8 +43,9 @@ class Probe:
  }
 
  @classmethod
- def configure(cls, mode):
+ def configure(cls, mode, probability):
  cls.mode = mode
+ cls.probability = probability
 
  def __init__(self, func, preds, length, entry):
  # length of call chain
@@ -68,15 +67,25 @@ def _get_if_top(self):
  if not chk:
  return ""
 
+ if Probe.probability == 1:
+ early_pred = "false"
+ else:
+ early_pred = "bpf_get_prandom_u32() > %s" % str(int((1<<32)*Probe.probability))
  # init the map
  # dont do an early exit here so the singular case works automatically
+ # have an early exit for probability option
  enter = """
+ /*
+ * Early exit for probability case
+ */
+ if (%s)
+ return 0;
  /*
  * Top level function init map
  */
  struct pid_struct p_struct = {0, 0};
  m.insert(&pid, &p_struct);
- """
+ """ % early_pred
 
  # kill the entry
  exit = """
@@ -221,9 +230,9 @@ def _generate_bottom(self):
  if (p->conds_met == %s && %s)
  bpf_override_return(ctx, %s);
  return 0;
-}""" % (self.prep, self.length, pred, self._get_err(), self.length - 1, pred,
-   self._get_err())
- return text
+}"""
+ return text % (self.prep, self.length, pred, self._get_err(),
+  self.length - 1, pred, self._get_err())
 
  # presently parses and replaces STRCMP
  # STRCMP exists because string comparison is inconvenient and somewhat buggy
@@ -302,6 +311,9 @@ def __init__(self):
  parser.add_argument("-I", "--include", action="append",
  metavar="header",
  help="additional header files to include in the BPF program")
+ parser.add_argument("-P", "--probability", default=1,
+ metavar="probability", type=float,
+ help="probability that this call chain will fail")
  parser.add_argument("-v", "--verbose", action="store_true",
  help="print BPF program")
  self.args = parser.parse_args()
@@ -315,7 +327,7 @@ def __init__(self):
  # create_probes and associated stuff
  def _create_probes(self):
  self._parse_spec()
- Probe.configure(self.args.mode)
+ Probe.configure(self.args.mode, self.args.probability)
  # self, func, preds, total, entry
 
  # create all the pair probes
@@ -425,8 +437,7 @@ def _attach_probes(self):
 
  def _generate_program(self):
  # leave out auto includes for now
-
- self.program += "#include <linux/mm.h>\n"
+ self.program += '#include <linux/mm.h>\n'
  for include in (self.args.include or []):
  self.program += "#include <%s>\n" % include