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 
do_reg_lookup(void * inner_map,u32 port)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 
do_inline_array_lookup(void * inner_map,u32 port)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 
do_inline_hash_lookup(void * inner_map,u32 port)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")
trace_sys_connect(struct pt_regs * ctx)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;
174