#!/bin/sh

ver="1.0.2"
bt_load=__bt_load
bt_read=__bt_read
bt_disym=__bt_disym
bt_get_addr=__bt_get_addr
syscall_n2i=__bt_syscall_name2id
module_dir=/lib/modules/`uname -r`/kernel/drivers/btrax
relfs_mod=`find $module_dir -name relayfs.ko`
djprobe_mod=`find $module_dir -name mod_djprobe.ko`
is_rhel=""
if [ -f /etc/redhat-release ]; then
	is_rhel=`grep '^Red Hat Enterprise Linux AS release 4' /etc/redhat-release`
fi
ctr_mod=$module_dir/ctr_mod.ko
ctr_loaded=`cat /proc/modules|grep '^ctr_mod ' 2> /dev/null`
on_off_mod=$module_dir/bt_on_off_mod.ko

usage() {
	local cmd
	cmd=`basename $0`
cat << EOF
$cmd $ver usage
$cmd [common_opt] [userland_trace_opt | all_trace_opt | syscall_trace_opt]
common_opt (default value):
   -S bts_size:   BTS buffer size (786432)
   -M bts_margin: BTS interrupt margin records (8192)
   -s size:       relayfs sub-buffer size (2097152)
   -n number:     number of relayfs sub-buffers (16)
   -N number:     target process sleep threshold of relayfs sub-buffers (8)
   -m mnt_dir:    relayfs mount directory (/mnt/relay)

userland_trace_opt:
   -p pid[,...] -d log_dir
   -d log_dir -c cmd [arg...]

all_trace_opt:
   [--start symbol] [--stop symbol] -d log_dir
   [--START addr,size] [--STOP addr,size] -d log_dir
   --fr -d log_dir
   --exit-fr [-d logdir]

syscall_trace_opt:
   --syscall fname[,...] -p pid[,...] -d log_dir
   --syscall fname[,...] -d log_dir -c cmd [arg...]

      -d log_dir:                 log write directory
      -p pid[,...]:               pid(s) of the target application(s)
      -c cmd [arg...]:            trace target application
      [--start|--stop] symbol:    trace start/stop symbol
      [--START|--STOP] addr,size: trace start/stop address and size
         symbol and addr shoud be exist in the kernel space.
      --syscall fname[,...]:       system call function name(s) for trace
EOF
}

terminate() {
	local temp
	if [ "$cmd_pid" ]; then
		temp=`ps -U root|grep $cmd_pid`
		if [ "$temp" ]; then
			kill $cmd_pid
		fi
	fi
	cat /proc/modules|grep '^bt_on_off_mod ' > /dev/null
	if [ $? -eq 0 ]; then
		rmmod bt_on_off_mod
	fi
	cat /proc/modules|grep '^bt_mod ' > /dev/null
	if [ $? -eq 0 ]; then
		rmmod bt_mod
	fi
	mount|grep '^relayfs\|^debugfs' > /dev/null
	if [ $? -eq 0 ]; then
		umount $1
	fi
	cat /proc/modules|grep '^relayfs ' > /dev/null
	if [ $? -eq 0 ]; then
		rmmod relayfs
	fi
	cat /proc/modules|grep '^mod_djprobe ' > /dev/null
	if [ $? -eq 0 ]; then
		rmmod mod_djprobe
	fi
}

kill_bt_read() {
	local pid_num pid
	pid_num=`ps -U root|grep '\<bt_read\>'|wc -l`
	pid_num=`expr $pid_num - 1`
	for pid in `ps -U root|grep '\<bt_read\>'|cut -c1-5|tail -n $pid_num`;do
		kill $pid
	done
}

wait_for_any_cpu_full() {
	local n_subbufs any_full
	n_subbufs=`cat /proc/btrax/n_subbufs`
	full_any=0
	while [ $full_any -eq 0 ]; do
		for produced in /proc/btrax/cpu*/produced; do
			if [ `od -A n -l $produced` -eq $n_subbufs ]; then
				full_any=1
				break
			fi
		done
		usleep 100000
	done
}

