samples/bpf/test_map_in_map_kern.c

   1 /*
   2  * Copyright (c) 2017 Facebook
   3  *
   4  * This program is free software; you can redistribute it and/or
   5  * modify it under the terms of version 2 of the GNU General Public
   6  * License as published by the Free Software Foundation.
   7  */
   8 #define KBUILD_MODNAME "foo"
   9 #include <linux/ptrace.h>
  10 #include <linux/version.h>
  11 #include <uapi/linux/bpf.h>
  12 #include <uapi/linux/in6.h>
  13 #include <bpf/bpf_helpers.h>
  14 #include "bpf_legacy.h"
  15 #include <bpf/bpf_tracing.h>
  16
  17 #define MAX_NR_PORTS 65536
  18
  19 /* map #0 */
  20 struct bpf_map_def_legacy SEC("maps") port_a = {
  21         .type = BPF_MAP_TYPE_ARRAY,
  22         .key_size = sizeof(u32),
  23         .value_size = sizeof(int),
  24         .max_entries = MAX_NR_PORTS,
  25 };
  26
  27 /* map #1 */
  28 struct bpf_map_def_legacy SEC("maps") port_h = {
  29         .type = BPF_MAP_TYPE_HASH,
  30         .key_size = sizeof(u32),
  31         .value_size = sizeof(int),
  32         .max_entries = 1,
  33 };
  34
  35 /* map #2 */
  36 struct bpf_map_def_legacy SEC("maps") reg_result_h = {
  37         .type = BPF_MAP_TYPE_HASH,
  38         .key_size = sizeof(u32),
  39         .value_size = sizeof(int),
  40         .max_entries = 1,
  41 };
  42
  43 /* map #3 */
  44 struct bpf_map_def_legacy SEC("maps") inline_result_h = {
  45         .type = BPF_MAP_TYPE_HASH,
  46         .key_size = sizeof(u32),
  47         .value_size = sizeof(int),
  48         .max_entries = 1,
  49 };
  50
  51 /* map #4 */ /* Test case #0 */
  52 struct bpf_map_def_legacy SEC("maps") a_of_port_a = {
  53         .type = BPF_MAP_TYPE_ARRAY_OF_MAPS,
  54         .key_size = sizeof(u32),
  55         .inner_map_idx = 0, /* map_fd[0] is port_a */
  56         .max_entries = MAX_NR_PORTS,
  57 };
  58
  59 /* map #5 */ /* Test case #1 */
  60 struct bpf_map_def_legacy SEC("maps") h_of_port_a = {
  61         .type = BPF_MAP_TYPE_HASH_OF_MAPS,
  62         .key_size = sizeof(u32),
  63         .inner_map_idx = 0, /* map_fd[0] is port_a */
  64         .max_entries = 1,
  65 };
  66
  67 /* map #6 */ /* Test case #2 */
  68 struct bpf_map_def_legacy SEC("maps") h_of_port_h = {
  69         .type = BPF_MAP_TYPE_HASH_OF_MAPS,
  70         .key_size = sizeof(u32),
  71         .inner_map_idx = 1, /* map_fd[1] is port_h */
  72         .max_entries = 1,
  73 };
  74
  75 static __always_inline int do_reg_lookup(void *inner_map, u32 port)
  76 {
  77         int *result;
  78
  79         result = bpf_map_lookup_elem(inner_map, &port);
  80         return result ? *result : -ENOENT;
  81 }
  82
  83 static __always_inline int do_inline_array_lookup(void *inner_map, u32 port)
  84 {
  85         int *result;
  86
  87         if (inner_map != &port_a)
  88                 return -EINVAL;
  89
  90         result = bpf_map_lookup_elem(&port_a, &port);
  91         return result ? *result : -ENOENT;
  92 }
  93
  94 static __always_inline int do_inline_hash_lookup(void *inner_map, u32 port)
  95 {
  96         int *result;
  97
  98         if (inner_map != &port_h)
  99                 return -EINVAL;
 100
 101         result = bpf_map_lookup_elem(&port_h, &port);
 102         return result ? *result : -ENOENT;
 103 }
 104
 105 SEC("kprobe/sys_connect")
 106 int trace_sys_connect(struct pt_regs *ctx)
 107 {
 108         struct sockaddr_in6 *in6;
 109         u16 test_case, port, dst6[8];
 110         int addrlen, ret, inline_ret, ret_key = 0;
 111         u32 port_key;
 112         void *outer_map, *inner_map;
 113         bool inline_hash = false;
 114
 115         in6 = (struct sockaddr_in6 *)PT_REGS_PARM2(ctx);
 116         addrlen = (int)PT_REGS_PARM3(ctx);
 117
 118         if (addrlen != sizeof(*in6))
 119                 return 0;
 120
 121         ret = bpf_probe_read_user(dst6, sizeof(dst6), &in6->sin6_addr);
 122         if (ret) {
 123                 inline_ret = ret;
 124                 goto done;
 125         }
 126
 127         if (dst6[0] != 0xdead || dst6[1] != 0xbeef)
 128                 return 0;
 129
 130         test_case = dst6[7];
 131
 132         ret = bpf_probe_read_user(&port, sizeof(port), &in6->sin6_port);
 133         if (ret) {
 134                 inline_ret = ret;
 135                 goto done;
 136         }
 137
 138         port_key = port;
 139
 140         ret = -ENOENT;
 141         if (test_case == 0) {
 142                 outer_map = &a_of_port_a;
 143         } else if (test_case == 1) {
 144                 outer_map = &h_of_port_a;
 145         } else if (test_case == 2) {
 146                 outer_map = &h_of_port_h;
 147         } else {
 148                 ret = __LINE__;
 149                 inline_ret = ret;
 150                 goto done;
 151         }
 152
 153         inner_map = bpf_map_lookup_elem(outer_map, &port_key);
 154         if (!inner_map) {
 155                 ret = __LINE__;
 156                 inline_ret = ret;
 157                 goto done;
 158         }
 159
 160         ret = do_reg_lookup(inner_map, port_key);
 161
 162         if (test_case == 0 || test_case == 1)
 163                 inline_ret = do_inline_array_lookup(inner_map, port_key);
 164         else
 165                 inline_ret = do_inline_hash_lookup(inner_map, port_key);
 166
 167 done:
 168         bpf_map_update_elem(&reg_result_h, &ret_key, &ret, BPF_ANY);
 169         bpf_map_update_elem(&inline_result_h, &ret_key, &inline_ret, BPF_ANY);
 170
 171         return 0;
 172 }
 173
 174 char _license[] SEC("license") = "GPL";
 175 u32 _version SEC("version") = LINUX_VERSION_CODE;