selftests/bpf: Cover stream capacity and partial read edge cases

Add coverage for the stream fixes requested on the capacity rollback
series:

- oversized bpf_stream_printk() returns -E2BIG and does not leak
  capacity for a subsequent successful write on the same program;
- bpf_prog_stream_read() returns the successfully copied prefix when
  the userspace buffer straddles an unmapped page.

Signed-off-by: Jianlin Shi <shijianlin11@foxmail.com>
Link: https://lore.kernel.org/r/c5a79efbd3cb52cc0668ab50db5379c907cfe6a6.1787492521.git.shijianlin11@foxmail.com
Signed-off-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>

2 files changed, 91 insertions, 0 deletions

diff --git a/tools/testing/selftests/bpf/prog_tests/stream.c b/tools/testing/selftests/bpf/prog_tests/stream.c
index c3cce5c292bd..e4caa60b4f9b 100644
--- a/tools/testing/selftests/bpf/prog_tests/stream.c
+++ b/tools/testing/selftests/bpf/prog_tests/stream.c
@@ -58,6 +58,69 @@ void test_stream_syscall(void)
 	stream__destroy(skel);
 }
 
+void test_stream_oversize(void)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	struct stream *skel;
+	int ret, prog_fd;
+
+	skel = stream__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "stream__open_and_load"))
+		return;
+
+	prog_fd = bpf_program__fd(skel->progs.stream_oversize);
+	ret = bpf_prog_test_run_opts(prog_fd, &opts);
+	ASSERT_OK(ret, "oversize run");
+	ASSERT_OK(opts.retval, "oversize retval");
+
+	stream__destroy(skel);
+}
+
+void test_stream_partial_read(void)
+{
+	LIBBPF_OPTS(bpf_test_run_opts, opts);
+	struct stream *skel;
+	int ret, prog_fd;
+	long page_size;
+	char *page, *buf;
+	char rest[8] = {};
+
+	skel = stream__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "stream__open_and_load"))
+		return;
+
+	prog_fd = bpf_program__fd(skel->progs.stream_syscall);
+	ret = bpf_prog_test_run_opts(prog_fd, &opts);
+	ASSERT_OK(ret, "ret");
+	ASSERT_OK(opts.retval, "retval");
+
+	page_size = sysconf(_SC_PAGESIZE);
+	page = mmap(NULL, page_size * 2, PROT_READ | PROT_WRITE,
+		    MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+	if (!ASSERT_NEQ(page, MAP_FAILED, "mmap")) {
+		stream__destroy(skel);
+		return;
+	}
+	/* Leave only the first page mapped so a straddling copy faults. */
+	if (!ASSERT_OK(munmap(page + page_size, page_size), "munmap second page")) {
+		munmap(page, page_size * 2);
+		stream__destroy(skel);
+		return;
+	}
+
+	buf = page + page_size - 1;
+	ret = bpf_prog_stream_read(prog_fd, BPF_STREAM_STDOUT, buf, 3, NULL);
+	ASSERT_EQ(ret, 1, "partial bytes");
+	ASSERT_EQ(buf[0], 'f', "first byte");
+
+	ret = bpf_prog_stream_read(prog_fd, BPF_STREAM_STDOUT, rest, sizeof(rest), NULL);
+	ASSERT_EQ(ret, 2, "remaining bytes");
+	ASSERT_OK(memcmp(rest, "oo", 2), "remaining data");
+
+	munmap(page, page_size);
+	stream__destroy(skel);
+}
+
 static void test_address(struct bpf_program *prog, unsigned long *fault_addr_p)
 {
 	LIBBPF_OPTS(bpf_test_run_opts, opts);
diff --git a/tools/testing/selftests/bpf/progs/stream.c b/tools/testing/selftests/bpf/progs/stream.c
index 8d8e53d37266..3d9c0a63fab8 100644
--- a/tools/testing/selftests/bpf/progs/stream.c
+++ b/tools/testing/selftests/bpf/progs/stream.c
@@ -36,7 +36,12 @@ struct {
 } array SEC(".maps");
 
 #define ENOSPC 28
+#define E2BIG 7
 #define _STR "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+#define _X64 "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
+/* 1024 bytes: truncated by bstr_printf, must return -E2BIG. */
+#define _BIG_STR (_X64 _X64 _X64 _X64 _X64 _X64 _X64 _X64 \
+		  _X64 _X64 _X64 _X64 _X64 _X64 _X64 _X64)
 
 int size;
 u64 fault_addr;
@@ -119,6 +124,29 @@ int stream_syscall(void *ctx)
 }
 
 SEC("syscall")
+__success __retval(0)
+int stream_oversize(void *ctx)
+{
+	int ret;
+
+	ret = bpf_stream_printk(BPF_STDOUT, _BIG_STR);
+	if (ret != -E2BIG)
+		return ret ?: 1;
+
+	/* The oversized output must not reduce the remaining stream capacity. */
+	size = 0;
+	bpf_repeat(BPF_MAX_LOOPS) {
+		ret = bpf_stream_printk(BPF_STDOUT, _STR);
+		if (ret == -ENOSPC)
+			return size == 99954 ? 0 : 1;
+		if (ret)
+			return ret;
+		size += sizeof(_STR) - 1;
+	}
+	return 1;
+}
+
+SEC("syscall")
 __arch_x86_64
 __arch_arm64
 __success __retval(0)