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 */
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/__sys_connect")
107 int 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 
168 done:
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 
175 char _license[] SEC("license") = "GPL";
176 u32 _version SEC("version") = LINUX_VERSION_CODE;
177