samples/bpf: add a samples/bpf test for BPF_TASK_FD_QUERY

This is mostly to test kprobe/uprobe which needs kernel headers.

Signed-off-by: Yonghong Song <yhs@fb.com>
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

3 files changed, 405 insertions, 0 deletions

diff --git a/samples/bpf/Makefile b/samples/bpf/Makefile
index 62d1aa1a4cf3..7dc85ed0ce4b 100644
--- a/samples/bpf/Makefile
+++ b/samples/bpf/Makefile
@@ -51,6 +51,7 @@ hostprogs-y += cpustat
 hostprogs-y += xdp_adjust_tail
 hostprogs-y += xdpsock
 hostprogs-y += xdp_fwd
+hostprogs-y += task_fd_query
 
 # Libbpf dependencies
 LIBBPF = $(TOOLS_PATH)/lib/bpf/libbpf.a
@@ -105,6 +106,7 @@ cpustat-objs := bpf_load.o cpustat_user.o
 xdp_adjust_tail-objs := xdp_adjust_tail_user.o
 xdpsock-objs := bpf_load.o xdpsock_user.o
 xdp_fwd-objs := bpf_load.o xdp_fwd_user.o
+task_fd_query-objs := bpf_load.o task_fd_query_user.o $(TRACE_HELPERS)
 
 # Tell kbuild to always build the programs
 always := $(hostprogs-y)
@@ -160,6 +162,7 @@ always += cpustat_kern.o
 always += xdp_adjust_tail_kern.o
 always += xdpsock_kern.o
 always += xdp_fwd_kern.o
+always += task_fd_query_kern.o
 
 HOSTCFLAGS += -I$(objtree)/usr/include
 HOSTCFLAGS += -I$(srctree)/tools/lib/
@@ -175,6 +178,7 @@ HOSTCFLAGS_offwaketime_user.o += -I$(srctree)/tools/lib/bpf/
 HOSTCFLAGS_spintest_user.o += -I$(srctree)/tools/lib/bpf/
 HOSTCFLAGS_trace_event_user.o += -I$(srctree)/tools/lib/bpf/
 HOSTCFLAGS_sampleip_user.o += -I$(srctree)/tools/lib/bpf/
+HOSTCFLAGS_task_fd_query_user.o += -I$(srctree)/tools/lib/bpf/
 
 HOST_LOADLIBES          += $(LIBBPF) -lelf
 HOSTLOADLIBES_tracex4           += -lrt
