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