summaryrefslogtreecommitdiff
path: root/samples/bpf/cookie_uid_helper_example.c
blob: f6e5e58931c521f540ea91a51a33e20ea154116b (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
/* This test is a demo of using get_socket_uid and get_socket_cookie
 * helper function to do per socket based network traffic monitoring.
 * It requires iptables version higher then 1.6.1. to load pinned eBPF
 * program into the xt_bpf match.
 *
 * TEST:
 * ./run_cookie_uid_helper_example.sh
 * Then generate some traffic in variate ways. ping 0 -c 10 would work
 * but the cookie and uid in this case could both be 0. A sample output
 * with some traffic generated by web browser is shown below:
 *
 * cookie: 877, uid: 0x3e8, Pakcet Count: 20, Bytes Count: 11058
 * cookie: 132, uid: 0x0, Pakcet Count: 2, Bytes Count: 286
 * cookie: 812, uid: 0x3e8, Pakcet Count: 3, Bytes Count: 1726
 * cookie: 802, uid: 0x3e8, Pakcet Count: 2, Bytes Count: 104
 * cookie: 877, uid: 0x3e8, Pakcet Count: 20, Bytes Count: 11058
 * cookie: 831, uid: 0x3e8, Pakcet Count: 2, Bytes Count: 104
 * cookie: 0, uid: 0x0, Pakcet Count: 6, Bytes Count: 712
 * cookie: 880, uid: 0xfffe, Pakcet Count: 1, Bytes Count: 70
 *
 * Clean up: if using shell script, the script file will delete the iptables
 * rule and unmount the bpf program when exit. Else the iptables rule need
 * to be deleted by hand, see run_cookie_uid_helper_example.sh for detail.
 */

#define _GNU_SOURCE

#define offsetof(type, member)	__builtin_offsetof(type, member)
#define ARRAY_SIZE(x) (sizeof(x) / sizeof(*(x)))

#include <arpa/inet.h>
#include <errno.h>
#include <error.h>
#include <limits.h>
#include <linux/bpf.h>
#include <linux/if_ether.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
#include <bpf/bpf.h>
#include "libbpf.h"

struct stats {
	uint32_t uid;
	uint64_t packets;
	uint64_t bytes;
};

static int map_fd, prog_fd;

static void maps_create(void)
{
	map_fd = bpf_create_map(BPF_MAP_TYPE_HASH, sizeof(uint32_t),
				sizeof(struct stats), 100, 0);
	if (map_fd < 0)
		error(1, errno, "map create failed!\n");
}

static void prog_load(void)
{
	static char log_buf[1 << 16];

	struct bpf_insn prog[] = {
		/*
		 * Save sk_buff for future usage. value stored in R6 to R10 will
		 * not be reset after a bpf helper function call.
		 */
		BPF_MOV64_REG(BPF_REG_6, BPF_REG_1),
		/*
		 * pc1: BPF_FUNC_get_socket_cookie takes one parameter,
		 * R1: sk_buff
		 */
		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
				BPF_FUNC_get_socket_cookie),
		/* pc2-4: save &socketCookie to r7 for future usage*/
		BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0, -8),
		BPF_MOV64_REG(BPF_REG_7, BPF_REG_10),
		BPF_ALU64_IMM(BPF_ADD, BPF_REG_7, -8),
		/*
		 * pc5-8: set up the registers for BPF_FUNC_map_lookup_elem,
		 * it takes two parameters (R1: map_fd,  R2: &socket_cookie)
		 */
		BPF_LD_MAP_FD(BPF_REG_1, map_fd),
		BPF_MOV64_REG(BPF_REG_2, BPF_REG_7),
		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
				BPF_FUNC_map_lookup_elem),
		/*
		 * pc9. if r0 != 0x0, go to pc+14, since we have the cookie
		 * stored already
		 * Otherwise do pc10-22 to setup a new data entry.
		 */
		BPF_JMP_IMM(BPF_JNE, BPF_REG_0, 0, 14),
		BPF_MOV64_REG(BPF_REG_1, BPF_REG_6),
		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
				BPF_FUNC_get_socket_uid),
		/*
		 * Place a struct stats in the R10 stack and sequentially
		 * place the member value into the memory. Packets value
		 * is set by directly place a IMM value 1 into the stack.
		 */
		BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_0,
				-32 + offsetof(struct stats, uid)),
		BPF_ST_MEM(BPF_DW, BPF_REG_10,
				-32 + offsetof(struct stats, packets), 1),
		/*
		 * __sk_buff is a special struct used for eBPF program to
		 * directly access some sk_buff field.
		 */
		BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_6,
				offsetof(struct __sk_buff, len)),
		BPF_STX_MEM(BPF_DW, BPF_REG_10, BPF_REG_1,
				-32 + offsetof(struct stats, bytes)),
		/*
		 * add new map entry using BPF_FUNC_map_update_elem, it takes
		 * 4 parameters (R1: map_fd, R2: &socket_cookie, R3: &stats,
		 * R4: flags)
		 */
		BPF_LD_MAP_FD(BPF_REG_1, map_fd),
		BPF_MOV64_REG(BPF_REG_2, BPF_REG_7),
		BPF_MOV64_REG(BPF_REG_3, BPF_REG_10),
		BPF_ALU64_IMM(BPF_ADD, BPF_REG_3, -32),
		BPF_MOV64_IMM(BPF_REG_4, 0),
		BPF_RAW_INSN(BPF_JMP | BPF_CALL, 0, 0, 0,
				BPF_FUNC_map_update_elem),
		BPF_JMP_IMM(BPF_JA, 0, 0, 5),
		/*
		 * pc24-30 update the packet info to a exist data entry, it can
		 * be done by directly write to pointers instead of using
		 * BPF_FUNC_map_update_elem helper function
		 */
		BPF_MOV64_REG(BPF_REG_9, BPF_REG_0),
		BPF_MOV64_IMM(BPF_REG_1, 1),
		BPF_STX_XADD(BPF_DW, BPF_REG_9, BPF_REG_1,
				offsetof(struct stats, packets)),
		BPF_LDX_MEM(BPF_W, BPF_REG_1, BPF_REG_6,
				offsetof(struct __sk_buff, len)),
		BPF_STX_XADD(BPF_DW, BPF_REG_9, BPF_REG_1,
				offsetof(struct stats, bytes)),
		BPF_LDX_MEM(BPF_W, BPF_REG_0, BPF_REG_6,
				offsetof(struct __sk_buff, len)),
		BPF_EXIT_INSN(),
	};
	prog_fd = bpf_load_program(BPF_PROG_TYPE_SOCKET_FILTER, prog,
					ARRAY_SIZE(prog), "GPL", 0,
					log_buf, sizeof(log_buf));
	if (prog_fd < 0)
		error(1, errno, "failed to load prog\n%s\n", log_buf);
}