wait_for_bts_disable() {
	while [ `cat /proc/btrax/enable` -eq 1 ]; do
		usleep 100000
	done
}

add_bt_mod_mode() {
	if [ "$start" ] && [ "$stop" ]; then
		mod_param=`echo $mod_param mode=0`
	elif [ "$start" ]; then
		mod_param=`echo $mod_param mode=1`
	elif [ "$stop" ]; then
		mod_param=`echo $mod_param mode=2`
	elif [ "$fr" ]; then
		mod_param=`echo $mod_param mode=3`
	elif [ "$syscall" ]; then
		mod_param=`echo $mod_param mode=4 chk_sctime=$chk_sct`
	else
		mod_param=`echo $mod_param mode=5`
	fi
}

chk_next() {
	if [ $# -lt 2 ]; then
		usage
		exit 1
	fi
}

parse_addr_and_size() {
	local addr size
	echo $2|grep ',' > /dev/null 2>&1
	if [ $? -ne 0 ]; then
		echo ""
		return
	fi
	addr=`echo "${2%%,*}"`
	size=`echo "${2##*,}"`
	if [ "$addr" = "" ] || [ "$size" = "" ]; then
		echo ""
	else
		printf "%s_addr=%s %s_size=%s" $1 $addr $1 $size
	fi
}

chk_cpu_type() {
	local family model
	family=`cat /proc/cpuinfo|grep '^cpu family'|head -n 1|cut -d: -f2`
	model=`cat /proc/cpuinfo|grep '^model[^ ]'|head -n 1|cut -d: -f2`
	if [ $family == 6 ] && ([ $model == 9 ] || [ $model == 13 ]); then
		touch $logdir/from_is_next
	fi
}

bt_insmod() {
	insmod $* > /dev/null 2>&1
	if [ $? -ne 0 ]; then
		echo "!!! Maybe you need root permission (insmod '$1' failed)."
		exit 1
	fi
}

#-------- first, ignore signal --------
trap '' HUP INT QUIT TERM

#-------- check parameters --------
logdir=""
pids=""
mod_param="pid_max=`cat /proc/sys/kernel/pid_max`"
relay_mnt_point=/mnt/relay
start=""
stop=""
cmd=""
cmd_arg=""
fr=""
exit_fr=""
syscall=""
syscall_filter=""
chk_sct="0"

while [ $# -gt 0 ]; do
	case $1 in
	-p)
		chk_next $*
		pids=$2
		;;
	--syscall)
		chk_next $*
		syscall="1"
		syscall_filter=`$syscall_n2i $2`
		if [ $? -ne 0 ]; then
			echo \"$2\" includes not system call function
			exit 1
		fi;
		;;
	--chk-sctime)
		chk_next $*
		chk_sct=$2
		;;
	-c)
		chk_next $*
		cmd=$2
		shift 2
		cmd_arg=$*
		break
		;;
	-m)
		chk_next $*
		relay_mnt_point=$2
		;;
	-d)
		chk_next $*
		logdir=$2
		;;
	-S)
		chk_next $*
		mod_param=`echo $mod_param btsbuf_size=$2`
		;;
	-M)
		chk_next $*
		mod_param=`echo $mod_param bts_int_margin_recs=$2`
		;;
	-s)
		chk_next $*
		mod_param=`echo $mod_param subbuf_size=$2`
		;;
	-n)
		chk_next $*
		mod_param=`echo $mod_param subbuf_num=$2`
		;;
	-N)
		chk_next $*
		mod_param=`echo $mod_param subbuf_sleep_threshold=$2`
		;;
	--start)
		chk_next $*
		start=`$bt_disym on $2`
		if [ "$start" = "" ]; then
			echo "symbol '$2' not found"
			exit 1
		fi
		;;
	--stop)
		chk_next $*
		stop=`$bt_disym off $2`
		if [ "$stop" = "" ]; then
			echo "symbol '$2' not found"
			exit 1
		fi
		;;
	--START)
		chk_next $*
		start=`parse_addr_and_size "on" $2`
		if [ "$start" = "" ]; then
			echo "'$2' format error"
			exit 1
		fi
		;;
	--STOP)
		chk_next $*
		stop=`parse_addr_and_size "off" $2`
		if [ "$stop" = "" ]; then
			echo "'$2' format error"
			exit 1
		fi
		;;
	--fr)
		fr=`$bt_disym off die`
		if [ -z "$fr" ]; then
			echo "symbol 'die' not found"
			exit 1
		fi
		shift
		continue
		;;
	--exit-fr)
		exit_fr="1"
		shift
		continue
		;;
	*)
		usage
		exit 1
	esac
	shift 2