diff --git a/samples/bpf/task_fd_query_kern.c b/samples/bpf/task_fd_query_kern.c
new file mode 100644
index 000000000000..f4b0a9ea674d
--- /dev/null
+++ b/samples/bpf/task_fd_query_kern.c
@@ -0,0 +1,19 @@
+// SPDX-License-Identifier: GPL-2.0
+#include <linux/version.h>
+#include <linux/ptrace.h>
+#include <uapi/linux/bpf.h>
+#include "bpf_helpers.h"
+
+SEC("kprobe/blk_start_request")
+int bpf_prog1(struct pt_regs *ctx)
+{
+        return 0;
+}
+
+SEC("kretprobe/blk_account_io_completion")
+int bpf_prog2(struct pt_regs *ctx)
+{
+        return 0;
+}
+char _license[] SEC("license") = "GPL";
+u32 _version SEC("version") = LINUX_VERSION_CODE;
diff --git a/samples/bpf/task_fd_query_user.c b/samples/bpf/task_fd_query_user.c
new file mode 100644
index 000000000000..8381d792f138
--- /dev/null
+++ b/samples/bpf/task_fd_query_user.c
@@ -0,0 +1,382 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include <stdio.h>
+#include <stdlib.h>
+#include <signal.h>
+#include <unistd.h>
+#include <stdbool.h>
+#include <string.h>
+#include <stdint.h>
+#include <fcntl.h>
+#include <linux/bpf.h>
+#include <sys/ioctl.h>
+#include <sys/resource.h>
+#include <sys/types.h>
+#include <sys/stat.h>
+
+#include "libbpf.h"
+#include "bpf_load.h"
+#include "bpf_util.h"
+#include "perf-sys.h"
+#include "trace_helpers.h"
+
+#define CHECK_PERROR_RET(condition) ({                  \
+        int __ret = !!(condition);                      \
+        if (__ret) {                                    \
+                printf("FAIL: %s:\n", __func__);        \
+                perror("    ");                 \
+                return -1;                              \
+        }                                               \
+})
+
+#define CHECK_AND_RET(condition) ({                     \
+        int __ret = !!(condition);                      \
+        if (__ret)                                      \
+                return -1;                              \
+})
+
+static __u64 ptr_to_u64(void *ptr)
+{
+        return (__u64) (unsigned long) ptr;
+}
+
+#define PMU_TYPE_FILE "/sys/bus/event_source/devices/%s/type"
+static int bpf_find_probe_type(const char *event_type)
+{
+        char buf[256];
+        int fd, ret;
+
+        ret = snprintf(buf, sizeof(buf), PMU_TYPE_FILE, event_type);
+        CHECK_PERROR_RET(ret < 0 || ret >= sizeof(buf));
+
+        fd = open(buf, O_RDONLY);
+        CHECK_PERROR_RET(fd < 0);
+
+        ret = read(fd, buf, sizeof(buf));
+        close(fd);
+        CHECK_PERROR_RET(ret < 0 || ret >= sizeof(buf));
+
+        errno = 0;
+        ret = (int)strtol(buf, NULL, 10);
+        CHECK_PERROR_RET(errno);
+        return ret;
+}
+
+#define PMU_RETPROBE_FILE "/sys/bus/event_source/devices/%s/format/retprobe"
+static int bpf_get_retprobe_bit(const char *event_type)
+{
+        char buf[256];
+        int fd, ret;
+
+        ret = snprintf(buf, sizeof(buf), PMU_RETPROBE_FILE, event_type);
+        CHECK_PERROR_RET(ret < 0 || ret >= sizeof(buf));
+
+        fd = open(buf, O_RDONLY);
+        CHECK_PERROR_RET(fd < 0);
+
+        ret = read(fd, buf, sizeof(buf));
+        close(fd);
+        CHECK_PERROR_RET(ret < 0 || ret >= sizeof(buf));
+        CHECK_PERROR_RET(strlen(buf) < strlen("config:"));
+
+        errno = 0;
+        ret = (int)strtol(buf + strlen("config:"), NULL, 10);
+        CHECK_PERROR_RET(errno);
+        return ret;
+}
+
+static int test_debug_fs_kprobe(int prog_fd_idx, const char *fn_name,
+                                __u32 expected_fd_type)
+{
+        __u64 probe_offset, probe_addr;
+        __u32 len, prog_id, fd_type;
+        char buf[256];
+        int err;
+
+        len = sizeof(buf);
+        err = bpf_task_fd_query(getpid(), event_fd[prog_fd_idx], 0, buf, &len,
+                                &prog_id, &fd_type, &probe_offset,
+                                &probe_addr);
+        if (err < 0) {
+                printf("FAIL: %s, for event_fd idx %d, fn_name %s\n",
+                       __func__, prog_fd_idx, fn_name);
+                perror("    :");
+                return -1;
+        }
+        if (strcmp(buf, fn_name) != 0 ||
+            fd_type != expected_fd_type ||
+            probe_offset != 0x0 || probe_addr != 0x0) {
+                printf("FAIL: bpf_trace_event_query(event_fd[%d]):\n",
+                       prog_fd_idx);
+                printf("buf: %s, fd_type: %u, probe_offset: 0x%llx,"
+                       " probe_addr: 0x%llx\n",
+                       buf, fd_type, probe_offset, probe_addr);
+                return -1;
+        }
+        return 0;
+}
+
+static int test_nondebug_fs_kuprobe_common(const char *event_type,
+        const char *name, __u64 offset, __u64 addr, bool is_return,
+        char *buf, __u32 *buf_len, __u32 *prog_id, __u32 *fd_type,
+        __u64 *probe_offset, __u64 *probe_addr)
+{
+        int is_return_bit = bpf_get_retprobe_bit(event_type);
+        int type = bpf_find_probe_type(event_type);
+        struct perf_event_attr attr = {};
+        int fd;
+
+        if (type < 0 || is_return_bit < 0) {
+                printf("FAIL: %s incorrect type (%d) or is_return_bit (%d)\n",
+                        __func__, type, is_return_bit);
+                return -1;
+        }
+
+        attr.sample_period = 1;
+        attr.wakeup_events = 1;
+        if (is_return)
+                attr.config |= 1 << is_return_bit;
+
+        if (name) {
+                attr.config1 = ptr_to_u64((void *)name);
+                attr.config2 = offset;