stap++

#!/usr/bin/env perl

use strict;
use warnings;

use Getopt::Long qw( GetOptions :config no_ignore_case);
use File::Spec ();
use File::Temp qw( tempdir );
use FindBin ();

sub quote_sh_args ($);

GetOptions("D=s@",          \(my $D_opts),
           "c=s",           \(my $shcmd),
           "d=s@",          \(my $d_opts),
           "arg=s%",        \(my $args),
           "args",          \(my $print_args),
           "dump-src",      \(my $dump_src),
           "sample-pid=i",  \(my $sample_pid),
           "exec=s",        \(my $exec_path),
           "help",          \(my $help),
           "master=i",      \(my $master_pid),
           "x=i",           \(my $pid),
           "e=s",           \(my $src),
           "v",             \(my $verbose),
           "skip-badvars",  \(my $skip_badvars),
           "I=s@",          \(my $Inc))
   or die usage();

push @$Inc, '.', "$FindBin::Bin/tapset";

#warn "Inc: @$Inc\n";

if ($^O ne 'linux') {
    die "Only linux is supported but I am on $^O.\n";
}

if ($print_args && $args && %$args) {
    die "No --arg is allowed when --args is specified.\n";
}

my $ver = `stap -V 2>&1`;
if (!defined $ver) {
    die "Systemtap not installed or its \"stap\" utility is not visible to the PATH environment: $!\n";
}

if ($ver =~ /version\s+(\d+\.\d+)/i) {
    my $v = $1;
    my $min_ver = 2.3;
    if ($v < $min_ver) {
        die "ERROR: at least systemtap $min_ver is required but found $v\n";
    }

} else {
    die "ERROR: unknown version of systemtap:\n$ver\n";
}

if ($help) {
    print usage();
    exit;
}

my %set_D_opts;
for my $opt (@$D_opts) {
    #warn "$opt\n";
    if ($opt =~ /^([_A-Z]+)=/) {
        $set_D_opts{$1} = 1;
    }
}

if (!$set_D_opts{MAXACTION}) {
    push @$D_opts, "MAXACTION=100000";
}

if (!$set_D_opts{MAXMAPENTRIES}) {
    push @$D_opts, "MAXMAPENTRIES=5000";
}

if (!$set_D_opts{MAXBACKTRACE}) {
    push @$D_opts, "MAXBACKTRACE=200";
}

my $infile;
if (!defined $src) {
    $infile = shift
        or die "No input file specified.\n";

    $src = read_src_file($infile);

} else {
    $infile = "-e";
}

my (%StdVars, %UsedLibs, %DSOs, %LibPaths, %UsedStdVars);

my @stap_opts;

if ($verbose) {
    push @stap_opts, "-v";
}

if ($d_opts) {
    for my $opt (@$d_opts) {
        push @stap_opts, '-d', $opt;
    }
}

if (defined $skip_badvars) {
    push @stap_opts, "--skip-badvars";
}

for my $opt (@$D_opts) {
    push @stap_opts, "-D$opt";
}

if (defined $exec_path) {
    if (!-f $exec_path) {
        die "$exec_path not found.\n";
    }

    $StdVars{exec_path} = $exec_path;
    push @stap_opts, '-d', $exec_path;

    open my $in, "ldd $exec_path|"
        or die "cannot run the command \"ldd $exec_path\": $!\n";
    while (<$in>) {
        if (m{\s+(\/\S+\.so(?:\.\d+)*)}) {
            my $path = $1;
            #warn "Found DSO $path.\n";
            $DSOs{$path} = 1;
        }
    }

    for my $path (sort keys %DSOs) {
        push @stap_opts, '-d', $path;
    }

    close $in;
}

if (defined $pid) {
    if (defined $master_pid) {
        die "-x <pid> and --master <pid> are exclusive.\n";
    }

    push @stap_opts, "-x", $pid;
    my $exec_file = "/proc/$pid/exe";
    if (!-f $exec_file) {
        die "Nginx process $pid is not running or ",
            "you do not have enough permissions.\n";
    }

    $StdVars{pid_ok} = gen_pid_test_condition();
    $StdVars{target} = $pid;

    if (!defined $exec_path) {
        $exec_path = readlink $exec_file;
        $StdVars{exec_path} = $exec_path;
        push @stap_opts, "-d", $exec_path;
    }

    process_dso($pid);
}

if (defined $master_pid) {
    #push @stap_opts, "-x", $master_pid;
    my $exec_file = "/proc/$master_pid/exe";
    if (!-f $exec_file) {
        die "Nginx process $master_pid is not running or ",
            "you do not have enough permissions.\n";
    }

    if (!defined $exec_path) {
        $exec_path = readlink $exec_file;
        $StdVars{exec_path} = $exec_path;
        push @stap_opts, "-d", $exec_path;
    }

    # process DSO files
    my @pids = get_child_processes($master_pid);
    if (@pids == 0) {
        die "No child processes found for $master_pid.\n";
    }

    $StdVars{pid_ok} = gen_pid_test_condition(\@pids);
    $StdVars{target} = "@pids";

    my $pid = $pids[0];
    process_dso($pid);
}

