path: root/stap-client

#!/bin/bash

# Compile server client for systemtap
#
# Copyright (C) 2008 Red Hat Inc.
#
# This file is part of systemtap, and is free software.  You can
# redistribute it and/or modify it under the terms of the GNU General
# Public License (GPL); either version 2, or (at your option) any
# later version.

# This script examines the systemtap command line and packages the files and
# information needed to execute the command. This is then sent to a trusted
# systemtap server which will process the request and return the resulting
# kernel module (if requested) and any other information generated by the
# request. If a kernel module is generated, this script will load the module
# and execute it using 'staprun', if requested.

# Catch ctrl-c and other termination signals
trap 'terminate' SIGTERM
trap 'interrupt' SIGINT
trap 'ignore_signal' SIGHUP SIGPIPE

#-----------------------------------------------------------------------------
# Helper functions.
#-----------------------------------------------------------------------------
# function: configuration
function configuration {
    tmpdir_prefix_client=stap.client
    tmpdir_prefix_server=stap.server
    avahi_service_tag=_stap._tcp
}

# function: initialization
function initialization {
    rc=0
    wd=`pwd`
    umask 0
    staprun_running=0

    # Default options settings
    p_phase=5
    v_level=0
    keep_temps=0
    b_specified=0

    # Create a temporary directory to package things in
    # Do this before parsing the command line so that there is a place
    # to put -I and -R directories.
    tmpdir_client=`mktemp -dt $tmpdir_prefix_client.XXXXXX` || \
	fatal "ERROR: cannot create temporary directory " $tmpdir_client
    tmpdir_env=`dirname $tmpdir_client`
}

# function: parse_options [ STAP-OPTIONS ]
#
# Examine the command line. We need not do much checking, but we do need to
# parse all options in order to discover the ones we're interested in.
# The server will take care of most situations and return the appropriate
# output.
#
function parse_options {
    cmdline=
    cmdline1=
    cmdline2=
    while test $# != 0
    do
	advance_p=0
	dash_seen=0

        # Start of a new token.
	first_token=$1
	until test $advance_p != 0
	do
            # Identify the next option
	    first_char=`expr "$first_token" : '\(.\).*'`
	    second_char=
	    if test $dash_seen = 0; then
		if test "$first_char" = "-"; then
		    if test "$first_token" != "-"; then
	                # It's not a lone dash, so it's an option.
			# Is it a long option (i.e. --option)?
			second_char=`expr "$first_token" : '.\(.\).*'`
			if test "$second_char" != "-"; then
	                    # It's not a lone dash, or a long option, so it's a short option string.
			    # Remove the dash.
			    first_token=`expr "$first_token" : '-\(.*\)'`
			    dash_seen=1
			    first_char=`expr "$first_token" : '\(.\).*'`
			    cmdline2="$cmdline2 -"
			fi
		    fi
		fi
		if test $dash_seen = 0; then
	            # The dash has not been seen. This is either the script file
	            # name, a long argument or an argument to be passed to the probe module.
	            # If this is the first time, and -e has not been specified,
	            # then it could be the name of the script file.
		    if test "X$second_char" = "X-"; then
			cmdline2="$cmdline2 $first_token"
		    elif test "X$e_script" = "X" -a "X$script_file" = "X"; then
			script_file=$first_token
			cmdline1="$cmdline2"
			cmdline2=
		    elif test "$first_char" != "'"; then
			cmdline2="$cmdline2 '$first_token'"
		    else
			cmdline2="$cmdline2 $first_token"
		    fi
		    advance_p=$(($advance_p + 1))
		    break
		fi
	    fi

            # We are at the start of an option. Look at the first character.
	    case $first_char in
		b)
		    b_specified=1
		    ;;
		c)
		    get_arg $first_token "$2"
		    process_c "$stap_arg"
		    ;;
		D)
		    get_arg $first_token $2
		    cmdline2="${cmdline2}D '$stap_arg'"
		    ;;
		e)
		    get_arg $first_token "$2"
		    process_e "$stap_arg"
		    ;;
		I)
		    get_arg $first_token $2
		    process_I $stap_arg
		    ;;	
		k)
		    keep_temps=1
		    ;;
		l)
		    get_arg $first_token $2
		    cmdline2="${cmdline2}l '$stap_arg'"
		    ;;
		m)
		    get_arg $first_token $2
		    process_m $stap_arg
		    ;;
		o)
		    get_arg $first_token $2
		    process_o $stap_arg
		    ;;
		p)
		    get_arg $first_token $2
		    process_p $stap_arg
		    ;;
		r)
		    get_arg $first_token $2
		    cmdline2="${cmdline2}r '$stap_arg'"
		    ;;	
		R)
		    get_arg $first_token $2
		    process_R $stap_arg
		    ;;	
		s)
		    get_arg $first_token $2
		    cmdline2="${cmdline2}s '$stap_arg'"
		    ;;	
		v)
		    v_level=$(($v_level + 1))
		    ;;
		x)
		    get_arg $first_token $2
		    cmdline2="${cmdline2}x '$stap_arg'"
		    ;;
		*)
		    # An unknown or unimportant flag. Ignore it, but pass it on to the server.
		    ;;
	    esac

	    if test $advance_p = 0; then
	        # Just another flag character. Consume it.
		cmdline2="$cmdline2$first_char"
		first_token=`expr "$first_token" : '.\(.*\)'`
		if test "X$first_token" = "X"; then
		    advance_p=$(($advance_p + 1))
		fi
	    fi
	done

        # Consume the arguments we just processed.
	while test $advance_p != 0
	do
	    shift
	    advance_p=$(($advance_p - 1))
	done
    done

    # If the script file was given and it's not '-', then replace it with its
    # client-temp-name in the command string.
    if test "X$script_file" != "X"; then
	local local_name
	if test "$script_file" != "-"; then
	    local_name=`generate_client_temp_name $script_file`
	else
	    local_name=$script_file
	fi
	cmdline="$cmdline1 script/$local_name $cmdline2"
    else
	cmdline="$cmdline1 $cmdline2"
    fi
}

