samples/bpf/test_map_in_map_kern.c

0001 /*
0002  * Copyright (c) 2017 Facebook
0003  *
0004  * This program is free software; you can redistribute it and/or
0005  * modify it under the terms of version 2 of the GNU General Public
0006  * License as published by the Free Software Foundation.
0007  */
0008 #define KBUILD_MODNAME "foo"
0009 #include <linux/ptrace.h>
0010 #include <linux/version.h>
0011 #include <uapi/linux/bpf.h>
0012 #include <uapi/linux/in6.h>
0013 #include <bpf/bpf_helpers.h>
0014 #include <bpf/bpf_tracing.h>
0015 #include <bpf/bpf_core_read.h>
0016 #include "trace_common.h"
0017
0018 #define MAX_NR_PORTS 65536
0019
0020 /* map #0 */
0021 struct inner_a {
0022     __uint(type, BPF_MAP_TYPE_ARRAY);
0023     __type(key, u32);
0024     __type(value, int);
0025     __uint(max_entries, MAX_NR_PORTS);
0026 } port_a SEC(".maps");
0027
0028 /* map #1 */
0029 struct inner_h {
0030     __uint(type, BPF_MAP_TYPE_HASH);
0031     __type(key, u32);
0032     __type(value, int);
0033     __uint(max_entries, 1);
0034 } port_h SEC(".maps");
0035
0036 /* map #2 */
0037 struct {
0038     __uint(type, BPF_MAP_TYPE_HASH);
0039     __type(key, u32);
0040     __type(value, int);
0041     __uint(max_entries, 1);
0042 } reg_result_h SEC(".maps");
0043
0044 /* map #3 */
0045 struct {
0046     __uint(type, BPF_MAP_TYPE_HASH);
0047     __type(key, u32);
0048     __type(value, int);
0049     __uint(max_entries, 1);
0050 } inline_result_h SEC(".maps");
0051
0052 /* map #4 */ /* Test case #0 */
0053 struct {
0054     __uint(type, BPF_MAP_TYPE_ARRAY_OF_MAPS);
0055     __uint(max_entries, MAX_NR_PORTS);
0056     __uint(key_size, sizeof(u32));
0057     __array(values, struct inner_a); /* use inner_a as inner map */
0058 } a_of_port_a SEC(".maps");
0059
0060 /* map #5 */ /* Test case #1 */
0061 struct {
0062     __uint(type, BPF_MAP_TYPE_HASH_OF_MAPS);
0063     __uint(max_entries, 1);
0064     __uint(key_size, sizeof(u32));
0065     __array(values, struct inner_a); /* use inner_a as inner map */
0066 } h_of_port_a SEC(".maps");
0067
0068 /* map #6 */ /* Test case #2 */
0069 struct {
0070     __uint(type, BPF_MAP_TYPE_HASH_OF_MAPS);
0071     __uint(max_entries, 1);
0072     __uint(key_size, sizeof(u32));
0073     __array(values, struct inner_h); /* use inner_h as inner map */
0074 } h_of_port_h SEC(".maps");
0075
0076 static __always_inline int do_reg_lookup(void *inner_map, u32 port)
0077 {
0078     int *result;
0079
0080     result = bpf_map_lookup_elem(inner_map, &port);
0081     return result ? *result : -ENOENT;
0082 }
0083
0084 static __always_inline int do_inline_array_lookup(void *inner_map, u32 port)
0085 {
0086     int *result;
0087
0088     if (inner_map != &port_a)
0089         return -EINVAL;
0090
0091     result = bpf_map_lookup_elem(&port_a, &port);
0092     return result ? *result : -ENOENT;
0093 }
0094
0095 static __always_inline int do_inline_hash_lookup(void *inner_map, u32 port)
0096 {
0097     int *result;
0098
0099     if (inner_map != &port_h)
0100         return -EINVAL;
0101
0102     result = bpf_map_lookup_elem(&port_h, &port);
0103     return result ? *result : -ENOENT;
0104 }
0105
0106 SEC("kprobe/__sys_connect")
0107 int trace_sys_connect(struct pt_regs *ctx)
0108 {
0109     struct sockaddr_in6 *in6;
0110     u16 test_case, port, dst6[8];
0111     int addrlen, ret, inline_ret, ret_key = 0;
0112     u32 port_key;
0113     void *outer_map, *inner_map;
0114     bool inline_hash = false;
0115
0116     in6 = (struct sockaddr_in6 *)PT_REGS_PARM2_CORE(ctx);
0117     addrlen = (int)PT_REGS_PARM3_CORE(ctx);
0118
0119     if (addrlen != sizeof(*in6))
0120         return 0;
0121
0122     ret = bpf_probe_read_user(dst6, sizeof(dst6), &in6->sin6_addr);
0123     if (ret) {
0124         inline_ret = ret;
0125         goto done;
0126     }
0127
0128     if (dst6[0] != 0xdead || dst6[1] != 0xbeef)
0129         return 0;
0130
0131     test_case = dst6[7];
0132
0133     ret = bpf_probe_read_user(&port, sizeof(port), &in6->sin6_port);
0134     if (ret) {
0135         inline_ret = ret;
0136         goto done;
0137     }
0138
0139     port_key = port;
0140
0141     ret = -ENOENT;
0142     if (test_case == 0) {
0143         outer_map = &a_of_port_a;
0144     } else if (test_case == 1) {
0145         outer_map = &h_of_port_a;
0146     } else if (test_case == 2) {
0147         outer_map = &h_of_port_h;
0148     } else {
0149         ret = __LINE__;
0150         inline_ret = ret;
0151         goto done;
0152     }
0153
0154     inner_map = bpf_map_lookup_elem(outer_map, &port_key);
0155     if (!inner_map) {
0156         ret = __LINE__;
0157         inline_ret = ret;
0158         goto done;
0159     }
0160
0161     ret = do_reg_lookup(inner_map, port_key);
0162
0163     if (test_case == 0 || test_case == 1)
0164         inline_ret = do_inline_array_lookup(inner_map, port_key);
0165     else
0166         inline_ret = do_inline_hash_lookup(inner_map, port_key);
0167
0168 done:
0169     bpf_map_update_elem(&reg_result_h, &ret_key, &ret, BPF_ANY);
0170     bpf_map_update_elem(&inline_result_h, &ret_key, &inline_ret, BPF_ANY);
0171
0172     return 0;
0173 }
0174
0175 char _license[] SEC("license") = "GPL";
0176 u32 _version SEC("version") = LINUX_VERSION_CODE;