LXR linux/samples/bpf/cgroup

   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 */
  43int 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
  75static 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
  84static 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 */
 119int 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 */
 140void 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 */
 159int 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}
 178