[PATCH v3 2/2] usdt, stapsdt: fix argument handling for multi-location probes

[PATCH v3 2/2] usdt, stapsdt: fix argument handling for multi-location probes

[Kris Van Hees]

When a USDT probe has multiple call sites in the same function, only one
of the probe programs retrieves the argument values correctly.  The
probe argument data is collected for the first underlying probe, and then
skipped for any subsequent underlying probes.

Since it is not guaranteed that the argument description is identical for
all underlying probes (although it is almost always that way), we need to
process the argument description string for each underlying probe.

Orabug: 38922360
Signed-off-by: Kris Van Hees <kris.van.hees@oracle.com>
---
 libdtrace/dt_pid.c                            |  9 ++-
 libdtrace/dt_prov_uprobe.c                    | 10 +--
 test/unittest/usdt/tst.multiloc.r             |  4 ++
 test/unittest/usdt/tst.multiloc.sh            | 69 +++++++++++++++++++
 .../usdt/tst.stapsdt-notes-bug38922360.r      |  4 ++
 .../usdt/tst.stapsdt-notes-bug38922360.sh     | 54 +++++++++++++++
 6 files changed, 140 insertions(+), 10 deletions(-)
 create mode 100644 test/unittest/usdt/tst.multiloc.r
 create mode 100755 test/unittest/usdt/tst.multiloc.sh
 create mode 100644 test/unittest/usdt/tst.stapsdt-notes-bug38922360.r
 create mode 100755 test/unittest/usdt/tst.stapsdt-notes-bug38922360.sh

