1 files changed, 164 insertions, 72 deletions

diff --git a/src/libstore/lock.cc b/src/libstore/lock.cc
index fa718f55d..2858137d6 100644
--- a/src/libstore/lock.cc
+++ b/src/libstore/lock.cc
@@ -2,105 +2,197 @@
 #include "globals.hh"
 #include "pathlocks.hh"
 
-#include <grp.h>
 #include <pwd.h>
-
-#include <fcntl.h>
-#include <unistd.h>
+#include <grp.h>
 
 namespace nix {
 
-UserLock::UserLock()
+struct SimpleUserLock : UserLock
 {
-    assert(settings.buildUsersGroup != "");
-    createDirs(settings.nixStateDir + "/userpool");
-}
+    AutoCloseFD fdUserLock;
+    uid_t uid;
+    gid_t gid;
+    std::vector<gid_t> supplementaryGIDs;
+
+    uid_t getUID() override { assert(uid); return uid; }
+    uid_t getUIDCount() override { return 1; }
+    gid_t getGID() override { assert(gid); return gid; }
+
+    std::vector<gid_t> getSupplementaryGIDs() override { return supplementaryGIDs; }
+
+    static std::unique_ptr<UserLock> acquire()
+    {
+        assert(settings.buildUsersGroup != "");
+        createDirs(settings.nixStateDir + "/userpool");
+
+        /* Get the members of the build-users-group. */
+        struct group * gr = getgrnam(settings.buildUsersGroup.get().c_str());
+        if (!gr)
+            throw Error("the group '%s' specified in 'build-users-group' does not exist", settings.buildUsersGroup);
+
+        /* Copy the result of getgrnam. */
+        Strings users;
+        for (char * * p = gr->gr_mem; *p; ++p) {
+            debug("found build user '%s'", *p);
+            users.push_back(*p);
+        }
 
-bool UserLock::findFreeUser() {
-    if (enabled()) return true;
-
-    /* Get the members of the build-users-group. */
-    struct group * gr = getgrnam(settings.buildUsersGroup.get().c_str());
-    if (!gr)
-        throw Error("the group '%1%' specified in 'build-users-group' does not exist",
-            settings.buildUsersGroup);
-    gid = gr->gr_gid;
-
-    /* Copy the result of getgrnam. */
-    Strings users;
-    for (char * * p = gr->gr_mem; *p; ++p) {
-        debug("found build user '%1%'", *p);
-        users.push_back(*p);
+        if (users.empty())
+            throw Error("the build users group '%s' has no members", settings.buildUsersGroup);
+
+        /* Find a user account that isn't currently in use for another
+           build. */
+        for (auto & i : users) {
+            debug("trying user '%s'", i);
+
+            struct passwd * pw = getpwnam(i.c_str());
+            if (!pw)
+                throw Error("the user '%s' in the group '%s' does not exist", i, settings.buildUsersGroup);
+
+            auto fnUserLock = fmt("%s/userpool/%s", settings.nixStateDir,pw->pw_uid);
+
+            AutoCloseFD fd = open(fnUserLock.c_str(), O_RDWR | O_CREAT | O_CLOEXEC, 0600);
+            if (!fd)
+                throw SysError("opening user lock '%s'", fnUserLock);
+
+            if (lockFile(fd.get(), ltWrite, false)) {
+                auto lock = std::make_unique<SimpleUserLock>();
+
+                lock->fdUserLock = std::move(fd);
+                lock->uid = pw->pw_uid;
+                lock->gid = gr->gr_gid;
+
+                /* Sanity check... */
+                if (lock->uid == getuid() || lock->uid == geteuid())
+                    throw Error("the Nix user should not be a member of '%s'", settings.buildUsersGroup);
+
+                #if __linux__
+                /* Get the list of supplementary groups of this build
+                   user.  This is usually either empty or contains a
+                   group such as "kvm".  */
+                int ngroups = 32; // arbitrary initial guess
+                std::vector<gid_t> gids;
+                gids.resize(ngroups);
+
+                int err = getgrouplist(
+                    pw->pw_name, pw->pw_gid,
+                    gids.data(),
+                    &ngroups);
+
+                /* Our initial size of 32 wasn't sufficient, the
+                   correct size has been stored in ngroups, so we try
+                   again. */
+                if (err == -1) {
+                    gids.resize(ngroups);
+                    err = getgrouplist(
+                        pw->pw_name, pw->pw_gid,
+                        gids.data(),
+                        &ngroups);
+                }
+
+                // If it failed once more, then something must be broken.
+                if (err == -1)
+                    throw Error("failed to get list of supplementary groups for '%s'", pw->pw_name);
+
+                // Finally, trim back the GID list to its real size.
+                for (auto i = 0; i < ngroups; i++)
+                    if (gids[i] != lock->gid)
+                        lock->supplementaryGIDs.push_back(gids[i]);
+                #endif
+
+                return lock;
+            }
+        }
+
+        return nullptr;
     }
+};
 
-    if (users.empty())
-        throw Error("the build users group '%1%' has no members",
-            settings.buildUsersGroup);
+struct AutoUserLock : UserLock
+{
+    AutoCloseFD fdUserLock;
+    uid_t firstUid = 0;
+    gid_t firstGid = 0;
+    uid_t nrIds = 1;
 
-    /* Find a user account that isn't currently in use for another
-       build. */
-    for (auto & i : users) {
-        debug("trying user '%1%'", i);
+    uid_t getUID() override { assert(firstUid); return firstUid; }
 
-        struct passwd * pw = getpwnam(i.c_str());
-        if (!pw)
-            throw Error("the user '%1%' in the group '%2%' does not exist",
-                i, settings.buildUsersGroup);
+    gid_t getUIDCount() override { return nrIds; }
 
+    gid_t getGID() override { assert(firstGid); return firstGid; }
 
-        fnUserLock = (format("%1%/userpool/%2%") % settings.nixStateDir % pw->pw_uid).str();
+    std::vector<gid_t> getSupplementaryGIDs() override { return {}; }
 
-        AutoCloseFD fd = open(fnUserLock.c_str(), O_RDWR | O_CREAT | O_CLOEXEC, 0600);
-        if (!fd)
-            throw SysError("opening user lock '%1%'", fnUserLock);
+    static std::unique_ptr<UserLock> acquire(uid_t nrIds, bool useChroot)
+    {
+        settings.requireExperimentalFeature(Xp::AutoAllocateUids);
+        assert(settings.startId > 0);
+        assert(settings.uidCount % maxIdsPerBuild == 0);
+        assert((uint64_t) settings.startId + (uint64_t) settings.uidCount <= std::numeric_limits<uid_t>::max());
+        assert(nrIds <= maxIdsPerBuild);
 
-        if (lockFile(fd.get(), ltWrite, false)) {
-            fdUserLock = std::move(fd);
-            user = i;
-            uid = pw->pw_uid;
+        createDirs(settings.nixStateDir + "/userpool2");
 
-            /* Sanity check... */
-            if (uid == getuid() || uid == geteuid())
-                throw Error("the Nix user should not be a member of '%1%'",
-                    settings.buildUsersGroup);
+        size_t nrSlots = settings.uidCount / maxIdsPerBuild;
 
-#if __linux__
-            /* Get the list of supplementary groups of this build user.  This
-               is usually either empty or contains a group such as "kvm".  */
-            int ngroups = 32; // arbitrary initial guess
-            supplementaryGIDs.resize(ngroups);
+        for (size_t i = 0; i < nrSlots; i++) {
+            debug("trying user slot '%d'", i);
 
-            int err = getgrouplist(pw->pw_name, pw->pw_gid, supplementaryGIDs.data(),
-                                 &ngroups);
+            createDirs(settings.nixStateDir + "/userpool2");
 
-            // Our initial size of 32 wasn't sufficient, the correct size has
-            // been stored in ngroups, so we try again.
-            if (err == -1) {
-                supplementaryGIDs.resize(ngroups);
-                err = getgrouplist(pw->pw_name, pw->pw_gid, supplementaryGIDs.data(),
-                                   &ngroups);
-            }
+            auto fnUserLock = fmt("%s/userpool2/slot-%d", settings.nixStateDir, i);
+
+            AutoCloseFD fd = open(fnUserLock.c_str(), O_RDWR | O_CREAT | O_CLOEXEC, 0600);
+            if (!fd)
+                throw SysError("opening user lock '%s'", fnUserLock);
 
-            // If it failed once more, then something must be broken.
-            if (err == -1)
-                throw Error("failed to get list of supplementary groups for '%1%'",
-                            pw->pw_name);
+            if (lockFile(fd.get(), ltWrite, false)) {
 
-            // Finally, trim back the GID list to its real size
-            supplementaryGIDs.resize(ngroups);
-#endif
+                auto firstUid = settings.startId + i * maxIdsPerBuild;
 
-            isEnabled = true;
-            return true;
+                auto pw = getpwuid(firstUid);
+                if (pw)
+                    throw Error("auto-allocated UID %d clashes with existing user account '%s'", firstUid, pw->pw_name);
+
+                auto lock = std::make_unique<AutoUserLock>();
+                lock->fdUserLock = std::move(fd);
+                lock->firstUid = firstUid;
+                if (useChroot)
+                    lock->firstGid = firstUid;
+                else {
+                    struct group * gr = getgrnam(settings.buildUsersGroup.get().c_str());
+                    if (!gr)
+                        throw Error("the group '%s' specified in 'build-users-group' does not exist", settings.buildUsersGroup);
+                    lock->firstGid = gr->gr_gid;
+                }
+                lock->nrIds = nrIds;
+                return lock;
+            }
         }
+
+        return nullptr;
     }
+};
 
-    return false;
+std::unique_ptr<UserLock> acquireUserLock(uid_t nrIds, bool useChroot)
+{
+    if (settings.autoAllocateUids)
+        return AutoUserLock::acquire(nrIds, useChroot);
+    else
+        return SimpleUserLock::acquire();
 }
 
-void UserLock::kill()
+bool useBuildUsers()
 {
-    killUser(uid);
+    #if __linux__
+    static bool b = (settings.buildUsersGroup != "" || settings.startId.get() != 0) && getuid() == 0;
+    return b;
+    #elif __APPLE__
+    static bool b = settings.buildUsersGroup != "" && getuid() == 0;
+    return b;
+    #else
+    return false;
+    #endif
 }
 
 }