# function: get_arg FIRSTWORD SECONDWORD
#
# Collect an argument to the given option
function get_arg {
    # Remove first character.
    local opt=`expr "$1" : '\(.\).*'`
    local first=`expr "$1" : '.\(.*\)'`

    # Advance to the next token, if the first one is exhausted.
    if test "X$first" = "X"; then
	shift
	advance_p=$(($advance_p + 1))
	first=$1
    fi

    test "X$first" != "X" || \
	fatal "Missing argument to -$opt"

    stap_arg="$first"
    advance_p=$(($advance_p + 1))
}

# function: process_c ARGUMENT
#
# Process the -c flag.
function process_c {
    c_cmd="$1"
    cmdline2="${cmdline2}c '$1'"
}

# function: process_e ARGUMENT
#
# Process the -e flag.
function process_e {
    # Only the first -e option is recognized and it overrides any script file name
    # which may have already been identified.
    if test "X$e_script" = "X"; then
	e_script="$1"
	if test "X$script_file" != "X"; then
	    cmdline1="$cmdline1 $script_file $cmdline2"
	    cmdline2=
	    script_file=
	fi
    fi
    cmdline2="${cmdline2}e '$1'"
}

# function: process_I ARGUMENT
#
# Process the -I flag.
function process_I {
    local local_name=`include_file_or_directory tapsets $1`
    test "X$local_name" != "X" || return
    cmdline2="${cmdline2}I 'tapsets/$local_name'"
}

# function: process_m ARGUMENT
#
# Process the -m flag.
function process_m {
    m_name="$1"
    cmdline2="${cmdline2}m '$1'"
}

# function: process_o ARGUMENT
#
# Process the -o flag.
function process_o {
    stdout_redirection="$1"
    cmdline2="${cmdline2}o '$1'"
}

# function: process_p ARGUMENT
#
# Process the -p flag.
function process_p {
    p_phase=$1
    cmdline2="${cmdline2}p '$1'"
}

# function: process_R ARGUMENT
#
# Process the -R flag.
function process_R {
    local local_name=`include_file_or_directory runtime $1`
    test "X$local_name" != "X" || return
    cmdline2="${cmdline2}R 'runtime/$local_name'"
}

# function: include_file_or_directory PREFIX NAME
#
# Include the given file or directory in the client's temporary
# tree to be sent to the server.
function include_file_or_directory {
    # Add a symbolic link of the named file or directory to our temporary directory
    local local_name=`generate_client_temp_name $2`
    mkdir -p $tmpdir_client/$1/`dirname $local_name` || \
	fatal "ERROR: could not create $tmpdir_client/$1/`dirname $local_name`"
    ln -s /$local_name $tmpdir_client/$1/$local_name || \
	fatal "ERROR: could not link $tmpdir_client/$1/$local_name to /$local_name"
    echo "$local_name"
}

