#!/bin/bash
#
# Oracle Linux DTrace.
# Copyright (c) 2024, 2026, Oracle and/or its affiliates. All rights reserved.
# Licensed under the Universal Permissive License v 1.0 as shown at
# http://oss.oracle.com/licenses/upl.
#
# This test verifies various properties of USDT and pid probes sharing
# underlying probes.  With only one argument (as called by runtest.sh),
# it verifies that USDT and pid probes can share underlying probes.
#
# Other tests call back to this one to test arg-mapping behaviour,
# verify that usdt probes actually fire when overlapped by pid probes,
# etc.

dtrace=$1
usdt=${2:-}
mapping=${3:-}

# Set up test directory.

d=`pwd`
DIRNAME=$tmpdir/pidprobes.$$.$RANDOM
mkdir -p $DIRNAME
cd $DIRNAME

# Create test source files.

if [[ -z $mapping ]]; then
    cat > prov.d <<EOF
provider pyramid {
	probe entry(int, char, int, int);
};
EOF
else
    cat > prov.d <<EOF
provider pyramid {
	probe entry(int a, char b, int c, int d) : (int c, int d, int a, char b);
};
EOF
fi

cat > main.c <<'EOF'
#include <stdio.h>
#include "prov.h"

void foo() {
	int n = 0;

	PYRAMID_ENTRY(2, 'a', 16, 128);
	if (PYRAMID_ENTRY_ENABLED())
		n += 2;
	PYRAMID_ENTRY(4, 'b', 32, 256);
	if (PYRAMID_ENTRY_ENABLED())
		n += 8;
	printf("my result: %d\n", n);
}

int
main(int argc, char **argv)
{
	foo();
}
EOF

# Build the test program.

$dtrace $dt_flags -h -s prov.d
if [ $? -ne 0 ]; then
	echo "failed to generate header file" >&2
	exit 1
fi
$CC $test_cppflags -c main.c
if [ $? -ne 0 ]; then
	echo "failed to compile test" >&2
	exit 1
fi
$dtrace $dt_flags -G -64 -s prov.d main.o
if [ $? -ne 0 ]; then
	echo "failed to create DOF" >&2
	exit 1
fi
$CC $test_ldflags -o main main.o prov.o
if [ $? -ne 0 ]; then
	echo "failed to link final executable" >&2
	exit 1
fi

# Check that the program output is 0 when the USDT probe is not enabled.
# That is, the PYRAMID_ENTRY_ENABLED() is-enabled checks should not pass.

./main > main.out
echo "my result: 0" > main.out.expected
if ! diff -q main.out main.out.expected > /dev/null; then
	echo '"my result"' looks wrong when not using DTrace
	echo === got ===
	cat main.out
	echo === expected ===
	cat main.out.expected
	exit 1
fi

# Disassemble foo().  (simplify with --disassemble=foo)

$OBJDUMP -d main | awk '
BEGIN { use = 0 }             # start by not printing lines
use == 1 && NF == 0 { exit }  # if printing lines but hit a blank, then exit
use == 1 { print }            # print lines
/<foo>:/ { use = 1 }          # turn on printing when we hit "<foo>:" (without printing this line itself)
' > disasm_foo.txt
if [ $? -ne 0 ]; then
	echo cannot objdump main
	$OBJDUMP -d main
	exit 1
fi

# From the disassembly, get the PCs for foo()'s instructions.

pcs=`awk '{print strtonum("0x"$1)}' disasm_foo.txt`
pc0=`echo $pcs | awk '{print $1}'`

# Construct D script:  add a pid$pid::-:$pc0 probe to determine base address.
printf 'p*d$target::-:%x\n' $pc0 >> pidprobes.d
printf '{\n\tbase = uregs[R_PC] - %s;\n}\n' $pc0 >> pidprobes.d

# Construct D script:  add a pid$pid::-:$absoff probe for each PC in foo.
for pc in $pcs; do
	printf 'p*d$target::-:%x,\n' $pc >> pidprobes.d
done

# Construct D script:  add a glob for all pid and USDT pyramid probes in foo.

cat >> pidprobes.d <<'EOF'
p*d$target::foo:
{
	printf("%d %s:%s:%s:%s %x\n", pid, probeprov, probemod, probefunc, probename, uregs[R_PC] - base);
}
EOF

# Construct D script:  add a glob for all USDT pyramid probes, dumping args.

if [[ -n $usdt ]]; then
	echo 'pyramid$target::foo: {' >> pidprobes.d

	if [[ -n $mapping ]]; then
		echo 'printf("%d %s:%s:%s:%s %i %i %i %c\n", pid, probeprov, probemod, probefunc, probename, args[0], args[1], args[2], args[3]);' >> pidprobes.d
	else
		echo 'printf("%d %s:%s:%s:%s %i %c %i %i\n", pid, probeprov, probemod, probefunc, probename, args[0], args[1], args[2], args[3]);' >> pidprobes.d
	fi
	echo '}' >> pidprobes.d
fi

# Run dtrace.

$dtrace $dt_flags -q -c ./main -o dtrace.out -s pidprobes.d > main.out2
if [ $? -ne 0 ]; then
	echo "failed to run dtrace" >&2
	cat pidprobes.d
	cat main.out2
	cat dtrace.out
	exit 1
fi

# Check that the program output is 10 when the USDT probe is enabled.
# That is, the PYRAMID_ENTRY_ENABLED() is-enabled checks should pass.

echo "my result: 10" > main.out2.expected

