LXR linux/samples/bpf/test_map_in_map

   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 */
  21struct 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 */
  29struct 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 */
  37struct {
  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 */
  45struct {
  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 */
  53struct {
  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 */
  61struct {
  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 */
  69struct {
  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
  76static __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
  84static __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
  95static __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
 106SEC("kprobe/__sys_connect")
 107int trace_sys_connect(struct pt_regs *ctx)
 108{
 109        struct sockaddr_in6 *in6;
 110        u16 test_case, port, dst6[8];
 111        int addrlen, ret, inline_ret, ret_key = 0;
 112        u32 port_key;
 113        void *outer_map, *inner_map;
 114        bool inline_hash = false;
 115
 116        in6 = (struct sockaddr_in6 *)PT_REGS_PARM2_CORE(ctx);
 117        addrlen = (int)PT_REGS_PARM3_CORE(ctx);
 118
 119        if (addrlen != sizeof(*in6))
 120                return 0;
 121
 122        ret = bpf_probe_read_user(dst6, sizeof(dst6), &in6->sin6_addr);
 123        if (ret) {
 124                inline_ret = ret;
 125                goto done;
 126        }
 127
 128        if (dst6[0] != 0xdead || dst6[1] != 0xbeef)
 129                return 0;
 130
 131        test_case = dst6[7];
 132
 133        ret = bpf_probe_read_user(&port, sizeof(port), &in6->sin6_port);
 134        if (ret) {
 135                inline_ret = ret;
 136                goto done;
 137        }
 138
 139        port_key = port;
 140
 141        ret = -ENOENT;
 142        if (test_case == 0) {
 143                outer_map = &a_of_port_a;
 144        } else if (test_case == 1) {
 145                outer_map = &h_of_port_a;
 146        } else if (test_case == 2) {
 147                outer_map = &h_of_port_h;
 148        } else {
 149                ret = __LINE__;
 150                inline_ret = ret;
 151                goto done;
 152        }
 153
 154        inner_map = bpf_map_lookup_elem(outer_map, &port_key);
 155        if (!inner_map) {
 156                ret = __LINE__;
 157                inline_ret = ret;
 158                goto done;
 159        }
 160
 161        ret = do_reg_lookup(inner_map, port_key);
 162
 163        if (test_case == 0 || test_case == 1)
 164                inline_ret = do_inline_array_lookup(inner_map, port_key);
 165        else
 166                inline_ret = do_inline_hash_lookup(inner_map, port_key);
 167
 168done:
 169        bpf_map_update_elem(&reg_result_h, &ret_key, &ret, BPF_ANY);
 170        bpf_map_update_elem(&inline_result_h, &ret_key, &inline_ret, BPF_ANY);
 171
 172        return 0;
 173}
 174
 175char _license[] SEC("license") = "GPL";
 176u32 _version SEC("version") = LINUX_VERSION_CODE;
 177