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

[KP Singh <kpsingh@chromium.org>]

On Fri, Nov 20, 2020 at 1:32 AM Martin KaFai Lau <kafai@fb.com> wrote:
>
> 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?

Yeah, one such way would be to set the socket storage on the socket
from userspace and then "search" for the socket in the iterator and
mark it as found in a shared global variable.