samples/bpf/cgroup_helpers.c

   1 // SPDX-License-Identifier: GPL-2.0
   2 #define _GNU_SOURCE
   3 #include <sched.h>
   4 #include <sys/mount.h>
   5 #include <sys/stat.h>
   6 #include <sys/types.h>
   7 #include <linux/limits.h>
   8 #include <stdio.h>
   9 #include <linux/sched.h>
  10 #include <fcntl.h>
  11 #include <unistd.h>
  12 #include <ftw.h>
  13
  14
  15 #include "cgroup_helpers.h"
  16
  17 /*
  18  * To avoid relying on the system setup, when setup_cgroup_env is called
  19  * we create a new mount namespace, and cgroup namespace. The cgroup2
  20  * root is mounted at CGROUP_MOUNT_PATH
  21  *
  22  * Unfortunately, most people don't have cgroupv2 enabled at this point in time.
  23  * It's easier to create our own mount namespace and manage it ourselves.
  24  *
  25  * We assume /mnt exists.
  26  */
  27
  28 #define WALK_FD_LIMIT                   16
  29 #define CGROUP_MOUNT_PATH               "/mnt"
  30 #define CGROUP_WORK_DIR                 "/cgroup-test-work-dir"
  31 #define format_cgroup_path(buf, path) \
  32         snprintf(buf, sizeof(buf), "%s%s%s", CGROUP_MOUNT_PATH, \
  33                  CGROUP_WORK_DIR, path)
  34
  35 /**
  36  * setup_cgroup_environment() - Setup the cgroup environment
  37  *
  38  * After calling this function, cleanup_cgroup_environment should be called
  39  * once testing is complete.
  40  *
  41  * This function will print an error to stderr and return 1 if it is unable
  42  * to setup the cgroup environment. If setup is successful, 0 is returned.
  43  */
  44 int setup_cgroup_environment(void)
  45 {
  46         char cgroup_workdir[PATH_MAX + 1];
  47
  48         format_cgroup_path(cgroup_workdir, "");
  49
  50         if (unshare(CLONE_NEWNS)) {
  51                 log_err("unshare");
  52                 return 1;
  53         }
  54
  55         if (mount("none", "/", NULL, MS_REC | MS_PRIVATE, NULL)) {
  56                 log_err("mount fakeroot");
  57                 return 1;
  58         }
  59
  60         if (mount("none", CGROUP_MOUNT_PATH, "cgroup2", 0, NULL)) {
  61                 log_err("mount cgroup2");
  62                 return 1;
  63         }
  64
  65         /* Cleanup existing failed runs, now that the environment is setup */
  66         cleanup_cgroup_environment();
  67
  68         if (mkdir(cgroup_workdir, 0777) && errno != EEXIST) {
  69                 log_err("mkdir cgroup work dir");
  70                 return 1;
  71         }
  72
  73         return 0;
  74 }
  75
  76 static int nftwfunc(const char *filename, const struct stat *statptr,
  77                     int fileflags, struct FTW *pfwt)
  78 {
  79         if ((fileflags & FTW_D) && rmdir(filename))
  80                 log_err("Removing cgroup: %s", filename);
  81         return 0;
  82 }
  83
  84
  85 static int join_cgroup_from_top(char *cgroup_path)
  86 {
  87         char cgroup_procs_path[PATH_MAX + 1];
  88         pid_t pid = getpid();
  89         int fd, rc = 0;
  90
  91         snprintf(cgroup_procs_path, sizeof(cgroup_procs_path),
  92                  "%s/cgroup.procs", cgroup_path);
  93
  94         fd = open(cgroup_procs_path, O_WRONLY);
  95         if (fd < 0) {
  96                 log_err("Opening Cgroup Procs: %s", cgroup_procs_path);
  97                 return 1;
  98         }
  99
 100         if (dprintf(fd, "%d\n", pid) < 0) {
 101                 log_err("Joining Cgroup");
 102                 rc = 1;
 103         }
 104
 105         close(fd);
 106         return rc;
 107 }
 108
 109 /**
 110  * join_cgroup() - Join a cgroup
 111  * @path: The cgroup path, relative to the workdir, to join
 112  *
 113  * This function expects a cgroup to already be created, relative to the cgroup
 114  * work dir, and it joins it. For example, passing "/my-cgroup" as the path
 115  * would actually put the calling process into the cgroup
 116  * "/cgroup-test-work-dir/my-cgroup"
 117  *
 118  * On success, it returns 0, otherwise on failure it returns 1.
 119  */
 120 int join_cgroup(char *path)
 121 {
 122         char cgroup_path[PATH_MAX + 1];
 123
 124         format_cgroup_path(cgroup_path, path);
 125         return join_cgroup_from_top(cgroup_path);
 126 }
 127
 128 /**
 129  * cleanup_cgroup_environment() - Cleanup Cgroup Testing Environment
 130  *
 131  * This is an idempotent function to delete all temporary cgroups that
 132  * have been created during the test, including the cgroup testing work
 133  * directory.
 134  *
 135  * At call time, it moves the calling process to the root cgroup, and then
 136  * runs the deletion process. It is idempotent, and should not fail, unless
 137  * a process is lingering.
 138  *
 139  * On failure, it will print an error to stderr, and try to continue.
 140  */
 141 void cleanup_cgroup_environment(void)
 142 {
 143         char cgroup_workdir[PATH_MAX + 1];
 144
 145         format_cgroup_path(cgroup_workdir, "");
 146         join_cgroup_from_top(CGROUP_MOUNT_PATH);
 147         nftw(cgroup_workdir, nftwfunc, WALK_FD_LIMIT, FTW_DEPTH | FTW_MOUNT);
 148 }
 149
 150 /**
 151  * create_and_get_cgroup() - Create a cgroup, relative to workdir, and get the FD
 152  * @path: The cgroup path, relative to the workdir, to join
 153  *
 154  * This function creates a cgroup under the top level workdir and returns the
 155  * file descriptor. It is idempotent.
 156  *
 157  * On success, it returns the file descriptor. On failure it returns 0.
 158  * If there is a failure, it prints the error to stderr.
 159  */
 160 int create_and_get_cgroup(char *path)
 161 {
 162         char cgroup_path[PATH_MAX + 1];
 163         int fd;
 164
 165         format_cgroup_path(cgroup_path, path);
 166         if (mkdir(cgroup_path, 0777) && errno != EEXIST) {
 167                 log_err("mkdiring cgroup");
 168                 return 0;
 169         }
 170
 171         fd = open(cgroup_path, O_RDONLY);
 172         if (fd < 0) {
 173                 log_err("Opening Cgroup");
 174                 return 0;
 175         }
 176
 177         return fd;
 178 }