if ($sample_pid) {
    process_dso($sample_pid);
}

if (!defined $StdVars{pid_ok} && defined $exec_path) {
    $StdVars{pid_ok} = '1';
}

if (!defined $StdVars{target} && defined $exec_path) {
    $StdVars{target} = 'ANY';
}

if (!$print_args) {
    while (my ($k, $v) = each %$args) {
        #warn "getting $k => $v";
        $StdVars{"arg_$k"} = $v;
    }
}

my %used_args;

my $stap_src = process_src($infile, $src);
if ($dump_src) {
    print $stap_src;
    exit;
}

if ($print_args) {
    for my $name (sort keys %used_args) {
        my $default = $used_args{$name};
        print "\t--arg $name=VALUE";
        if (defined $default) {
            print " (default: $default)\n";

        } else {
            print "\n";
        }
    }
    exit;
}

for my $key (keys %$args) {
    if (!$UsedStdVars{"arg_$key"}) {
        my $val = $args->{$key};
        if (!defined $val) {
            $val = '';
        }
        warn "WARNING: \$^arg_$key is defined by \"--arg $key=$val\", "
             ."but never used.\n";
    }
}

my $tmpdir = tempdir("stapxx-XXXXXXXX", CLEANUP => 1);
while (my ($lib, $src) = each %UsedLibs) {
    my $outfile = "$tmpdir/$lib.stp";
    open my $out, ">$outfile"
        or die "Cannot open $outfile for writing: $!\n";
    print $out $src;
    close $out;
}

push @stap_opts, "-I", $tmpdir;

if (defined $shcmd) {
    push @stap_opts, "-c", $shcmd;
}

my $cmd = "stap " . quote_sh_args(\@stap_opts) . " -";
#warn $cmd;
open my $in, "|$cmd"
    or die "Cannot run stap: $!\n";
print $in $stap_src;
close $in;

sub eval_usr_var {
    my ($file, $usr_vars, $var) = @_;
    if (defined $usr_vars->{$var}) {
        return $usr_vars->{$var};
    }

    die "$file: line $.: Undefined user varaible \$*$var.\n";
}

sub eval_std_var {
    my ($file, $var, $trait_name, $trait_val) = @_;

    $UsedStdVars{$var} = 1;

    if (defined $StdVars{$var}) {
        return $StdVars{$var};
    }

    if (defined $trait_name) {
        #warn "trait: $trait_name";
        if ($trait_name eq 'default') {
            if ($print_args && $var =~ /^arg_(\w+)$/) {
                $used_args{$1} = $trait_val;
            }
            $StdVars{$var} = $trait_val;
            return $trait_val;

        } else {
            die "$file: line $.: unknown trait name: $trait_name\n";
        }
    }

    if ($print_args) {
        if ($var =~ /^arg_(\w+)$/) {
            $used_args{$1} = undef;
        }
        return '';
    }

    if ($var eq 'exec_path') {
        die "$file: line $.: \$^exec_path is used but neither -x <pid> ",
            "nor --exec <path> is specified.\n";

    } elsif ($var =~ /^arg_(\w+)$/) {
        die "$file: line $.: \$^$var is used but no --arg $1=VAL option is specified.\n";

    } elsif ($var =~ /^(lib\w+)_path$/) {
        my $prefix = $1;
        my $libpath = find_dso_path($prefix);
        if (!$libpath) {
            warn "$file: line $.: $prefix is not found, assuming it is statically linked.\n";
            if (!defined $StdVars{exec_path}) {
                die "No -x <pid> option is specified.\n";
            }

            $LibPaths{$prefix} = $StdVars{exec_path};
            return $StdVars{exec_path};
        }

        return $libpath;

    } else {
        die "$file: line $.: Undefined built-in variable \$^$var.\n";
    }
}

sub find_dso_path {
    my $pat = shift;

    my $path = $LibPaths{$pat};
    if ($path) {
        return $path;
    }

    my $name;
    if ($pat !~ /^lib(\S+)/) {
        die "bad pattern: $pat";
    }

    $name = $1;

    my $found_path;
    for my $path (sort keys %DSOs) {
        #warn "checking $path against $pat";
        if ($path =~ m{\blib\Q$name\E[-.\d]*\.so(?:\.\d+)*$}) {
            $LibPaths{$pat} = $path;
            $found_path = $path;
            warn "Found exact match for $pat: $path\n";
            last;
	}

        if ($path =~ m{\b(?:lib)?\Q$name\E[^/\s]*?\.so(?:\.\d+)*$}) {
            if ($found_path) {
                warn "Ignored ambiguous library $path for \"$pat\"\n";
                next;
            }

            $LibPaths{$pat} = $path;
            $found_path = $path;
        }
    }

    return $found_path;
}

sub read_src_file {
    my $infile = shift;
    open my $in, $infile
        or die "Cannot open $infile for reading: $!\n";
    my $src = do { local $/; <$in> };
    close $in;
    return $src;
}

