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