# function: generate_client_temp_name NAME
#
# Generate the name to be used for the given file/directory relative to the
# client's temporary directory.
function generate_client_temp_name {
    # Transform the name into a fully qualified path name
    local full_name=`echo "$1" | sed "s,^\\\([^/]\\\),$wd/\\\\1,"`

    # The same name without the initial / or trailing /
    local local_name=`echo "$full_name" | sed 's,^/\(.*\),\1,'`
    local_name=`echo "$local_name" | sed 's,\(.*\)/$,\1,'`
    echo "$local_name"
}

# function: create_request
#
# Add information to the client's temp directory representing the request
# to the server.
function create_request {
    # Work in our temporary directory
    cd $tmpdir_client

    if test "X$script_file" != "X"; then
	if test "$script_file" = "-"; then
	    mkdir -p $tmpdir_client/script || \
		fatal "ERROR: cannot create temporary directory " $tmpdir_client/script
	    cat > $tmpdir_client/script/$script_file
	else
	    include_file_or_directory script $script_file > /dev/null
	fi
    fi

    # Add the necessary info to special files in our temporary directory.
    echo "cmdline: $cmdline" > cmdline
    echo "sysinfo: `client_sysinfo`" > sysinfo
}

# function client_sysinfo
#
# Generate the client's sysinfo and echo it to stdout
function client_sysinfo {
    if test "X$sysinfo_client" = "X"; then
	# Add some info from uname
	sysinfo_client="`uname -rvm`"
    fi
    echo "$sysinfo_client"
}

# function: package_request
#
# Package the client's temp directory into a form suitable for sending to the
# server.
function package_request {
    # Package up the temporary directory into a tar file
    cd $tmpdir_env

    local tmpdir_client_base=`basename $tmpdir_client`
    tar_client=$tmpdir_env/`mktemp $tmpdir_client_base.tgz.XXXXXX` || \
	fatal "ERROR: cannot create temporary file " $tar_client

    tar -czhf $tar_client $tmpdir_client_base || \
	fatal "ERROR: tar of request tree, $tmpdir_client, failed"
}

# function: send_request
#
# Notify the server and then send $tar_client to the server
# The protocol is:
#   client -> "request:"
#   client -> $tar_client
function send_request {
    # Send the request file.
    for ((attempt=0; $attempt < 10; ++attempt))
    do
	if nc -w10 $server $(($port+1)) < $tar_client > /dev/null 2>&1; then
	    return;
	fi
	sleep 1
    done
    fatal "ERROR: Unable to connect to server while sending request file"
}

# function: receive_response
#
# Wait for a response from the server indicating the results of our request.
function receive_response {
    # Make a place to receive the response file.
    tar_server=`mktemp -t $tmpdir_prefix_client.server.tgz.XXXXXX` || \
	fatal "ERROR: cannot create temporary file " $tar_server

    # Retrieve the file. Wait for up to 5 minutes for a response.
    for ((attempt=0; $attempt < 300; ++attempt))
    do
	if nc -d $server $(($port+1)) > $tar_server 2>/dev/null; then
	    return;
	fi
	sleep 1
    done
    fatal "ERROR: Unable to connect to server while receiving response file"
}

