162306a36Sopenharmony_ci// SPDX-License-Identifier: GPL-2.0-only
262306a36Sopenharmony_ci/* Copyright (c) 2016 Facebook
362306a36Sopenharmony_ci */
462306a36Sopenharmony_ci#include <linux/unistd.h>
562306a36Sopenharmony_ci#include <linux/bpf.h>
662306a36Sopenharmony_ci
762306a36Sopenharmony_ci#include <stdio.h>
862306a36Sopenharmony_ci#include <stdint.h>
962306a36Sopenharmony_ci#include <unistd.h>
1062306a36Sopenharmony_ci#include <string.h>
1162306a36Sopenharmony_ci#include <errno.h>
1262306a36Sopenharmony_ci#include <fcntl.h>
1362306a36Sopenharmony_ci
1462306a36Sopenharmony_ci#include <bpf/bpf.h>
1562306a36Sopenharmony_ci
1662306a36Sopenharmony_cistatic void usage(void)
1762306a36Sopenharmony_ci{
1862306a36Sopenharmony_ci	printf("Usage: test_cgrp2_array_pin [...]\n");
1962306a36Sopenharmony_ci	printf("       -F <file>   File to pin an BPF cgroup array\n");
2062306a36Sopenharmony_ci	printf("       -U <file>   Update an already pinned BPF cgroup array\n");
2162306a36Sopenharmony_ci	printf("       -v <value>  Full path of the cgroup2\n");
2262306a36Sopenharmony_ci	printf("       -h          Display this help\n");
2362306a36Sopenharmony_ci}
2462306a36Sopenharmony_ci
2562306a36Sopenharmony_ciint main(int argc, char **argv)
2662306a36Sopenharmony_ci{
2762306a36Sopenharmony_ci	const char *pinned_file = NULL, *cg2 = NULL;
2862306a36Sopenharmony_ci	int create_array = 1;
2962306a36Sopenharmony_ci	int array_key = 0;
3062306a36Sopenharmony_ci	int array_fd = -1;
3162306a36Sopenharmony_ci	int cg2_fd = -1;
3262306a36Sopenharmony_ci	int ret = -1;
3362306a36Sopenharmony_ci	int opt;
3462306a36Sopenharmony_ci
3562306a36Sopenharmony_ci	while ((opt = getopt(argc, argv, "F:U:v:")) != -1) {
3662306a36Sopenharmony_ci		switch (opt) {
3762306a36Sopenharmony_ci		/* General args */
3862306a36Sopenharmony_ci		case 'F':
3962306a36Sopenharmony_ci			pinned_file = optarg;
4062306a36Sopenharmony_ci			break;
4162306a36Sopenharmony_ci		case 'U':
4262306a36Sopenharmony_ci			pinned_file = optarg;
4362306a36Sopenharmony_ci			create_array = 0;
4462306a36Sopenharmony_ci			break;
4562306a36Sopenharmony_ci		case 'v':
4662306a36Sopenharmony_ci			cg2 = optarg;
4762306a36Sopenharmony_ci			break;
4862306a36Sopenharmony_ci		default:
4962306a36Sopenharmony_ci			usage();
5062306a36Sopenharmony_ci			goto out;
5162306a36Sopenharmony_ci		}
5262306a36Sopenharmony_ci	}
5362306a36Sopenharmony_ci
5462306a36Sopenharmony_ci	if (!cg2 || !pinned_file) {
5562306a36Sopenharmony_ci		usage();
5662306a36Sopenharmony_ci		goto out;
5762306a36Sopenharmony_ci	}
5862306a36Sopenharmony_ci
5962306a36Sopenharmony_ci	cg2_fd = open(cg2, O_RDONLY);
6062306a36Sopenharmony_ci	if (cg2_fd < 0) {
6162306a36Sopenharmony_ci		fprintf(stderr, "open(%s,...): %s(%d)\n",
6262306a36Sopenharmony_ci			cg2, strerror(errno), errno);
6362306a36Sopenharmony_ci		goto out;
6462306a36Sopenharmony_ci	}
6562306a36Sopenharmony_ci
6662306a36Sopenharmony_ci	if (create_array) {
6762306a36Sopenharmony_ci		array_fd = bpf_map_create(BPF_MAP_TYPE_CGROUP_ARRAY, NULL,
6862306a36Sopenharmony_ci					  sizeof(uint32_t), sizeof(uint32_t),
6962306a36Sopenharmony_ci					  1, NULL);
7062306a36Sopenharmony_ci		if (array_fd < 0) {
7162306a36Sopenharmony_ci			fprintf(stderr,
7262306a36Sopenharmony_ci				"bpf_create_map(BPF_MAP_TYPE_CGROUP_ARRAY,...): %s(%d)\n",
7362306a36Sopenharmony_ci				strerror(errno), errno);
7462306a36Sopenharmony_ci			goto out;
7562306a36Sopenharmony_ci		}
7662306a36Sopenharmony_ci	} else {
7762306a36Sopenharmony_ci		array_fd = bpf_obj_get(pinned_file);
7862306a36Sopenharmony_ci		if (array_fd < 0) {
7962306a36Sopenharmony_ci			fprintf(stderr, "bpf_obj_get(%s): %s(%d)\n",
8062306a36Sopenharmony_ci				pinned_file, strerror(errno), errno);
8162306a36Sopenharmony_ci			goto out;
8262306a36Sopenharmony_ci		}
8362306a36Sopenharmony_ci	}
8462306a36Sopenharmony_ci
8562306a36Sopenharmony_ci	ret = bpf_map_update_elem(array_fd, &array_key, &cg2_fd, 0);
8662306a36Sopenharmony_ci	if (ret) {
8762306a36Sopenharmony_ci		perror("bpf_map_update_elem");
8862306a36Sopenharmony_ci		goto out;
8962306a36Sopenharmony_ci	}
9062306a36Sopenharmony_ci
9162306a36Sopenharmony_ci	if (create_array) {
9262306a36Sopenharmony_ci		ret = bpf_obj_pin(array_fd, pinned_file);
9362306a36Sopenharmony_ci		if (ret) {
9462306a36Sopenharmony_ci			fprintf(stderr, "bpf_obj_pin(..., %s): %s(%d)\n",
9562306a36Sopenharmony_ci				pinned_file, strerror(errno), errno);
9662306a36Sopenharmony_ci			goto out;
9762306a36Sopenharmony_ci		}
9862306a36Sopenharmony_ci	}
9962306a36Sopenharmony_ci
10062306a36Sopenharmony_ciout:
10162306a36Sopenharmony_ci	if (array_fd != -1)
10262306a36Sopenharmony_ci		close(array_fd);
10362306a36Sopenharmony_ci	if (cg2_fd != -1)
10462306a36Sopenharmony_ci		close(cg2_fd);
10562306a36Sopenharmony_ci	return ret;
10662306a36Sopenharmony_ci}
107