done
if [ -z "$exit_fr" ] && [ -z "$logdir" ]; then
	usage
	exit 1
fi
if ([ "$fr" ] || [ "$exit_fr" ]) \
   && ([ "$start" ]||[ "$stop" ]||[ "$syscall" ]||[ "$pids" ]||[ "$cmd" ]); then
	usage
	exit 1
fi

in_fr=""
cat /proc/modules|grep '^bt_on_off_mod ' > /dev/null
if [ $? -eq 0 ] && [ "`cat /proc/btrax/enable 2> /dev/null`" = "1" ]; then
	in_fr="1"
fi
if [ "$fr" ] && [ "$in_fr" ]; then
	echo "Already flight recorder mode"
	exit 1
fi
if [ "$exit_fr" ] && [ -z "$in_fr" ]; then
	echo "None flight recorder mode"
	exit 1
fi

#echo logdir $logdir
#echo pids $pids

#-------- check module and proc file existance --------
if [ -z "$ctr_loaded" ] && [ ! -f $ctr_mod ]; then
	echo "!!! Btrax kernel module(ctr_mod.ko) not installed."
	exit 1
fi
if [ -z "$exit_fr" ] && [ ! -d $relay_mnt_point ]; then
	echo "!!! Relayfs mount point '$relay_mnt_point' not found or is not a directory"
	exit 1
fi
if ([ "$start" ] || [ "$stop" ] || [ "$fr" ]) && [ ! -f $on_off_mod ]; then
	echo "!!! Btrax kernel module(bt_on_off_mod.ko) not installed."
	exit 1
fi
if [ ! -f /proc/modules ]; then
	echo "/proc/modules not found"
	exit 1
fi
if [ ! -f /proc/kallsyms ]; then
	echo "/proc/kallsyms not found"
	exit 1
fi
if [ "$cmd" ]; then
	if [ "$pids" ]; then
		echo "pid($pids) and command($cmd) both defined"
		exit 1
	fi
	if [ "$start" ] || [ "$stop" ]; then
		echo "start/stop and command($cmd) both defined"
	fi
	cmd_fullpath=`which $cmd 2> /dev/null`
	if [ ! -x "$cmd_fullpath" ]; then
		echo "command($cmd) is not executable"
		exit 1
	fi
fi

#-------- load module --------
if [ -z "$exit_fr" ]; then
	if [ -z "$ctr_loaded" ]; then
		sys_call_table_addr=0x`$bt_get_addr sys_call_table`
		if [ -z "$sys_call_table_addr" ]; then
			echo "sys_call_table not found in System.map"
			exit 1
		fi
		bt_insmod $ctr_mod addr=$sys_call_table_addr
	fi
	if [ "$relfs_mod" ]; then
		bt_insmod $relfs_mod
	fi
	mount -t relayfs relayfs $relay_mnt_point > /dev/null 2>&1
	if [ $? -ne 0 ]; then
		mount -t debugfs debugfs $relay_mnt_point > /dev/null 2>&1
	fi
	if [ "$djprobe_mod" ]; then
		bt_insmod $djprobe_mod
	fi
	add_bt_mod_mode
	$bt_load $mod_param
	if [ $? -ne 0 ]; then
		terminate $relay_mnt_point
		exit 1
	fi