# function: unpack_response
#
# Unpack the tar file received from the server and make the contents available
# for printing the results and/or running 'staprun'.
function unpack_response {
    tmpdir_server=`mktemp -dt $tmpdir_prefix_client.server.XXXXXX` || \
	fatal "ERROR: cannot create temporary file " $tmpdir_server

    # Unpack the server output directory
    cd $tmpdir_server
    tar -xzf $tar_server || \
	fatal "ERROR: Unpacking of server response, $tar_server, failed"

    # Identify the server's response tree. The tar file should have expanded
    # into a single directory named to match $tmpdir_prefix_server.??????
    # which should now be the only item in the current directory.
    test "`ls | wc -l`" = 1 || \
	fatal "ERROR: Wrong number of files after expansion of server's tar file"

    tmpdir_server=`ls`
    tmpdir_server=`expr "$tmpdir_server" : "\\\($tmpdir_prefix_server\\\\.......\\\)"`

    test "X$tmpdir_server" != "X" || \
	fatal "ERROR: server tar file did not expand as expected"

    # Check the contents of the expanded directory. It should contain:
    # 1) a file called stdout
    # 2) a file called stderr
    # 3) a file called rc
    # 4) optionally a directory named to match stap??????
    local num_files=`ls $tmpdir_server | wc -l`
    test $num_files = 4 -o $num_files = 3 || \
	fatal "ERROR: Wrong number of files in server's temp directory"
    test -f $tmpdir_server/stdout || \
	fatal "ERROR: `pwd`/$tmpdir_server/stdout does not exist or is not a regular file"
    test -f $tmpdir_server/stderr || \
	fatal "ERROR: `pwd`/$tmpdir_server/stderr does not exist or is not a regular file"
    test -f $tmpdir_server/rc || \
	fatal "ERROR: `pwd`/$tmpdir_server/rc does not exist or is not a regular file"

    # See if there is a systemtap temp directory
    tmpdir_stap=`ls $tmpdir_server | grep stap`
    tmpdir_stap=`expr "$tmpdir_stap" : "\\\(stap......\\\)"`
    if test "X$tmpdir_stap" != "X"; then
	test -d $tmpdir_server/$tmpdir_stap || \
	    fatal "ERROR: `pwd`/$tmpdir_server/$tmpdir_stap is not a directory"

        # Move the systemtap temp directory to a local temp location, if -k
        # was specified.
	if test $keep_temps = 1; then
	    local local_tmpdir_stap=`mktemp -dt stapXXXXXX` || \
		fatal "ERROR: cannot create temporary directory " $local_tmpdir_stap
	    mv $tmpdir_server/$tmpdir_stap/* $local_tmpdir_stap 2>/dev/null
	    rm -fr $tmpdir_server/$tmpdir_stap

	    # Correct the name of the temp directory in the server's stderr output
	    sed -i "s,^Keeping temporary directory.*,Keeping temporary directory \"$local_tmpdir_stap\"," $tmpdir_server/stderr
	    tmpdir_stap=$local_tmpdir_stap
	else
            # Make sure we own the systemtap temp directory if we are root.
	    test $EUID = 0 && chown $EUID:$EUID $tmpdir_server/$tmpdir_stap
	    # The temp directory will be moved to here below.
	    tmpdir_stap=`pwd`/$tmpdir_stap
	fi
    fi

    # Move the contents of the server's tmpdir down one level to the
    # current directory (our local server tmpdir)
    mv $tmpdir_server/* . 2>/dev/null
    rm -fr $tmpdir_server
    tmpdir_server=`pwd`
}

# function: find_and_connect_to_server
#
# Find and establish connection with a compatible stap server.
function find_and_connect_to_server {
    # Use a temp file here instead of a pipeline so that the side effects
    # of choose_server are seen by the rest of this script.
    cd $tmpdir_client
    stap-find-servers > servers
    choose_server < servers
    rm -fr servers
}

# function: choose_server
#
# Examine each line from stdin and attempt to connect to each server
# specified until successful.
function choose_server {
    local num_servers=0
    local name
    while read name server port remain
    do
	num_servers=$(($num_servers + 1))

	if test "X$server" = "X"; then
	    fatal "ERROR: server ip address not provided"
	fi

	if test "X$port" = "X"; then
	    fatal "ERROR: server port not provided"
	fi

	if connect_to_server $server $port; then
	    return 0
	fi
    done

    if test $num_servers = 0; then
	fatal "ERROR: unable to find a server"
    fi

    fatal "ERROR: unable to connect to a server"
}

# function: connect_to_server IP PORT
#
# Establish connection with the given server
function connect_to_server {
    for ((attempt=0; $attempt < 10; ++attempt))
    do
	if echo "request:" | nc -w10 $1 $2 >/dev/null 2>&1; then
	    return 0
	fi
	sleep 1
    done

    return 1
}

# function: disconnect_from_server
#
# Disconnect from the server.
function disconnect_from_server {
    :
}

# function: process_response
#
# Write the stdout and stderr from the server to stdout and stderr respectively.
function process_response {
    # Pick up the results of running stap on the server.
    rc=`cat rc`

    # Copy the module to the current directory, if -m was specified
    if test "X$m_name" != "X"; then
	if test -f $tmpdir_stap/$m_name.ko; then
	    cp $tmpdir_stap/$m_name.ko $wd
	else
	    stream_output
	    fatal "module $tmpdir_stap/$m_name.ko does not exist"
	fi
    fi

    # Output stdout and stderr as directed
    stream_output
}

# function: stream_output
#
# Output stdout and stderr as directed
function stream_output {
    cd $tmpdir_server
    cat stderr >&2
    cat stdout
}

# function: maybe_call_staprun
#
# Call staprun using the module returned from the server, if requested.
function maybe_call_staprun {
    if test $rc != 0; then
	# stap run on the server failed, so don't bother
	return
    fi

    if test $p_phase -ge 4; then
        # There should be a systemtap temporary directory.
	if test "X$tmpdir_stap" = "X"; then
	    # OK if no script specified
	    if test "X$e_script" != "X" -o "X$script_file" != "X"; then
		fatal "ERROR: systemtap temporary directory is missing in server response"
	    fi
	    return
	fi

        # There should be a module.
	local mod_name=`ls $tmpdir_stap | grep '.ko$'`
	if test "X$mod_name" = "X"; then
	    fatal "ERROR: no module was found in $tmpdir_stap"
	fi

	if test $p_phase = 5; then
	    test $v_level -gt 0 && echo "Pass 5: starting run." >&2

	    # We have a module. Try to run it
	    # If a -c command was specified, pass it along.
	    if test "X$c_cmd" != "X"; then
		staprun_opts="-c '$c_cmd'"
	    fi

	    # The -v level will be one less than what was specified
	    # for us.
	    for ((vl=$((v_level - 1)); $vl > 0; --vl))
	    do
		staprun_opts="$staprun_opts -v"
	    done

	    # if -o was specified, pass it along
	    if test "X$stdout_redirection" != "X"; then
		staprun_opts="$staprun_opts -o $stdout_redirection"
	    fi

	    # Run it in the background and wait for it. This
	    # way any signals send to us can be caught.
	    if test $v_level -ge 2; then
		echo "running `which staprun` $staprun_opts $tmpdir_stap/`ls $tmpdir_stap | grep '.ko$'`" >&2
	    fi
	    PATH=`staprun_PATH` eval staprun "$staprun_opts" \
		$tmpdir_stap/`ls $tmpdir_stap | grep '.ko$'`
	    staprun_running=1
	    wait '%?staprun' > /dev/null 2>&1
	    rc=$?
	    staprun_running=0
	    # 127 from wait means that the job was already finished.
	    test $rc=127 && rc=0

	    # Wait until the job actually disappears so that its output is complete.
	    while jobs '%?staprun' >/dev/null 2>&1
	    do
		sleep 1
	    done

	    test $v_level -gt 0 && echo "Pass 5: run completed in 0usr/0sys/0real ms." >&2
	fi
    fi
}

# function: staprun_PATH
#
# Compute a PATH suitable for running staprun.
function staprun_PATH {
    # staprun may invoke 'stap'. So we can use the current PATH if we were
    # not invoked as 'stap' or we are not the first 'stap' on the PATH.
    local first_stap=`which stap 2>/dev/null`

    if test "X$first_stap" = "X"; then
	echo "$PATH"
	return
    fi

    if test `which $0 2>/dev/null` != $first_stap; then
	echo "$PATH"
	return
    fi

    # Otherwise, remove the PATH component where we live from the PATH
    local PATH_component=`dirname $first_stap`
    echo "$PATH" | sed "s,$PATH_component,,g"
}

# function: fatal [ MESSAGE ]
#
# Fatal error
# Prints its arguments to stderr and exits
function fatal {
    echo "$0:" "$@" >&2
    disconnect_from_server
    cleanup
    exit 1
}

# function cleanup
#
# Cleanup work files unless asked to keep them.
function cleanup {
    # Clean up.
    cd $tmpdir_env
    if test $keep_temps != 1; then
	rm -fr $tmpdir_client
	rm -f  $tar_client
	rm -f  $tar_server
	rm -fr $tmpdir_server
    fi
}

# function: terminate
#
# Terminate gracefully.
function terminate {
    # Clean up
    echo "$0: terminated by signal" >&2
    cleanup

    # Kill any running staprun job
    kill -s SIGTERM '%?staprun' 2>/dev/null

    exit 1
}

# function: interrupt
#
# Pass an interrupt (ctrl-C) to staprun
function interrupt {
    # Pass the signal on to any running staprun job
    if test $staprun_running = 1; then
	kill -s SIGINT '%?staprun' 2>/dev/null
	return
    fi

    # If staprun was not running, then exit.
    cleanup
    exit 1
}

# function: ignore_signal
#
# Called in order to ignore a signal
function ignore_signal {
    :
}

#-----------------------------------------------------------------------------
# Beginning of main line execution.
#-----------------------------------------------------------------------------
configuration
initialization
parse_options "$@"
create_request
package_request
find_and_connect_to_server
send_request
receive_response
disconnect_from_server
unpack_response
process_response
maybe_call_staprun
cleanup

exit $rc