samples/bpf: Update sock test to allow setting mark and priority

Update sock test to set mark and priority on socket create.

Signed-off-by: David Ahern <dsahern@gmail.com>
Acked-by: Alexei Starovoitov <ast@kernel.org>
Signed-off-by: David S. Miller <davem@davemloft.net>

diff --git a/samples/bpf/test_cgrp2_sock.c b/samples/bpf/test_cgrp2_sock.c
index c3cfb23..681abbe 100644 (file)
--- a/samples/bpf/test_cgrp2_sock.c
+++ b/samples/bpf/test_cgrp2_sock.c
@@ -19,59 +19,161 @@
 #include <errno.h>
 #include <fcntl.h>
 #include <net/if.h>
+#include <inttypes.h>
 #include <linux/bpf.h>
 
 #include "libbpf.h"
 
 char bpf_log_buf[BPF_LOG_BUF_SIZE];
 
-static int prog_load(int idx)
+static int prog_load(__u32 idx, __u32 mark, __u32 prio)
 {
-       struct bpf_insn prog[] = {
+       /* save pointer to context */
+       struct bpf_insn prog_start[] = {
                BPF_MOV64_REG(BPF_REG_6, BPF_REG_1),
+       };
+       struct bpf_insn prog_end[] = {
+               BPF_MOV64_IMM(BPF_REG_0, 1), /* r0 = verdict */
+               BPF_EXIT_INSN(),
+       };
+
+       /* set sk_bound_dev_if on socket */
+       struct bpf_insn prog_dev[] = {
                BPF_MOV64_IMM(BPF_REG_3, idx),
                BPF_MOV64_IMM(BPF_REG_2, offsetof(struct bpf_sock, bound_dev_if)),
                BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3, offsetof(struct bpf_sock, bound_dev_if)),
-               BPF_MOV64_IMM(BPF_REG_0, 1), /* r0 = verdict */
-               BPF_EXIT_INSN(),
        };
-       size_t insns_cnt = sizeof(prog) / sizeof(struct bpf_insn);
 
-       return bpf_load_program(BPF_PROG_TYPE_CGROUP_SOCK, prog, insns_cnt,
+       /* set mark on socket */
+       struct bpf_insn prog_mark[] = {
+               BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
+               BPF_MOV64_IMM(BPF_REG_3, mark),
+               BPF_MOV64_IMM(BPF_REG_2, offsetof(struct bpf_sock, mark)),
+               BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3, offsetof(struct bpf_sock, mark)),
+       };
+
+       /* set priority on socket */
+       struct bpf_insn prog_prio[] = {
+               BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
+               BPF_MOV64_IMM(BPF_REG_3, prio),
+               BPF_MOV64_IMM(BPF_REG_2, offsetof(struct bpf_sock, priority)),
+               BPF_STX_MEM(BPF_W, BPF_REG_1, BPF_REG_3, offsetof(struct bpf_sock, priority)),
+       };
+
+       struct bpf_insn *prog;
+       size_t insns_cnt;
+       void *p;
+       int ret;
+
+       insns_cnt = sizeof(prog_start) + sizeof(prog_end);
+       if (idx)
+               insns_cnt += sizeof(prog_dev);
+
+       if (mark)
+               insns_cnt += sizeof(prog_mark);
+
+       if (prio)
+               insns_cnt += sizeof(prog_prio);
+
+       p = prog = malloc(insns_cnt);
+       if (!prog) {
+               fprintf(stderr, "Failed to allocate memory for instructions\n");
+               return EXIT_FAILURE;
+       }
+
+       memcpy(p, prog_start, sizeof(prog_start));
+       p += sizeof(prog_start);
+
+       if (idx) {
+               memcpy(p, prog_dev, sizeof(prog_dev));
+               p += sizeof(prog_dev);
+       }
+
+       if (mark) {
+               memcpy(p, prog_mark, sizeof(prog_mark));
+               p += sizeof(prog_mark);
+       }
+
+       if (prio) {
+               memcpy(p, prog_prio, sizeof(prog_prio));
+               p += sizeof(prog_prio);
+       }
+
+       memcpy(p, prog_end, sizeof(prog_end));
+       p += sizeof(prog_end);
+
+       insns_cnt /= sizeof(struct bpf_insn);
+
+       ret = bpf_load_program(BPF_PROG_TYPE_CGROUP_SOCK, prog, insns_cnt,
                                "GPL", 0, bpf_log_buf, BPF_LOG_BUF_SIZE);
+
+       free(prog);
+
+       return ret;
 }
 
 static int usage(const char *argv0)
 {
-       printf("Usage: %s cg-path device-index\n", argv0);
+       printf("Usage: %s -b bind-to-dev -m mark -p prio cg-path\n", argv0);
        return EXIT_FAILURE;
 }
 
 int main(int argc, char **argv)
 {
+       __u32 idx = 0, mark = 0, prio = 0;
+       const char *cgrp_path = NULL;
        int cg_fd, prog_fd, ret;
-       unsigned int idx;
+       int rc;
+
+       while ((rc = getopt(argc, argv, "b:m:p:")) != -1) {
+               switch (rc) {
+               case 'b':
+                       idx = if_nametoindex(optarg);
+                       if (!idx) {
+                               idx = strtoumax(optarg, NULL, 0);
+                               if (!idx) {
+                                       printf("Invalid device name\n");
+                                       return EXIT_FAILURE;
+                               }
+                       }
+                       break;
+               case 'm':
+                       mark = strtoumax(optarg, NULL, 0);
+                       break;
+               case 'p':
+                       prio = strtoumax(optarg, NULL, 0);
+                       break;
+               default:
+                       return usage(argv[0]);
+               }
+       }
 
-       if (argc < 2)
+       if (optind == argc)
                return usage(argv[0]);
 
-       idx = if_nametoindex(argv[2]);
-       if (!idx) {
-               printf("Invalid device name\n");
+       cgrp_path = argv[optind];
+       if (!cgrp_path) {
+               fprintf(stderr, "cgroup path not given\n");
                return EXIT_FAILURE;
        }
 
-       cg_fd = open(argv[1], O_DIRECTORY | O_RDONLY);
+       if (!idx && !mark && !prio) {
+               fprintf(stderr,
+                       "One of device, mark or priority must be given\n");
+               return EXIT_FAILURE;
+       }
+
+       cg_fd = open(cgrp_path, O_DIRECTORY | O_RDONLY);
        if (cg_fd < 0) {
                printf("Failed to open cgroup path: '%s'\n", strerror(errno));
                return EXIT_FAILURE;
        }
 
-       prog_fd = prog_load(idx);
-       printf("Output from kernel verifier:\n%s\n-------\n", bpf_log_buf);
-
+       prog_fd = prog_load(idx, mark, prio);
        if (prog_fd < 0) {
                printf("Failed to load prog: '%s'\n", strerror(errno));
+               printf("Output from kernel verifier:\n%s\n-------\n",
+                      bpf_log_buf);
                return EXIT_FAILURE;
        }
 

diff --git a/samples/bpf/test_cgrp2_sock.sh b/samples/bpf/test_cgrp2_sock.sh
index 925fd46..1153c33 100755 (executable)
--- a/samples/bpf/test_cgrp2_sock.sh
+++ b/samples/bpf/test_cgrp2_sock.sh
@@ -20,7 +20,7 @@ function attach_bpf {
        mkdir -p /tmp/cgroupv2
        mount -t cgroup2 none /tmp/cgroupv2
        mkdir -p /tmp/cgroupv2/foo
-       test_cgrp2_sock /tmp/cgroupv2/foo foo
+       test_cgrp2_sock -b foo /tmp/cgroupv2/foo
        echo $$ >> /tmp/cgroupv2/foo/cgroup.procs
 }