sub use_libs {
    my ($file, $libs) = @_;
    #warn "libs: $libs";
    my @libs = split /\s*,\s*/, $libs;
    for my $lib (@libs) {
        #warn "processing $lib...";
        my $path = join("/", split /\./, $lib) . ".sxx";
        #warn $path;
        for my $dir (@$Inc) {
            my $abspath = File::Spec->catfile($dir, $path);
            #warn "Testing $abspath\n";
            #warn "exist: ", (-f $abspath) ? 1 : 0;
            if (-f $abspath) {
                my $src = read_src_file($abspath);
                $UsedLibs{$lib} = process_src($abspath, $src);
                goto next_lib;
            }
        }
        die "$file: line $.: cannot find \@use library $lib\n";
next_lib:
    }
    return "";
}

sub process_src {
    my ($file, $src) = @_;

    my %usr_vars;
    my @bits;

    my $libname_pat = qr/(?:\w+(?:\.\w+)*)/;
    # process the input file
    open my $in, '<', \$src or die $!;

    while (<$in>) {
        if ($. == 1 && /^\#!/) {
            $_ = "#!/usr/bin/env stap\n";
            next;
        }

        s{\$\^(arg_\w+)(?:\s*:(\w+)\s*\((.*?)\))?}{eval_std_var($file, $1, $2, $3)}eg;
        s{\@pfunc\s*\(\s*(\w+)\s*\)}{process("\$^exec_path").function("$1")}g;
        s{\$\^(\w+|lib[-.\w]+_path)(?:\s*:(\w+)\s*\((.*?)\))?}{eval_std_var($file, $1, $2, $3)}eg;
        s{\$\*(\w+)\s*:=\s*(\&?\@(?:cast|var)\(.*?\)(?:\[\d+\])?)}{$usr_vars{$1} = $2; ""}eg;
        s{\$\*(\w+)}{eval_usr_var($file, \%usr_vars, $1)}eg;
        s{\@use\s+($libname_pat(?:\s*,\s*$libname_pat)*)}{use_libs($file, $1)}eg;

    } continue {
        push @bits, $_;
    }

    close $in;
    return join '', @bits;
}

sub usage {
    return <<'_EOC_';
Usage:
    stap++ [optoins] [infile]

Options:
    --arg NM=VAL    Specify extra user arguments (for $^arg_NM).
    --args          Print all available arguments for --arg.
    -c CMD     start the probes, run CMD, and exit when it finishes
    -d PATH         Load debug info for the specified objects
    -D NM=VAL       Emit macro definition into generated C code.
    -e SCRIPT       Run given script.
    --exec PATH     Specify the executable file path to be traced.
    --help          Print this help.
    -I PATH         Specify the stap++ tapset library search directory.
    --master PID    Specify the master pid whose child processes are traced.
    --sample-pid PID  Sample process to inspect DSO objects to load via -d
    -v              Be verbose.
    -x PID          Sets target() to PID (also for $^exec_path and $^libxxx_path).

Examples:
    stap++ -x 12345 -e 'probe begin { println("hello") exit() }'
    stap++ -x 12345 infile.ss
_EOC_
}

sub get_child_processes {
    my $pid = shift;
    my @files = glob "/proc/[0-9]*/stat";
    my @children;
    for my $file (@files) {
        #print "file: $file\n";
        if ($file =~ m{^/proc/$pid/}) {
            next;
        }

        open my $in, $file or next;
        my $line = <$in>;
        close $in;
        if ($line =~ /^(\d+) \S+ \S+ (\d+)/) {
            my ($child, $parent) = ($1, $2);
            if ($parent eq $pid) {
                push @children, $child;
            }
        }
    }

    @children = sort { $a <=> $b } @children;
    return @children;
}

sub gen_pid_test_condition {
    my $pids = shift;
    if (!$pids) {
        return "pid() == target()";
    }

    my @c;
    for my $pid (@$pids) {
        push @c, "pid() == $pid";
    }
    return '(' . join(" || ", @c) . ')';
}

sub process_dso {
    my $pid = shift;
    my $maps_file = "/proc/$pid/maps";
    open my $in, $maps_file
        or die "Cannot open $maps_file for reading: $!\n";

    while (<$in>) {
        if (m{\S+\.so(?:\.\d+)*$}) {
            my $path = $&;
            $DSOs{$path} = 1;
            #warn "seeing $path";
        }
    }

    for my $path (sort keys %DSOs) {
        push @stap_opts, "-d", $path;
    }
}

sub quote_sh_args ($) {
    my ($args) = @_;
    for my $arg (@$args) {
       if ($arg =~ m{^[- "&%;,|?*.+=\w:/()]*$}) {
          if ($arg =~ /[ "&%;,|?*()]/) {
             $arg = "'$arg'";
          }
          next;
       }
       $arg =~ s/\\/\\\\/g;
       $arg =~ s/'/\\'/g;
       $arg =~ s/\n/\\n/g;
       $arg =~ s/\r/\\r/g;
       $arg =~ s/\t/\\t/g;
       $arg = "\$'$arg'";
    }
    return "@$args";
}