[PATCH bpf-net 14/14] samples/bpf: extend test_cgrp2_attach2 test to use cgroup storage

[Roman Gushchin]

The test_cgrp2_attach test covers bpf cgroup attachment code well,
so let's re-use it for testing allocation/releasing of cgroup storage.

The extension is pretty straightforward: the bpf program will use
the cgroup storage to save the number of transmitted bytes.

Expected output:
  $ ./test_cgrp2_attach2
  Attached DROP prog. This ping in cgroup /foo should fail...
  ping: sendmsg: Operation not permitted
  Attached DROP prog. This ping in cgroup /foo/bar should fail...
  ping: sendmsg: Operation not permitted
  Attached PASS prog. This ping in cgroup /foo/bar should pass...
  Detached PASS from /foo/bar while DROP is attached to /foo.
  This ping in cgroup /foo/bar should fail...
  ping: sendmsg: Operation not permitted
  Attached PASS from /foo/bar and detached DROP from /foo.
  This ping in cgroup /foo/bar should pass...
  ### override:PASS
  ### multi:PASS

Signed-off-by: Roman Gushchin <g...@fb.com>
Cc: Alexei Starovoitov <a...@kernel.org>
Cc: Daniel Borkmann <dan...@iogearbox.net>
Acked-by: Martin KaFai Lau <ka...@fb.com>
---
 samples/bpf/test_cgrp2_attach2.c | 27 ++++++++++++++++++++++++++-
 1 file changed, 26 insertions(+), 1 deletion(-)

diff --git a/samples/bpf/test_cgrp2_attach2.c b/samples/bpf/test_cgrp2_attach2.c
index b453e6a161be..f682e0b8aa83 100644
--- a/samples/bpf/test_cgrp2_attach2.c
+++ b/samples/bpf/test_cgrp2_attach2.c
@@ -8,7 +8,8 @@
  *   information. The number of invocations of the program, which maps
  *   to the number of packets received, is stored to key 0. Key 1 is
  *   incremented on each iteration by the number of bytes stored in
- *   the skb.
+ *   the skb. The program also stores the number of received bytes
+ *   in the cgroup storage.
  *
  * - Attaches the new program to a cgroup using BPF_PROG_ATTACH
  *
@@ -21,12 +22,15 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <assert.h>
+#include <sys/resource.h>
+#include <sys/time.h>
 #include <unistd.h>
 
 #include <linux/bpf.h>
 #include <bpf/bpf.h>
 
 #include "bpf_insn.h"
+#include "bpf_rlimit.h"
 #include "cgroup_helpers.h"
 
 #define FOO            "/foo"
@@ -205,6 +209,8 @@ static int map_fd = -1;
 
 static int prog_load_cnt(int verdict, int val)
 {
+       int cgroup_storage_fd;
+
        if (map_fd < 0)
                map_fd = bpf_create_map(BPF_MAP_TYPE_ARRAY, 4, 8, 1, 0);
        if (map_fd < 0) {
@@ -212,6 +218,13 @@ static int prog_load_cnt(int verdict, int val)
                return -1;
        }
 
+       cgroup_storage_fd = bpf_create_map(BPF_MAP_TYPE_CGROUP_STORAGE,
+                               sizeof(struct bpf_cgroup_storage_key), 8, 0, 0);
+       if (cgroup_storage_fd < 0) {
+               printf("failed to create map '%s'\n", strerror(errno));
+               return -1;
+       }
+
        struct bpf_insn prog[] = {
                BPF_MOV32_IMM(BPF_REG_0, 0),
                BPF_STX_MEM(BPF_W, BPF_REG_10, BPF_REG_0, -4), /* *(u32 *)(fp - 
4) = r0 */
@@ -222,6 +235,11 @@ static int prog_load_cnt(int verdict, int val)
                BPF_JMP_IMM(BPF_JEQ, BPF_REG_0, 0, 2),
                BPF_MOV64_IMM(BPF_REG_1, val), /* r1 = 1 */
                BPF_RAW_INSN(BPF_STX | BPF_XADD | BPF_DW, BPF_REG_0, BPF_REG_1, 
0, 0), /* xadd r0 += r1 */
+               BPF_LD_MAP_FD(BPF_REG_1, cgroup_storage_fd),
+               BPF_MOV64_IMM(BPF_REG_2, 0),
+               BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0, 
BPF_FUNC_get_local_storage),
+               BPF_MOV64_IMM(BPF_REG_1, val),
+               BPF_RAW_INSN(BPF_STX | BPF_XADD | BPF_W, BPF_REG_0, BPF_REG_1, 
0, 0),
                BPF_MOV64_IMM(BPF_REG_0, verdict), /* r0 = verdict */
                BPF_EXIT_INSN(),
        };
@@ -237,6 +255,7 @@ static int prog_load_cnt(int verdict, int val)
                printf("Output from verifier:\n%s\n-------\n", bpf_log_buf);
                return 0;
        }
+       close(cgroup_storage_fd);
        return ret;
 }
 
@@ -414,6 +433,12 @@ static int test_multiprog(void)
 int main(int argc, char **argv)
 {
        int rc = 0;
+       struct rlimit r = {1024*1024, RLIM_INFINITY};
+
+       if (setrlimit(RLIMIT_MEMLOCK, &r)) {
+               log_err("Setrlimit(RLIMIT_MEMLOCK) failed");
+               return 1;
+       }
 
        rc = test_foo_bar();
        if (rc)
-- 
2.14.4