#include <stdbool.h>
 #include <stdarg.h>
 
+/*
+ * NOTES about this test:
+ * - requries libcap-dev to be installed on test system
+ * - requires securityfs to me mounted at /sys/kernel/security, e.g.:
+ * mount -n -t securityfs -o nodev,noexec,nosuid securityfs /sys/kernel/security
+ * - needs CONFIG_SECURITYFS and CONFIG_SAFESETID to be enabled
+ */
+
 #ifndef CLONE_NEWUSER
 # define CLONE_NEWUSER 0x10000000
 #endif
 
-#define ROOT_USER 0
-#define RESTRICTED_PARENT 1
-#define ALLOWED_CHILD1 2
-#define ALLOWED_CHILD2 3
-#define NO_POLICY_USER 4
+#define ROOT_UGID 0
+#define RESTRICTED_PARENT_UGID 1
+#define ALLOWED_CHILD1_UGID 2
+#define ALLOWED_CHILD2_UGID 3
+#define NO_POLICY_UGID 4
+
+#define UGID_POLICY_STRING "1:2\n1:3\n2:2\n3:3\n"
 
-char* add_whitelist_policy_file = "/sys/kernel/security/safesetid/add_whitelist_policy";
+char* add_uid_whitelist_policy_file = "/sys/kernel/security/safesetid/uid_allowlist_policy";
+char* add_gid_whitelist_policy_file = "/sys/kernel/security/safesetid/gid_allowlist_policy";
 
 static void die(char *fmt, ...)
 {
                        die("couldn't open file\n");
                if (fseek(fd, 0, SEEK_END))
                        die("couldn't fseek\n");
-               snprintf(name_str, 10, "%d", uid);
+               snprintf(name_str, 10, "user %d", uid);
                p.pw_name=name_str;
                p.pw_uid=uid;
                p.pw_gecos="Test account";
 
 static void ensure_securityfs_mounted(void)
 {
-       int fd = open(add_whitelist_policy_file, O_WRONLY);
+       int fd = open(add_uid_whitelist_policy_file, O_WRONLY);
        if (fd < 0) {
                if (errno == ENOENT) {
                        // Need to mount securityfs
        } else {
                if (close(fd) != 0) {
                        die("close of %s failed: %s\n",
-                               add_whitelist_policy_file, strerror(errno));
+                               add_uid_whitelist_policy_file, strerror(errno));
                }
        }
 }
 
-static void write_policies(void)
+static void write_uid_policies()
 {
-       static char *policy_str =
-               "1:2\n"
-               "1:3\n"
-               "2:2\n"
-               "3:3\n";
+       static char *policy_str = UGID_POLICY_STRING;
        ssize_t written;
        int fd;
 
-       fd = open(add_whitelist_policy_file, O_WRONLY);
+       fd = open(add_uid_whitelist_policy_file, O_WRONLY);
        if (fd < 0)
-               die("can't open add_whitelist_policy file\n");
+               die("can't open add_uid_whitelist_policy file\n");
        written = write(fd, policy_str, strlen(policy_str));
        if (written != strlen(policy_str)) {
                if (written >= 0) {
-                       die("short write to %s\n", add_whitelist_policy_file);
+                       die("short write to %s\n", add_uid_whitelist_policy_file);
                } else {
                        die("write to %s failed: %s\n",
-                               add_whitelist_policy_file, strerror(errno));
+                               add_uid_whitelist_policy_file, strerror(errno));
                }
        }
        if (close(fd) != 0) {
                die("close of %s failed: %s\n",
-                       add_whitelist_policy_file, strerror(errno));
+                       add_uid_whitelist_policy_file, strerror(errno));
        }
 }
 
 
 static void ensure_users_exist(void)
 {
-       ensure_user_exists(ROOT_USER);
-       ensure_user_exists(RESTRICTED_PARENT);
-       ensure_user_exists(ALLOWED_CHILD1);
-       ensure_user_exists(ALLOWED_CHILD2);
-       ensure_user_exists(NO_POLICY_USER);
+       ensure_user_exists(ROOT_UGID);
+       ensure_user_exists(RESTRICTED_PARENT_UGID);
+       ensure_user_exists(ALLOWED_CHILD1_UGID);
+       ensure_user_exists(ALLOWED_CHILD2_UGID);
+       ensure_user_exists(NO_POLICY_UGID);
 }
 
 static void drop_caps(bool setid_retained)
 {
        ensure_users_exist();
        ensure_securityfs_mounted();
-       write_policies();
+       write_uid_policies();
 
        if (prctl(PR_SET_KEEPCAPS, 1L))
                die("Error with set keepcaps\n");
 
-       // First test to make sure we can write userns mappings from a user
-       // that doesn't have any restrictions (as long as it has CAP_SETUID);
-       if (setuid(NO_POLICY_USER) < 0)
-               die("Error with set uid(%d)\n", NO_POLICY_USER);
-       if (setgid(NO_POLICY_USER) < 0)
-               die("Error with set gid(%d)\n", NO_POLICY_USER);
-
+       // First test to make sure we can write userns mappings from a non-root
+       // user that doesn't have any restrictions (as long as it has
+       // CAP_SETUID);
+       if (setgid(NO_POLICY_UGID) < 0)
+               die("Error with set gid(%d)\n", NO_POLICY_UGID);
+       if (setuid(NO_POLICY_UGID) < 0)
+               die("Error with set uid(%d)\n", NO_POLICY_UGID);
        // Take away all but setid caps
        drop_caps(true);
-
        // Need PR_SET_DUMPABLE flag set so we can write /proc/[pid]/uid_map
        // from non-root parent process.
        if (prctl(PR_SET_DUMPABLE, 1, 0, 0, 0))
                die("Error with set dumpable\n");
-
        if (!test_userns(true)) {
                die("test_userns failed when it should work\n");
        }
 
-       if (setuid(RESTRICTED_PARENT) < 0)
-               die("Error with set uid(%d)\n", RESTRICTED_PARENT);
-       if (setgid(RESTRICTED_PARENT) < 0)
-               die("Error with set gid(%d)\n", RESTRICTED_PARENT);
+       // Now switch to a user/group with restrictions
+       if (setgid(RESTRICTED_PARENT_UGID) < 0)
+               die("Error with set gid(%d)\n", RESTRICTED_PARENT_UGID);
+       if (setuid(RESTRICTED_PARENT_UGID) < 0)
+               die("Error with set uid(%d)\n", RESTRICTED_PARENT_UGID);
 
-       test_setuid(ROOT_USER, false);
-       test_setuid(ALLOWED_CHILD1, true);
-       test_setuid(ALLOWED_CHILD2, true);
-       test_setuid(NO_POLICY_USER, false);
+       test_setuid(ROOT_UGID, false);
+       test_setuid(ALLOWED_CHILD1_UGID, true);
+       test_setuid(ALLOWED_CHILD2_UGID, true);
+       test_setuid(NO_POLICY_UGID, false);
 
        if (!test_userns(false)) {
                die("test_userns worked when it should fail\n");
 
        // NOTE: this test doesn't clean up users that were created in
        // /etc/passwd or flush policies that were added to the LSM.
+       printf("test successful!\n");
        return EXIT_SUCCESS;
 }