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 "vmlinux.h"
10#include <linux/version.h>
11#include <bpf/bpf_helpers.h>
12#include <bpf/bpf_tracing.h>
13#include <bpf/bpf_core_read.h>
14
15#define MAX_NR_PORTS 65536
16
17#define EINVAL 22
18#define ENOENT 2
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("ksyscall/connect")
107int BPF_KSYSCALL(trace_sys_connect, unsigned int fd, struct sockaddr_in6 *in6, int addrlen)
108{
109	u16 test_case, port, dst6[8];
110	int ret, inline_ret, ret_key = 0;
111	u32 port_key;
112	void *outer_map, *inner_map;
113	bool inline_hash = false;
114
115	if (addrlen != sizeof(*in6))
116		return 0;
117
118	ret = bpf_probe_read_user(dst6, sizeof(dst6), &in6->sin6_addr);
119	if (ret) {
120		inline_ret = ret;
121		goto done;
122	}
123
124	if (dst6[0] != 0xdead || dst6[1] != 0xbeef)
125		return 0;
126
127	test_case = dst6[7];
128
129	ret = bpf_probe_read_user(&port, sizeof(port), &in6->sin6_port);
130	if (ret) {
131		inline_ret = ret;
132		goto done;
133	}
134
135	port_key = port;
136
137	ret = -ENOENT;
138	if (test_case == 0) {
139		outer_map = &a_of_port_a;
140	} else if (test_case == 1) {
141		outer_map = &h_of_port_a;
142	} else if (test_case == 2) {
143		outer_map = &h_of_port_h;
144	} else {
145		ret = __LINE__;
146		inline_ret = ret;
147		goto done;
148	}
149
150	inner_map = bpf_map_lookup_elem(outer_map, &port_key);
151	if (!inner_map) {
152		ret = __LINE__;
153		inline_ret = ret;
154		goto done;
155	}
156
157	ret = do_reg_lookup(inner_map, port_key);
158
159	if (test_case == 0 || test_case == 1)
160		inline_ret = do_inline_array_lookup(inner_map, port_key);
161	else
162		inline_ret = do_inline_hash_lookup(inner_map, port_key);
163
164done:
165	bpf_map_update_elem(&reg_result_h, &ret_key, &ret, BPF_ANY);
166	bpf_map_update_elem(&inline_result_h, &ret_key, &inline_ret, BPF_ANY);
167
168	return 0;
169}
170
171char _license[] SEC("license") = "GPL";
172u32 _version SEC("version") = LINUX_VERSION_CODE;
173