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