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 */
 18struct 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 */
 26struct 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 */
 34struct 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 */
 42struct 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 */
 50struct 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 */
 58struct 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 */
 66struct 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
 73static __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
 81static __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
 92static __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
103SEC("kprobe/sys_connect")
104int 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
165done:
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
172char _license[] SEC("license") = "GPL";
173u32 _version SEC("version") = LINUX_VERSION_CODE;