#!/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.

#-----------------------------------------------------------------------------
# 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 {
    wd=`pwd`
    umask 0

    # Default options settings
    p_phase=5
    v_level=0
    keep_temps=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=
    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" : '\(.\).*'`
	    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. Remove the dash.
			first_token=`expr "$first_token" : '-\(.*\)'`
			dash_seen=1
			first_char=`expr "$first_token" : '\(.\).*'`
			cmdline="$cmdline -"
		    fi
		fi
		if test $dash_seen = 0; then
	            # The dash has not been seen. This is either the script file
	            # name or an arument 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$e_script" = "X" -a "X$script_file" = "X"; then
			script_file=$first_token
		    fi
		    advance_p=$(($advance_p + 1))
		    cmdline="$cmdline $first_token"
		    break
		fi
	    fi

            # We are at the start of an option. Look at the first character.
	    case $first_char in
		c)
		    get_arg $first_token "$2"
		    process_c "$stap_arg"
		    ;;
		D)
		    get_arg $first_token $2
		    cmdline="${cmdline}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
		    cmdline="${cmdline}l '$stap_arg'"
		    ;;
		m)
		    get_arg $first_token $2
		    cmdline="${cmdline}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
		    cmdline="${cmdline}r $stap_arg"
		    ;;	
		R)
		    get_arg $first_token $2
		    process_R $stap_arg
		    ;;	
		s)
		    get_arg $first_token $2
		    cmdline="${cmdline}s $stap_arg"
		    ;;	
		v)
		    v_level=$(($v_level + 1))
		    ;;
		x)
		    get_arg $first_token $2
		    cmdline="${cmdline}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.
		cmdline="$cmdline$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" -a "$script_file" != "-"; then
	local local_name=`generate_client_temp_name $script_file`
	cmdline=`echo $cmdline | sed s,$script_file,script/$local_name,`
    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"
    cmdline="${cmdline}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"
	script_file=
    fi
    cmdline="${cmdline}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
    cmdline="${cmdline}I $local_name"
}

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

# function: process_p ARGUMENT
#
# Process the -p flag.
function process_p {
    p_phase=$1
    cmdline="${cmdline}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
    cmdline="${cmdline}R $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 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 diectory " $tmpdir_client/script
	    cat > $tmpdir_client/script/-
	else
	    include_file_or_directory script $script_file > /dev/null
	fi
    fi

    # Add the necessary info to special files in our temporary directory. Do this
    # after linking in -I and -R directories in order to guarantee no name clashes.
    tmpfile=`mktemp cmdline.XXXXXX` || \
	fatal "ERROR: cannot create temporary file "
    echo "cmdline: $cmdline" > $tmpfile

    tmpfile=`mktemp sysinfo.XXXXXX` || \
	fatal "ERROR: cannot create temporary file " $tmpfile
    echo "sysinfo: `client_sysinfo`" > $tmpfile
}

# function client_sysinfo
#
# Generate the client's sysinfo and echo it to stdout
function client_sysinfo {
    if test "X$sysinfo_client" = "X"; then
	# Get the stap version
	stap_version=`stap -V 2>&1 | grep version`
	# Remove the number before the first slash
	stap_version=`expr "$stap_version" : '.*version [^/]*/\([0-9./]*\).*'`
	# Add some info from uname
	sysinfo_client="stap $stap_version `uname -sr`"
    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"

    tar_server=$tmpdir_env/`mktemp $tmpdir_prefix_server.tgz.XXXXXX` || \
	fatal "ERROR: cannot create temporary file " $tar_server
}

# function: send_request
#
# Notify the server and then send $tar_client to the server as $tar_server
# The protocol is:
#   client -> "request: $tmpdir_client"
#   server -> "send: $tar_server"
#   client: rsync local:$tar_client server:$tar_server
#   client -> "waiting:"
#
# $tmpdir_client is provided on the request so that the server knows what
# the tar file will expand to.
function send_request {
    echo "request: `basename $tmpdir_client`" >&3

    read <&3
    local line=$REPLY
    check_server_error $line

    local tar_dest=`expr "$line" : 'send: \([^ ]*\)$'`
    test "X$tar_dest" == "X" && \
	fatal "ERROR: destination tar file not provided"

    rsync -essh -a --delete $tar_client root@$server:$tar_dest || \
	fatal "ERROR: rsync of client request, $tar_client to $server:$tar_dest, failed"

    echo "waiting:" >&3
}

# function: receive_response
#
# Wait for a response from the server indicating the results of our request.
# The protocol is:
#   server -> "sending: remote-tar-name server-tempdir-name stap-tempdir-name"
#   client -> "OK"
function receive_response {
    read <&3
    local line=$REPLY
    check_server_error $line

    tar_dest=`expr "$line" : 'sending: \([^ ]*\) [^ ]* [^ ]*$'`
    test "X$tar_dest" == "X" && \
	fatal "ERROR: server response remote file is missing"

    tmpdir_server=`expr "$line" : 'sending: [^ ]* \([^ ]*\) [^ ]*$'`
    test "X$tmpdir_server" == "X" && \
	fatal "ERROR: server response temp dir is missing"

    tmpdir_stap=`expr "$line" : 'sending: [^ ]* [^ ]* \([^ ]*\)$'`
    test "X$tmpdir_stap" == "X" && \
	fatal "ERROR: server response stap temp dir is missing"


    # Retrieve the file
    rsync -essh -a --delete root@$server:$tar_dest $tar_server || \
	fatal "ERROR: rsync of server response, $server:$tar_dest to $tar_server, failed"
    echo "OK" >&3
}

