1a3f74617SMartin KaFai Lau /* Copyright (c) 2016 Facebook
2a3f74617SMartin KaFai Lau  *
3a3f74617SMartin KaFai Lau  * This program is free software; you can redistribute it and/or
4a3f74617SMartin KaFai Lau  * modify it under the terms of version 2 of the GNU General Public
5a3f74617SMartin KaFai Lau  * License as published by the Free Software Foundation.
6a3f74617SMartin KaFai Lau  */
7a3f74617SMartin KaFai Lau #include <linux/unistd.h>
8a3f74617SMartin KaFai Lau #include <linux/bpf.h>
9a3f74617SMartin KaFai Lau 
10a3f74617SMartin KaFai Lau #include <stdio.h>
11a3f74617SMartin KaFai Lau #include <stdint.h>
12a3f74617SMartin KaFai Lau #include <unistd.h>
13a3f74617SMartin KaFai Lau #include <string.h>
14a3f74617SMartin KaFai Lau #include <errno.h>
15a3f74617SMartin KaFai Lau #include <fcntl.h>
16a3f74617SMartin KaFai Lau 
172bf3e2efSJakub Kicinski #include <bpf/bpf.h>
18a3f74617SMartin KaFai Lau 
19a3f74617SMartin KaFai Lau static void usage(void)
20a3f74617SMartin KaFai Lau {
21a3f74617SMartin KaFai Lau 	printf("Usage: test_cgrp2_array_pin [...]\n");
22a3f74617SMartin KaFai Lau 	printf("       -F <file>   File to pin an BPF cgroup array\n");
23a3f74617SMartin KaFai Lau 	printf("       -U <file>   Update an already pinned BPF cgroup array\n");
24a3f74617SMartin KaFai Lau 	printf("       -v <value>  Full path of the cgroup2\n");
25a3f74617SMartin KaFai Lau 	printf("       -h          Display this help\n");
26a3f74617SMartin KaFai Lau }
27a3f74617SMartin KaFai Lau 
28a3f74617SMartin KaFai Lau int main(int argc, char **argv)
29a3f74617SMartin KaFai Lau {
30a3f74617SMartin KaFai Lau 	const char *pinned_file = NULL, *cg2 = NULL;
31a3f74617SMartin KaFai Lau 	int create_array = 1;
32a3f74617SMartin KaFai Lau 	int array_key = 0;
33a3f74617SMartin KaFai Lau 	int array_fd = -1;
34a3f74617SMartin KaFai Lau 	int cg2_fd = -1;
35a3f74617SMartin KaFai Lau 	int ret = -1;
36a3f74617SMartin KaFai Lau 	int opt;
37a3f74617SMartin KaFai Lau 
38a3f74617SMartin KaFai Lau 	while ((opt = getopt(argc, argv, "F:U:v:")) != -1) {
39a3f74617SMartin KaFai Lau 		switch (opt) {
40a3f74617SMartin KaFai Lau 		/* General args */
41a3f74617SMartin KaFai Lau 		case 'F':
42a3f74617SMartin KaFai Lau 			pinned_file = optarg;
43a3f74617SMartin KaFai Lau 			break;
44a3f74617SMartin KaFai Lau 		case 'U':
45a3f74617SMartin KaFai Lau 			pinned_file = optarg;
46a3f74617SMartin KaFai Lau 			create_array = 0;
47a3f74617SMartin KaFai Lau 			break;
48a3f74617SMartin KaFai Lau 		case 'v':
49a3f74617SMartin KaFai Lau 			cg2 = optarg;
50a3f74617SMartin KaFai Lau 			break;
51a3f74617SMartin KaFai Lau 		default:
52a3f74617SMartin KaFai Lau 			usage();
53a3f74617SMartin KaFai Lau 			goto out;
54a3f74617SMartin KaFai Lau 		}
55a3f74617SMartin KaFai Lau 	}
56a3f74617SMartin KaFai Lau 
57a3f74617SMartin KaFai Lau 	if (!cg2 || !pinned_file) {
58a3f74617SMartin KaFai Lau 		usage();
59a3f74617SMartin KaFai Lau 		goto out;
60a3f74617SMartin KaFai Lau 	}
61a3f74617SMartin KaFai Lau 
62a3f74617SMartin KaFai Lau 	cg2_fd = open(cg2, O_RDONLY);
63a3f74617SMartin KaFai Lau 	if (cg2_fd < 0) {
64a3f74617SMartin KaFai Lau 		fprintf(stderr, "open(%s,...): %s(%d)\n",
65a3f74617SMartin KaFai Lau 			cg2, strerror(errno), errno);
66a3f74617SMartin KaFai Lau 		goto out;
67a3f74617SMartin KaFai Lau 	}
68a3f74617SMartin KaFai Lau 
69a3f74617SMartin KaFai Lau 	if (create_array) {
70a3f74617SMartin KaFai Lau 		array_fd = bpf_create_map(BPF_MAP_TYPE_CGROUP_ARRAY,
71a3f74617SMartin KaFai Lau 					  sizeof(uint32_t), sizeof(uint32_t),
72a3f74617SMartin KaFai Lau 					  1, 0);
73a3f74617SMartin KaFai Lau 		if (array_fd < 0) {
74a3f74617SMartin KaFai Lau 			fprintf(stderr,
75a3f74617SMartin KaFai Lau 				"bpf_create_map(BPF_MAP_TYPE_CGROUP_ARRAY,...): %s(%d)\n",
76a3f74617SMartin KaFai Lau 				strerror(errno), errno);
77a3f74617SMartin KaFai Lau 			goto out;
78a3f74617SMartin KaFai Lau 		}
79a3f74617SMartin KaFai Lau 	} else {
80a3f74617SMartin KaFai Lau 		array_fd = bpf_obj_get(pinned_file);
81a3f74617SMartin KaFai Lau 		if (array_fd < 0) {
82a3f74617SMartin KaFai Lau 			fprintf(stderr, "bpf_obj_get(%s): %s(%d)\n",
83a3f74617SMartin KaFai Lau 				pinned_file, strerror(errno), errno);
84a3f74617SMartin KaFai Lau 			goto out;
85a3f74617SMartin KaFai Lau 		}
86a3f74617SMartin KaFai Lau 	}
87a3f74617SMartin KaFai Lau 
88d40fc181SJoe Stringer 	ret = bpf_map_update_elem(array_fd, &array_key, &cg2_fd, 0);
89a3f74617SMartin KaFai Lau 	if (ret) {
90d40fc181SJoe Stringer 		perror("bpf_map_update_elem");
91a3f74617SMartin KaFai Lau 		goto out;
92a3f74617SMartin KaFai Lau 	}
93a3f74617SMartin KaFai Lau 
94a3f74617SMartin KaFai Lau 	if (create_array) {
95a3f74617SMartin KaFai Lau 		ret = bpf_obj_pin(array_fd, pinned_file);
96a3f74617SMartin KaFai Lau 		if (ret) {
97a3f74617SMartin KaFai Lau 			fprintf(stderr, "bpf_obj_pin(..., %s): %s(%d)\n",
98a3f74617SMartin KaFai Lau 				pinned_file, strerror(errno), errno);
99a3f74617SMartin KaFai Lau 			goto out;
100a3f74617SMartin KaFai Lau 		}
101a3f74617SMartin KaFai Lau 	}
102a3f74617SMartin KaFai Lau 
103a3f74617SMartin KaFai Lau out:
104a3f74617SMartin KaFai Lau 	if (array_fd != -1)
105a3f74617SMartin KaFai Lau 		close(array_fd);
106a3f74617SMartin KaFai Lau 	if (cg2_fd != -1)
107a3f74617SMartin KaFai Lau 		close(cg2_fd);
108a3f74617SMartin KaFai Lau 	return ret;
109a3f74617SMartin KaFai Lau }
110