[PATCH bpf-next v3 4/5] bpftool: Add option to wait for program stream output

[Kumar Kartikeya Dwivedi <memxor@gmail.com>]

bpftool prog tracelog { stdout | stderr } PROG dumps the output buffered
in a program stream and exits, which is all BPF_PROG_STREAM_READ_BY_FD
allows. Waiting for further output means rerunning the command.

Add a -w/--wait option that opens the stream with bpf_prog_stream_open()
in its default blocking mode and keeps printing output as the program
produces it. Flush each chunk as it arrives so redirected output is not
held in stdio buffers while the next read blocks. Waiting is opt-in: the
default dump keeps using BPF_PROG_STREAM_READ_BY_FD and exits once the
buffered output is drained, so existing scripts behave the same on old and
new kernels.

A stream descriptor does not keep its program alive, so bpftool drops the
program descriptor once the stream is open. Waiting then ends with EOF
when the program is unloaded and otherwise runs until interrupted, as
bpftool prog tracelog already does for the trace pipe. Exit from the
SIGINT, SIGHUP and SIGTERM handlers like that command does. A flag set by
the handler and checked before each read would miss a signal that lands
between the check and the blocking read(), leaving bpftool asleep until the
next print.

Waiting requires BPF_PROG_STREAM_OPEN. The bpf() syscall fails with EINVAL
for an unknown command, which bpftool reports as missing kernel support
rather than degrading into a dump. Document the option and offer it in the
bash completion.

Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
---
 .../bpftool/Documentation/bpftool-prog.rst    | 12 +++-
 tools/bpf/bpftool/bash-completion/bpftool     |  2 +-
 tools/bpf/bpftool/main.c                      |  7 ++-
 tools/bpf/bpftool/main.h                      |  1 +
 tools/bpf/bpftool/prog.c                      | 62 ++++++++++++++++---
 5 files changed, 73 insertions(+), 11 deletions(-)