# 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 {
    # Unpack the server output directory
    cd $tmpdir_client
    tar -xzf $tar_server || \
	fatal "ERROR: Unpacking of server response, $tar_server, failed"

    # Create a local location for the server response.
    local local_tmpdir_server_base=`basename $tar_server | sed 's,\.tgz,,'`
    local local_tmpdir_server=`mktemp -dt $local_tmpdir_server_base.XXXXXX` || \
	fatal "ERROR: cannot create temporary directory " $local_tmpdir_server

    # Move the systemtap temp directory to our 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
	tmpdir_stap=$local_tmpdir_server/$tmpdir_stap
    fi

    # Move the extracted tree to our local location
    mv $tmpdir_server/* $local_tmpdir_server
    rm -fr $tmpdir_server
    tmpdir_server=$local_tmpdir_server
}

# function: find_and_connect_to_server
#
# Find and establish connection with a compatibale stap server.
function find_and_connect_to_server {
    # Find a server
    server=`avahi-browse $avahi_service_tag --terminate -r 2>/dev/null | match_server`
    test "X$server" != "X" || \
	fatal "ERROR: cannot find a server"

    port=`expr "$server" : '[^/]*/\(.*\)'`
    server=`expr "$server" : '\([^/]*\)/.*'`

    # Open a connection to the server
    if ! exec 3<> /dev/tcp/$server/$port; then
	fatal "ERROR: cannot connect to server at /dev/tcp/$server/$port"
    fi
}

# function: match_server
#
# Find a suitable server using the avahi-browse output provided on stdin.
function match_server {
    local server_ip

    # Loop over the avahi service descriptors.
    while read
    do
	# Examine the next service descriptor
	# Is it a stap server?
	(echo $REPLY | grep -q "^=.*_stap") || continue
	
	# Get the details of the service
	local service_tag equal data
	while read service_tag equal service_data
	do
	    case $service_tag in
		'=' )
		    break ;;
		hostname )
		    server_name=`expr "$service_data" : '\[\([^]]*\)\]'`
		    ;;
		address )
		    # Sometimes (seems random), avahi-resolve-host-name resolves a local server to its
		    # hardware address rather its ip address. Keep trying until we get
		    # an ip address.
		    server_ip=`expr "$service_data" : '\[\([^]]*\)\]'`
		    local attempt
		    for ((attempt=0; $attempt < 5; ++attempt))
		    do
			server_ip=`expr "$server_ip" : '^\([0-9]*\.[0-9]*\.[0-9]*\.[0-9]*\)$'`
			if test "X$server_ip" != "X"; then
			    break
			fi
		        # Resolve the server.domain to an ip address.
			server_ip=`avahi-resolve-host-name $hostname`
			server_ip=`expr "$server_ip" : '.*	\(.*\)$'`
		    done
		    ;;
		port )
		    port=`expr "$service_data" : '\[\([^]]*\)\]'`
		    ;;
		txt )
		    sysinfo_server=`expr "$service_data" : '\[\"\([^]]*\)\"\]'`
		    sysinfo_server=`expr "$sysinfo_server" : '[^/]*/\(.*\)'`
		    ;;
		* )
		    break ;;
	    esac
	done

	# It is a stap server, but is it compatible?
	sysinfo_server="stap $sysinfo_server"
	if test "$sysinfo_server" != "`client_sysinfo`"; then
	    continue
	fi

	if test "X$server_ip" != "X"; then
	    break
	fi
    done

    echo $server_ip/$port
}

# function: disconnect_from_server
#
# Disconnect from the server.
function disconnect_from_server {
    # Close the connection to the server.
    exec 3<&-
}

# function: stream_output
#
# Write the stdout and stderr from the server to stdout and stderr respectively.
function stream_output {
    # Output stdout and stderr as directed
    cd $local_tmpdir_server
    cat ${tmpdir_server}/stderr >&2
    eval cat ${tmpdir_server}/stdout $stdout_redirection
}

# function: maybe_call_staprun
#
# Call staprun using the module returned from the server, if requested.
function maybe_call_staprun {
    if test $p_phase = 5; then
	for ((--v_level; $v_level > 0; --v_level))
	do
	    staprun_opts="$staprun_opts -v"
	done
	staprun $staprun_opts \
	    $tmpdir_stap/`ls $tmpdir_stap | grep '.ko$'`
    fi
}

# function: check_server_error SERVER_RESPONSE
#
# Check the given server response for an error message.
function check_server_error {
    echo $1 | grep -q "^ERROR:" && \
	fatal "Server:" "$@"
}

# function: fatal [ MESSAGE ]
#
# Fatal error
# Prints its arguments to stderr and exits
function fatal {
    echo $0: "$@" >&2
    cat <&3 >&2
    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
}

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

exit 0