diff --git a/libdtrace/dt_pid.c b/libdtrace/dt_pid.c
index 6f93123f..a379a0fe 100644
--- a/libdtrace/dt_pid.c
+++ b/libdtrace/dt_pid.c
@@ -1350,9 +1350,16 @@ dt_stapsdt_parse(dtrace_hdl_t *dtp, dt_proc_t *dpr, dtrace_probedesc_t *pdp,
 		if (strcmp(pdp->prb, "*") != 0 &&
 		    (strlen(pdp->prb) > 0 && strcmp(pdp->prb, prbname) != 0))
 			continue;
-		if (prb + strlen(prb) + 1 < dbuf + doff + nhdr.n_descsz)
+		if (prb + strlen(prb) + 1 < dbuf + doff + nhdr.n_descsz) {
+			char	*p;
+
 			psp.pps_sargv = prb + strlen(prb) + 1;
 
+			for (p = psp.pps_sargv; (p = strchr(p, '@')) != NULL;
+			     p++)
+				psp.pps_nargc++;
+		}
+
 		psp.pps_type = DTPPT_STAPSDT;
 		psp.pps_prv = prvname;
 		psp.pps_mod = mod;
diff --git a/libdtrace/dt_prov_uprobe.c b/libdtrace/dt_prov_uprobe.c
index 6d1c0f65..e575b072 100644
--- a/libdtrace/dt_prov_uprobe.c
+++ b/libdtrace/dt_prov_uprobe.c
@@ -1659,15 +1659,7 @@ static int probe_info_stap(dtrace_hdl_t *dtp, const dt_probe_t *prp,
 	if (!upp || upp->sargv == NULL)
 		goto done;
 
-	/* First count the arguments. */
-	for (p = upp->sargv; p != NULL; argc++) {
-		p = strchr(p, '@');
-		if (p++ == NULL)
-			break;
-	}
-
-	/* Record number of arguments, and allocate descriptors. */
-	upp->sargc = argc;
+	argc = upp->sargc;
 	if (argc == 0)
 		goto done;
 
diff --git a/test/unittest/usdt/tst.multiloc.r b/test/unittest/usdt/tst.multiloc.r
new file mode 100644
index 00000000..a0bbc7a8
--- /dev/null
+++ b/test/unittest/usdt/tst.multiloc.r
@@ -0,0 +1,4 @@
+main
+test:main:args main
+test:main:args main
+
diff --git a/test/unittest/usdt/tst.multiloc.sh b/test/unittest/usdt/tst.multiloc.sh
new file mode 100755
index 00000000..19d5d6c1
--- /dev/null
+++ b/test/unittest/usdt/tst.multiloc.sh
@@ -0,0 +1,69 @@
+#!/bin/bash
+#
+# Oracle Linux DTrace.
+# Copyright (c) 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 covers USDT probes firing from multiple locations in the same
+# function, verifying that argument data is correct.
+
+if [ $# != 1 ]; then
+	echo expected one argument: '<'dtrace-path'>'
+	exit 2
+fi
+
+dtrace=$1
+CFLAGS="-std=gnu89 $test_cppflags"
+LDFLAGS="$test_ldflags"
+
+DIRNAME="$tmpdir/usdt-multiloc.$$.$RANDOM"
+mkdir -p $DIRNAME
+cd $DIRNAME
+
+cat > test.c <<EOF
+#include <sys/sdt.h>
+#include <stdio.h>
+
+int
+main(int argc, char **argv)
+{
+	printf("%s\n", __func__);
+	DTRACE_PROBE1(test_prov, args, __func__);
+	DTRACE_PROBE1(test_prov, args, __func__);
+}
+EOF
+
+cat > prov.d <<EOF
+provider test_prov {
+	probe args(char *);
+};
+EOF
+
+${CC} ${CFLAGS} -c test.c
+if [ $? -ne 0 ]; then
+	echo "failed to compile test.c" >& 2
+	exit 1
+fi
+$dtrace $dt_flags -G -s prov.d test.o
+if [ $? -ne 0 ]; then
+	echo "failed to create DOF" >& 2
+	exit 1
+fi
+${CC} ${LDFLAGS} -o test test.o prov.o
+if [ $? -ne 0 ]; then
+	echo "failed to link final executable" >& 2
+	exit 1
+fi
+
+$dtrace $dt_flags -c './test arg1val' -qs /dev/stdin <<EOF
+test_prov\$target:::args
+{
+	printf("%s:%s:%s %s\n", probemod, probefunc, probename,
+	       copyinstr(arg0));
+}
+
+EOF
+status=$?
+
+exit $status
diff --git a/test/unittest/usdt/tst.stapsdt-notes-bug38922360.r b/test/unittest/usdt/tst.stapsdt-notes-bug38922360.r
new file mode 100644
index 00000000..a0bbc7a8
--- /dev/null
+++ b/test/unittest/usdt/tst.stapsdt-notes-bug38922360.r
@@ -0,0 +1,4 @@
+main
+test:main:args main
+test:main:args main
+
diff --git a/test/unittest/usdt/tst.stapsdt-notes-bug38922360.sh b/test/unittest/usdt/tst.stapsdt-notes-bug38922360.sh
new file mode 100755
index 00000000..ea0d6880
--- /dev/null
+++ b/test/unittest/usdt/tst.stapsdt-notes-bug38922360.sh
@@ -0,0 +1,54 @@
+#!/bin/bash
+#
+# Oracle Linux DTrace.
+# Copyright (c) 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 covers stapsdt probes fired by the STAP_PROBEn macros, verifying
+# that argument data is processed correctly when the same probe is used in
+# multiple locations.
+
+if [ $# != 1 ]; then
+	echo expected one argument: '<'dtrace-path'>'
+	exit 2
+fi
+
+dtrace=$1
+CFLAGS="-std=gnu89 -I${PWD}/test/unittest/usdt $test_cppflags"
+LDFLAGS="$test_ldflags"
+
+DIRNAME="$tmpdir/usdt-notes-bug38922360.$$.$RANDOM"
+mkdir -p $DIRNAME
+cd $DIRNAME
+
+cat > test.c <<EOF
+#include <sdt_notes.h>
+#include <stdio.h>
+
+int
+main(int argc, char **argv)
+{
+	printf("%s\n", __func__);
+	STAP_PROBE1(test_prov, args, __func__);
+	STAP_PROBE1(test_prov, args, __func__);
+}
+EOF
+
+${CC} ${CFLAGS} -o test test.c
+if [ $? -ne 0 ]; then
+	echo "failed to compile test.c" >& 2
+	exit 1
+fi
+
+$dtrace -c './test arg1val' -qs /dev/stdin <<EOF
+test_prov\$target:::args
+{
+	printf("%s:%s:%s %s\n", probemod, probefunc, probename,
+	       copyinstr(arg0));
+}
+
+EOF
+status=$?
+
+exit $status
-- 
2.51.0