if ! diff -q main.out2 main.out2.expected > /dev/null; then
	echo '"my result"' looks wrong when using DTrace
	echo === got ===
	cat main.out2
	echo === expected ===
	cat main.out2.expected
	exit 1
fi

# Get the reported pid.

if [ `awk 'NF != 0 { print $1 }' dtrace.out | uniq | wc -l` -ne 1 ]; then
	echo no unique pid
	cat dtrace.out
	exit 1
fi
pid=`awk 'NF != 0 { print $1 }' dtrace.out | uniq`

# From showUSDT output, get the PCs for USDT probes.  We look for output like:
#     Note usdt, type N:
#       Offset 0xoffset
#       Function Offset 0xfoffset
#       Probe pyramid::foo:entry
$d/test/utils/showUSDT main | awk '
/^ *Note usdt, type / {
	getline;
	if (!match($0, /^ *Offset *0x[0-9a-f]* *$/)) {
		print "ERROR: expect Offset";
		exit(1);
	}
	off = strtonum($2);

	getline;
	if (!match($0, /^ *Function Offset *0x[0-9a-f]* *$/)) {
		print "ERROR: expect Function Offset";
		exit(1);
	}

	getline;
	if (!match($0, /^ *Probe pyramid::foo:entry/)) {
		print "ERROR: expect Probe pyramid::foo:entry";
		exit(1);
	}

	print off, $0;
} ' > usdt_pcs.txt
if [ $? -ne 0 ]; then
	echo ERROR: showUSDT output to awk
	$d/test/utils/showUSDT main
	exit 1
fi
usdt_pcs=`awk '!/is-enabled/ { sub("0x", ""); print $1}' usdt_pcs.txt`
usdt_pcs_isenabled=`awk '/is-enabled/ { sub("0x", ""); print $1}' usdt_pcs.txt`

# We expect 2 USDT probes plus 2 is-enabled.

if [ `echo $usdt_pcs           | awk '{print NF}'` -ne 2 ]; then
	echo ERROR: expected 2 USDT regular probes but got $usdt_pcs
	exit 1
fi
if [ `echo $usdt_pcs_isenabled | awk '{print NF}'` -ne 2 ]; then
	echo ERROR: expected 2 USDT is-enabled probes but got $usdt_pcs_isenabled
	exit 1
fi

# We expect all of the USDT probe PCs to be among the PCs in objdump output.

for pc in $usdt_pcs $usdt_pcs_isenabled; do
	if echo $pcs | grep -q -vw $pc ; then
		echo ERROR: cannot find USDT PC $pc in $pcs
		exit 1
	fi
done

# Get the PC for the pid return probe.  (Just keep it in hex.)

pc_return=`awk '/'$pid' pid'$pid':main:foo:return/ { print $NF }' dtrace.out`

$OBJDUMP -d main | awk '
/^[0-9a-f]* <.*>:$/ { myfunc = $NF }         # enter a new function
/^ *'$pc_return'/ { print myfunc; exit(0) }  # report the function $pc_return is in
' > return_func.out

echo "<main>:" > return_func.out.expected    # since we use uretprobe for pid return probes, the PC will be in the caller

if ! diff -q return_func.out return_func.out.expected > /dev/null; then
	echo ERROR: return PC looks to be in the wrong function
	echo === got ===
	cat return_func.out
	echo === expected ===
	cat return_func.out.expected
	exit 1
fi

# Build up a list of expected dtrace output:
# - a blank line
# - pid entry
# - pid return
# - pid offset (relative -- that is, pid$pid:main:foo:$reloff)
# - pid offset (absolute -- that is, pid$pid:main:-:$absoff)
# - two USDT probes (ignore is-enabled probes)

echo > dtrace.out.expected
printf "$pid pid$pid:main:foo:entry %x\n" $pc0   >> dtrace.out.expected
echo   "$pid pid$pid:main:foo:return $pc_return" >> dtrace.out.expected
for pc in $pcs; do
	printf "$pid pid$pid:main:foo:%x %x\n" $(($pc - $pc0)) $pc >> dtrace.out.expected
	printf "$pid pid$pid:main:-:%x %x\n"      $pc          $pc >> dtrace.out.expected
done
echo $usdt_pcs | awk '{printf("'$pid' pyramid'$pid':main:foo:entry %x\n", $1);}' >> dtrace.out.expected
echo $usdt_pcs | awk '{printf("'$pid' pyramid'$pid':main:foo:entry %x\n", $2);}' >> dtrace.out.expected
if [[ -n $usdt ]]; then
	if [[ -z $mapping ]]; then
		echo "$pid pyramid$pid:main:foo:entry 2 a 16 128" >> dtrace.out.expected
		echo "$pid pyramid$pid:main:foo:entry 4 b 32 256" >> dtrace.out.expected
	else
		echo "$pid pyramid$pid:main:foo:entry 16 128 2 a" >> dtrace.out.expected
		echo "$pid pyramid$pid:main:foo:entry 32 256 4 b" >> dtrace.out.expected
	fi
fi

# Sort and check.

sort dtrace.out          > dtrace.out.sorted
sort dtrace.out.expected > dtrace.out.expected.sorted

if ! diff -q dtrace.out.sorted dtrace.out.expected.sorted ; then
	echo ERROR: dtrace output looks wrong
	echo === got ===
	cat dtrace.out.sorted
	echo === expected ===
	cat dtrace.out.expected.sorted
	echo === diff ===
	diff dtrace.out.sorted dtrace.out.expected.sorted
	exit 1
fi

echo success
exit 0
