Re: [PATCH v2 5/5] bpf: Add an iterator selftest for bpf_sk_storage_get

[Martin KaFai Lau <kafai@fb.com>]

On Thu, Nov 19, 2020 at 05:26:54PM +0100, Florent Revest wrote:
> From: Florent Revest <revest@google.com>
> 
> The eBPF program iterates over all files and tasks. For all socket
> files, it stores the tgid of the last task it encountered with a handle
> to that socket. This is a heuristic for finding the "owner" of a socket
> similar to what's done by lsof, ss, netstat or fuser. Potentially, this
> information could be used from a cgroup_skb/*gress hook to try to
> associate network traffic with processes.
> 
> The test makes sure that a socket it created is tagged with prog_tests's
> pid.
> 
> Signed-off-by: Florent Revest <revest@google.com>
> ---
>  .../selftests/bpf/prog_tests/bpf_iter.c       | 35 +++++++++++++++++++
>  .../progs/bpf_iter_bpf_sk_storage_helpers.c   | 26 ++++++++++++++
>  2 files changed, 61 insertions(+)
> 
> diff --git a/tools/testing/selftests/bpf/prog_tests/bpf_iter.c b/tools/testing/selftests/bpf/prog_tests/bpf_iter.c
> index bb4a638f2e6f..4d0626003c03 100644
> --- a/tools/testing/selftests/bpf/prog_tests/bpf_iter.c
> +++ b/tools/testing/selftests/bpf/prog_tests/bpf_iter.c
> @@ -975,6 +975,39 @@ static void test_bpf_sk_storage_delete(void)
>  	bpf_iter_bpf_sk_storage_helpers__destroy(skel);
>  }
>  
> +/* The BPF program stores in every socket the tgid of a task owning a handle to
> + * it. The test verifies that a locally-created socket is tagged with its pid
> + */
> +static void test_bpf_sk_storage_get(void)
> +{
> +	struct bpf_iter_bpf_sk_storage_helpers *skel;
> +	int err, map_fd, val = -1;
> +	int sock_fd = -1;
> +
> +	skel = bpf_iter_bpf_sk_storage_helpers__open_and_load();
> +	if (CHECK(!skel, "bpf_iter_bpf_sk_storage_helpers__open_and_load",
> +		  "skeleton open_and_load failed\n"))
> +		return;
> +
> +	sock_fd = socket(AF_INET6, SOCK_STREAM, 0);
> +	if (CHECK(sock_fd < 0, "socket", "errno: %d\n", errno))
> +		goto out;
> +
> +	do_dummy_read(skel->progs.fill_socket_owners);
> +
> +	map_fd = bpf_map__fd(skel->maps.sk_stg_map);
> +
> +	err = bpf_map_lookup_elem(map_fd, &sock_fd, &val);
> +	CHECK(err || val != getpid(), "bpf_map_lookup_elem",
> +	      "map value wasn't set correctly (expected %d, got %d, err=%d)\n",
> +	      getpid(), val, err);
> +
> +	if (sock_fd >= 0)
> +		close(sock_fd);
> +out:
> +	bpf_iter_bpf_sk_storage_helpers__destroy(skel);
> +}
> +
>  static void test_bpf_sk_storage_map(void)
>  {
>  	DECLARE_LIBBPF_OPTS(bpf_iter_attach_opts, opts);
> @@ -1131,6 +1164,8 @@ void test_bpf_iter(void)
>  		test_bpf_sk_storage_map();
>  	if (test__start_subtest("bpf_sk_storage_delete"))
>  		test_bpf_sk_storage_delete();
> +	if (test__start_subtest("bpf_sk_storage_get"))
> +		test_bpf_sk_storage_get();
>  	if (test__start_subtest("rdonly-buf-out-of-bound"))
>  		test_rdonly_buf_out_of_bound();
>  	if (test__start_subtest("buf-neg-offset"))
> diff --git a/tools/testing/selftests/bpf/progs/bpf_iter_bpf_sk_storage_helpers.c b/tools/testing/selftests/bpf/progs/bpf_iter_bpf_sk_storage_helpers.c
> index 01ff3235e413..7206fd6f09ab 100644
> --- a/tools/testing/selftests/bpf/progs/bpf_iter_bpf_sk_storage_helpers.c
> +++ b/tools/testing/selftests/bpf/progs/bpf_iter_bpf_sk_storage_helpers.c
> @@ -21,3 +21,29 @@ int delete_bpf_sk_storage_map(struct bpf_iter__bpf_sk_storage_map *ctx)
>  
>  	return 0;
>  }
> +
> +SEC("iter/task_file")
> +int fill_socket_owners(struct bpf_iter__task_file *ctx)
> +{
> +	struct task_struct *task = ctx->task;
> +	struct file *file = ctx->file;
> +	struct socket *sock;
> +	int *sock_tgid;
> +
> +	if (!task || !file || task->tgid != task->pid)
> +		return 0;
> +
> +	sock = bpf_sock_from_file(file);
> +	if (!sock)
> +		return 0;
> +
> +	sock_tgid = bpf_sk_storage_get(&sk_stg_map, sock->sk, 0,
> +				       BPF_SK_STORAGE_GET_F_CREATE);
Does it affect all sk(s) in the system?  Can it be limited to
the sk that the test is testing?