diff --git a/tools/bpf/bpftool/Documentation/bpftool-prog.rst b/tools/bpf/bpftool/Documentation/bpftool-prog.rst
index 90fe8c61bf42..17aa0d560e32 100644
--- a/tools/bpf/bpftool/Documentation/bpftool-prog.rst
+++ b/tools/bpf/bpftool/Documentation/bpftool-prog.rst
@@ -18,7 +18,8 @@ SYNOPSIS
 
 *OPTIONS* := { |COMMON_OPTIONS| |
 { **-f** | **--bpffs** } | { **-m** | **--mapcompat** } | { **-n** | **--nomount** } |
-{ **-L** | **--use-loader** } | [ { **-S** | **--sign** } **-k** <private_key.pem> **-i** <certificate.x509> ] }
+{ **-L** | **--use-loader** } | [ { **-S** | **--sign** } **-k** <private_key.pem> **-i** <certificate.x509> ] |
+{ **-w** | **--wait** } }
 
 *COMMANDS* :=
 { **show** | **list** | **dump xlated** | **dump jited** | **pin** | **load** |
@@ -186,6 +187,11 @@ bpftool prog tracelog { stdout | stderr } *PROG*
     error messages to the standard error stream. This facility should be used
     only for debugging purposes.
 
+    By default, bpftool prints the output buffered so far and exits. With
+    **-w** or **--wait**, it keeps printing new output as the program produces
+    it, until the program is unloaded or <Ctrl+C> is hit. Waiting requires a
+    kernel that supports opening a stream as a file descriptor.
+
 bpftool prog run *PROG* data_in *FILE* [data_out *FILE* [data_size_out *L*]] [ctx_in *FILE* [ctx_out *FILE* [ctx_size_out *M*]]] [repeat *N*]
     Run BPF program *PROG* in the kernel testing infrastructure for BPF,
     meaning that the program works on the data and context provided by the
@@ -267,6 +273,10 @@ OPTIONS
     Path to the X.509 certificate file in PEM or DER format, required when
     signing.
 
+-w, --wait
+    When dumping a program stream with **bpftool prog tracelog**, wait for new
+    output instead of exiting once the buffered output has been printed.
+
 EXAMPLES
 ========
 **# bpftool prog show**
diff --git a/tools/bpf/bpftool/bash-completion/bpftool b/tools/bpf/bpftool/bash-completion/bpftool
index 9d9ced270685..a28af31937ad 100644
--- a/tools/bpf/bpftool/bash-completion/bpftool
+++ b/tools/bpf/bpftool/bash-completion/bpftool
@@ -269,7 +269,7 @@ _bpftool()
     # Deal with options
     if [[ ${words[cword]} == -* ]]; then
         local c='--version --json --pretty --bpffs --mapcompat --debug \
-            --use-loader --base-btf --sign -i -k'
+            --use-loader --base-btf --sign -i -k --wait'
         COMPREPLY=( $( compgen -W "$c" -- "$cur" ) )
         return 0
     fi
diff --git a/tools/bpf/bpftool/main.c b/tools/bpf/bpftool/main.c
index 83884b21e708..d15fec71b64e 100644
--- a/tools/bpf/bpftool/main.c
+++ b/tools/bpf/bpftool/main.c
@@ -31,6 +31,7 @@ bool block_mount;
 bool verifier_logs;
 bool relaxed_maps;
 bool use_loader;
+bool wait_output;
 struct btf *base_btf;
 struct hashmap *refs_table;
 bool sign_progs;
@@ -462,6 +463,7 @@ int main(int argc, char **argv)
 		{ "nomount",	no_argument,	NULL,	'n' },
 		{ "debug",	no_argument,	NULL,	'd' },
 		{ "use-loader",	no_argument,	NULL,	'L' },
+		{ "wait",	no_argument,	NULL,	'w' },
 		{ "sign",	no_argument,	NULL,	'S' },
 		{ "base-btf",	required_argument, NULL, 'B' },
 		{ 0 }
@@ -480,7 +482,7 @@ int main(int argc, char **argv)
 	bin_name = "bpftool";
 
 	opterr = 0;
-	while ((opt = getopt_long(argc, argv, "VhpjfLmndSi:k:B:l",
+	while ((opt = getopt_long(argc, argv, "VhpjfLmndSi:k:B:lw",
 				  options, NULL)) >= 0) {
 		switch (opt) {
 		case 'V':
@@ -530,6 +532,9 @@ int main(int argc, char **argv)
 		case 'L':
 			use_loader = true;
 			break;
+		case 'w':
+			wait_output = true;
+			break;
 		case 'S':
 			sign_progs = true;
 			use_loader = true;
diff --git a/tools/bpf/bpftool/main.h b/tools/bpf/bpftool/main.h
index 48eedcb9f6c0..fa2c877a3cbf 100644
--- a/tools/bpf/bpftool/main.h
+++ b/tools/bpf/bpftool/main.h
@@ -89,6 +89,7 @@ extern bool block_mount;
 extern bool verifier_logs;
 extern bool relaxed_maps;
 extern bool use_loader;
+extern bool wait_output;
 extern struct btf *base_btf;
 extern struct hashmap *refs_table;
 extern bool sign_progs;
diff --git a/tools/bpf/bpftool/prog.c b/tools/bpf/bpftool/prog.c
index 24e40dfab469..5e27b22444f2 100644
--- a/tools/bpf/bpftool/prog.c
+++ b/tools/bpf/bpftool/prog.c
@@ -1119,21 +1119,66 @@ enum prog_tracelog_mode {
 	TRACE_STDERR,
 };
 
+static void exit_stream(int signo)
+{
+	exit(0);
+}
+
+/* Consumes prog_fd. */
 static int
 prog_tracelog_stream(int prog_fd, enum prog_tracelog_mode mode)
 {
+	const struct sigaction act = { .sa_handler = exit_stream };
+	const int signals[] = { SIGHUP, SIGINT, SIGTERM };
 	FILE *file = mode == TRACE_STDOUT ? stdout : stderr;
 	int stream_id = mode == TRACE_STDOUT ? 1 : 2;
 	char buf[512];
-	int ret;
+	unsigned int i;
+	int fd, ret;
+
+	if (!wait_output) {
+		do {
+			ret = bpf_prog_stream_read(prog_fd, stream_id, buf, sizeof(buf), NULL);
+			if (ret > 0)
+				fwrite(buf, sizeof(buf[0]), ret, file);
+		} while (ret > 0);
+		if (ret < 0)
+			p_err("failed to read stream: %s", strerror(-ret));
+		close(prog_fd);
+		goto out;
+	}
 
-	ret = 0;
-	do {
-		ret = bpf_prog_stream_read(prog_fd, stream_id, buf, sizeof(buf), NULL);
-		if (ret > 0)
-			fwrite(buf, sizeof(buf[0]), ret, file);
-	} while (ret > 0);
+	fd = bpf_prog_stream_open(prog_fd, stream_id, NULL);
+	/*
+	 * The stream descriptor does not keep the program alive. Drop the
+	 * program reference so that reads return EOF once the program is gone.
+	 */
+	close(prog_fd);
+	if (fd == -EINVAL) {
+		/* The bpf() syscall rejects unknown commands with EINVAL. */
+		p_err("waiting for stream output is not supported by this kernel");
+		return -1;
+	}
+	if (fd < 0) {
+		p_err("failed to open stream: %s", strerror(-fd));
+		return -1;
+	}
 
+	/*
+	 * Exit from the handler like the trace pipe variant does. A flag checked
+	 * between reads would miss a signal that lands before read() blocks and
+	 * only end the loop at the next print.
+	 */
+	for (i = 0; i < ARRAY_SIZE(signals); i++)
+		sigaction(signals[i], &act, NULL);
+	while ((ret = read(fd, buf, sizeof(buf))) > 0) {
+		fwrite(buf, sizeof(buf[0]), ret, file);
+		fflush(file);
+	}
+	if (ret < 0)
+		p_err("failed to read stream: %s", strerror(errno));
+	close(fd);
+out:
 	fflush(file);
 	return ret ? -1 : 0;
 }
@@ -2659,7 +2704,8 @@ static int do_help(int argc, char **argv)
 		"       METRIC := { cycles | instructions | l1d_loads | llc_misses | itlb_misses | dtlb_misses }\n"
 		"       " HELP_SPEC_OPTIONS " |\n"
 		"                    {-f|--bpffs} | {-m|--mapcompat} | {-n|--nomount} |\n"
-		"                    {-L|--use-loader} | [ {-S|--sign } {-k} <private_key.pem> {-i} <certificate.x509> ] \n"
+		"                    {-L|--use-loader} | [ {-S|--sign } {-k} <private_key.pem> {-i} <certificate.x509> ] |\n"
+		"                    {-w|--wait} }\n"
 		"",
 		bin_name, argv[-2]);
 
-- 
2.53.0