fi

#-------- copy maps/modules files --------
if [ -z "$exit_fr" ]; then
	mkdir -p $logdir
	cat /proc/modules > $logdir/modules
	cat /proc/kallsyms > $logdir/kallsyms
	if [ "$pids" != "" ]; then
		echo "$pids" > /proc/btrax/pid
		for pid in `echo "$pids"|sed -e 's/,/ /g'`; do
			cat /proc/$pid/maps > $logdir/$pid.maps
		done
	fi
	chk_cpu_type
fi

#-------- ready for user-termination --------
trap '
	echo 0 > /proc/btrax/enable
	sync; sync; sync
	if [ "$start" ] && [ -z "$stop" ]; then
		$bt_read $relay_mnt_point/btrax $logdir -T &
		wait $!
	else
		kill_bt_read
	fi
	if [ -z "$fr" ]; then
		cat /proc/btrax/dropped > $logdir/dropped
	fi
	if [ "$syscall" ]; then
		cat /proc/btrax/exec_syscall > $logdir/syscall
	fi
	terminate $relay_mnt_point
	exit 0
' HUP INT QUIT TERM

#-------- start log collect --------
if [ "$start" ] && [ "$stop" ]; then
	bt_insmod $on_off_mod $start $stop enable_repeat=1
	$bt_read $relay_mnt_point/btrax $logdir &
	kill_pid=$!
	sleep 1
	ps aux|grep $kill_pid|grep -v 'grep' > /dev/null
	if [ $? -ne 0 ]; then
		terminate $relay_mnt_point
		exit 1
	fi
	while : ; do
		sleep 1
	done
elif [ "$start" ]; then
	bt_insmod $on_off_mod $start
	wait_for_any_cpu_full
	echo 0 > /proc/btrax/enable
	$bt_read $relay_mnt_point/btrax $logdir -T &
	wait $!
	cat /proc/btrax/dropped > $logdir/dropped
	terminate $relay_mnt_point
elif [ "$stop" ]; then
	bt_insmod $on_off_mod $stop
	echo 1 > /proc/btrax/enable
	wait_for_bts_disable
	$bt_read $relay_mnt_point/btrax $logdir -T &
	wait $!
	cat /proc/btrax/dropped > $logdir/dropped
	terminate $relay_mnt_point
elif [ "$fr" ]; then
	bt_insmod $on_off_mod $fr
	echo 1 > /proc/btrax/enable
elif [ "$exit_fr" ]; then
	if [ "$logdir" ]; then
		echo 0 > /proc/btrax/enable
		$bt_read $relay_mnt_point/btrax $logdir -T &
		wait $!
		cat /proc/btrax/dropped > $logdir/dropped
	fi
	terminate $relay_mnt_point
else
	$bt_read $relay_mnt_point/btrax $logdir &
	kill_pid=$!
	sleep 1
	ps aux|grep $kill_pid|grep -v 'grep' > /dev/null
	if [ $? -ne 0 ]; then
		terminate $relay_mnt_point
		exit 1
	fi
	#ps -u root|grep bt_read > bt_read.pid # for debug
	if [ "$cmd" ]; then
		if [ "$syscall_filter" ]; then
			BT_OUTPUT_DIR="$logdir" \
			__bt_execcmd -s $syscall_filter \
				-c $cmd_fullpath $cmd_arg &
		else
			BT_OUTPUT_DIR="$logdir" $cmd_fullpath $cmd_arg &
		fi
		cmd_pid=$!
	else
		if [ "$syscall_filter" ]; then
			echo $syscall_filter > /proc/btrax/filter_syscall
		fi
		echo 1 > /proc/btrax/enable
	fi
	while : ; do
		sleep 1
	done
fi
exit 0
