samples/bpf/cgroup_helpers.c

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