static void prog_attach_iptables(char *file)
{
	int ret;
	char rules[100];

	if (bpf_obj_pin(prog_fd, file))
		error(1, errno, "bpf_obj_pin");
	if (strlen(file) > 50) {
		printf("file path too long: %s\n", file);
		exit(1);
	}
	sprintf(rules, "iptables -A INPUT -m bpf --object-pinned %s -j ACCEPT",
		file);
	ret = system(rules);
	if (ret < 0) {
		printf("iptables rule update failed: %d/n", WEXITSTATUS(ret));
		exit(1);
	}
}

static void print_table(void)
{
	struct stats curEntry;
	uint32_t curN = UINT32_MAX;
	uint32_t nextN, res;

	while (bpf_map_get_next_key(map_fd, &curN, &nextN) > -1) {
		curN = nextN;
		res = bpf_map_lookup_elem(map_fd, &curN, &curEntry);
		if (res < 0) {
			error(1, errno, "fail to get entry value of Key: %u\n",
				curN);
		} else {
			printf("cookie: %u, uid: 0x%x, Packet Count: %lu,"
				" Bytes Count: %lu\n", curN, curEntry.uid,
				curEntry.packets, curEntry.bytes);
		}
	}
}

int main(int argc, char *argv[])
{
	if (argc > 2) {
		printf("Too many argument provided\n");
		return 1;
	} else if (argc < 2) {
		printf("Usage: %s bpfObjName\n", argv[0]);
		return 1;
	}

	maps_create();
	prog_load();
	prog_attach_iptables(argv[1]);

	while (true) {
		print_table();
		printf("\n");
		sleep(1);
	};

	return 0;
}