#!/bin/bash
#
#  kprobestest: Kprobes stress test tool
#  Written by Masami Hiramatsu <mhiramat@redhat.com>
#
#  Usage:
#     $ kprobestest [-s SYMLIST] [-b BLACKLIST] [-w WHITELIST]
#        Run stress test. If SYMLIST file is specified, use it as 
#        an initial symbol list (This is useful for verifying white list
#        after diagnosing all symbols).
#
#     $ kprobestest cleanup
#        Cleanup all lists

DEBUGFS=/debug
INITNR=64
DIV=4
SYMFILE=syms.list
FAILFILE=black.list
FAILBKUP=1

function do_test () {
  # Do some benchmark
  for i in {1..4} ; do
  sleep 0.5
  echo -n "."
  done
}

function usage () {
  echo "Usage: kprobestest [cleanup] [-s SYMLIST] [-b BLACKLIST] [-w WHITELIST]"
  exit 0
}

function cleanup_test () {
  rm -rf $SYMFILE failed passed testing unset backup
  exit 0
}


# Parse arguments
WHITELIST=
BLACKLIST=
SYMLIST=

while [ "$1" ]; do
case $1 in
  cleanup)
    cleanup_test
    ;;
  run)
    # ignore
    ;;
  -s)
    SYMLIST=$2
    shift 1
    ;;
  -b)
    BLACKLIST=$2
    shift 1
    ;;
  -w)
    WHITELIST=$2
    shift 1
    ;;
  *)
    usage
    ;;
esac
shift 1
done

# Show configurations
echo "Kprobe stress test starting."
[ -f "$BLACKLIST" ] && echo "Blacklist: $BLACKLIST" || BLACKLIST=""
[ -f "$WHITELIST" ] && echo "Whitelist: $WHITELIST" || WHITELIST=""
[ -f "$SYMLIST" ] && echo "Symlist: $SYMLIST" || SYMLIST=""

function make_filter () {
local EXP=""
if [ -z "$WHITELIST" -a -z "$BLACKLIST" ]; then
  echo "s/^$//g"
else
  for i in `cat $WHITELIST $BLACKLIST` ;do
    [ -z "$EXP" ] && EXP="^$i\$" || EXP="$EXP\\|^$i\$"
  done ; EXP="s/$EXP//g"
  echo $EXP
fi
}

function list_allsyms () {
local sym
local out=0
for sym in `sort /proc/kallsyms | egrep '[0-9a-f]+ [Tt] [^[]*$' | cut -d\  -f 3`;do
  [ $sym  = "_stext" ] && out=1 && continue
  [ $sym  = "__kprobes_text_start" ] && out=0 && continue
  [ $sym  = "__kprobes_text_end" ] && out=1 && continue
  [ $sym  = "_etext" ] && break
  [ $out -eq 1 ] && echo $sym
done
}

function prep_testing () {
local i=0
local n=0
local NR=$1
local fname=

echo "Grouping symbols: $NR"

fname=`printf "list-%03d.%d" $i $NR`
cat $SYMFILE | while read ln; do
  [ -z "$ln" ] && continue
  echo "$ln" >> testing/$fname
  n=$((n+1))
  if [ $n -eq $NR ]; then
    n=0
    i=$((i+1))
    fname=`printf "list-%03d.%d" $i $NR`
  fi
done
sync
}

function init_first () {
local EXP
EXP=`make_filter`
if [ -f "$SYMLIST" ]; then
  cat $SYMLIST | sed $EXP > $SYMFILE
else
  echo -n "Generating symbol list from /proc/kallsyms..."
  list_allsyms | sed $EXP > $SYMFILE
  echo "done. " `wc -l $SYMFILE | cut -f1 -d\  ` "symbols listed."
fi
mkdir testing
mkdir failed
mkdir backup
mkdir unset
mkdir passed
mkdir passed/profiles
prep_testing $INITNR
}

function get_max_nr () {
wc -l failed/list-* unset/list-* 2>/dev/null |\
awk '/^ *[0-9]+ .*list.*$/{ if (nr < $1) nr=$1 } BEGIN { nr=0 } END { print nr}'
}

function init_next () {
local NR
NR=`get_max_nr`
[ $NR -eq 0 ] && return 1
[ $NR -eq 1 ] && return 2
[ $NR -le $DIV ] && NR=1 || NR=`expr $NR / $DIV`

cat failed/* unset/* > $SYMFILE
if [ $FAILBKUP ]; then
    cp failed/* unset/* backup/
fi
rm failed/* unset/*

prep_testing $NR
return 0
}


# Initialize symbols
if [ ! -d testing ]; then
  init_first
elif [ -z "`ls testing/`" ]; then
  init_next
fi

function set_probes () {
local s
for s in `cat $1`; do
  echo "p:$s" $s >> $DEBUGFS/tracing/kprobe_events
  [ $? -ne 0 ] && return 1
  echo 1 > $DEBUGFS/tracing/events/kprobes/$s/enable
done
return 0
}

function clear_probes () {
  echo > $DEBUGFS/tracing/kprobe_events
}

function save_profile () {
  cat $DEBUGFS/tracing/kprobe_profile > passed/profiles/$1.prof
}

clear_probes

echo "Starting tests.."

RET=0

while [ $RET -eq 0 ]; do

  for list in `cd testing/; ls`; do
  echo -n $list
  mv testing/$list failed/
  sync;sync
  echo -n "sync.."
  set_probes failed/$list
  if [ $? -ne 0 ]; then
    clear_probes
    sync;sync
    echo "can not set"
    mv failed/$list unset/
    sync;sync
  else
    do_test
    save_profile $list
    clear_probes
    sync;sync
    echo "done"
    mv failed/$list passed/
    sync;sync
  fi
  done

  init_next
  RET=$?
done

if [ $RET -eq 1 ];then
# No failed symbols
  echo "no failed symbols found."
else
  echo "found failed symbols:"
  cat failed/* | tee $FAILFILE
  rm failed/*
fi

cat unset/* > "unset.list"
rm unset/*

function profile_symbols () {
  local s h m
  rm -f tested.list missed.list untested.list
  cat passed/profiles/*.prof | while read s h m ;do
    if [ $h -ne 0 ]; then 
      echo $s >> tested.list
    elif [ $m -ne 0 ]; then
      echo $s >> missed.list
    else
      echo $s >> untested.list
    fi
  done
}

echo -n "Profiling symbols..."
profile_symbols
echo done
echo tested: `wc -l tested.list | cut -d\  -f1` symbols
echo missed: `wc -l missed.list | cut -d\  -f1` symbols
echo untested: `wc -l untested.list | cut -d\  -f1` symbols
echo unset: `wc -l unset.list | cut -d\  -f1` symbols