Implement RLIMIT_NPROC

This is simplified version without supporting hierarchical
user-namespace counters.

Signed-off-by: Andrei Vagin <avagin@gmail.com>

pkg/sentry/fsimpl/testutil/kernel.go

@@ -130,6 +130,7 @@ func CreateTask(ctx context.Context, name string, tc *kernel.ThreadGroup, mntns
 	m := mm.NewMemoryManager(k, k, k.SleepForAddressSpaceActivation)
 	m.SetExecutable(ctx, fsbridge.NewVFSFile(exe))
 
+	creds := auth.CredentialsFromContext(ctx)
 	config := &kernel.TaskConfig{
 		Kernel:                  k,
 		ThreadGroup:             tc,
@@ -143,6 +144,7 @@ func CreateTask(ctx context.Context, name string, tc *kernel.ThreadGroup, mntns
 		MountNamespaceVFS2:      mntns,
 		FSContext:               kernel.NewFSContextVFS2(root, cwd, 0022),
 		FDTable:                 k.NewFDTable(),
+		UserCounters:            k.GetUserCounters(creds.RealKUID),
 	}
 	t, err := k.TaskSet().NewTask(ctx, config)
 	if err != nil {

pkg/sentry/kernel/kernel.go

@@ -42,6 +42,7 @@ import (
 	"gvisor.dev/gvisor/pkg/cleanup"
 	"gvisor.dev/gvisor/pkg/context"
 	"gvisor.dev/gvisor/pkg/cpuid"
+	"gvisor.dev/gvisor/pkg/errors/linuxerr"
 	"gvisor.dev/gvisor/pkg/eventchannel"
 	"gvisor.dev/gvisor/pkg/fspath"
 	"gvisor.dev/gvisor/pkg/log"
@@ -94,6 +95,35 @@ var LISAFSEnabled = false
 // easy access everywhere. To be removed once FUSE is completed.
 var FUSEEnabled = false
 
+// userCounters is a set of user counters.
+//
+// +stateify savable
+type userCounters struct {
+	uid auth.KUID
+
+	// +checkatomic
+	rlimitNProc uint64
+}
+
+// incRLimitNProc increments the rlimitNProc counter.
+func (uc *userCounters) incRLimitNProc(ctx context.Context) error {
+	lim := limits.FromContext(ctx).Get(limits.ProcessCount)
+	creds := auth.CredentialsFromContext(ctx)
+	nproc := atomic.AddUint64(&uc.rlimitNProc, 1)
+	if nproc > lim.Cur &&
+		!creds.HasCapability(linux.CAP_SYS_ADMIN) &&
+		!creds.HasCapability(linux.CAP_SYS_RESOURCE) {
+		atomic.AddUint64(&uc.rlimitNProc, ^uint64(0))
+		return linuxerr.EAGAIN
+	}
+	return nil
+}
+
+// decRLimitNProc decrements the rlimitNProc counter.
+func (uc *userCounters) decRLimitNProc() {
+	atomic.AddUint64(&uc.rlimitNProc, ^uint64(0))
+}
+
 // Kernel represents an emulated Linux kernel. It must be initialized by calling
 // Init() or LoadFrom().
 //
@@ -301,6 +331,10 @@ type Kernel struct {
 	// system. It is controller by cgroupfs. Nil if cgroupfs is unavailable on
 	// the system.
 	cgroupRegistry *CgroupRegistry
+
+	// userCountersMa maps auth.KUID into a set of user counters.
+	userCountersMap   map[auth.KUID]*userCounters
+	userCountersMapMu sync.Mutex `state:"nosave"`
 }
 
 // InitKernelArgs holds arguments to Init.
@@ -401,6 +435,7 @@ func (k *Kernel) Init(args InitKernelArgs) error {
 	k.netlinkPorts = port.New()
 	k.ptraceExceptions = make(map[*Task]*Task)
 	k.YAMAPtraceScope = linux.YAMA_SCOPE_RELATIONAL
+	k.userCountersMap = make(map[auth.KUID]*userCounters)
 
 	if VFS2Enabled {
 		ctx := k.SupervisorContext()
@@ -1032,6 +1067,7 @@ func (k *Kernel) CreateProcess(args CreateProcessArgs) (*ThreadGroup, ThreadID,
 		AbstractSocketNamespace: args.AbstractSocketNamespace,
 		MountNamespaceVFS2:      mntnsVFS2,
 		ContainerID:             args.ContainerID,
+		UserCounters:            k.GetUserCounters(args.Credentials.RealKUID),
 	}
 	t, err := k.tasks.NewTask(ctx, config)
 	if err != nil {
@@ -1871,3 +1907,16 @@ func (k *Kernel) ReplaceFSContextRoots(ctx context.Context, oldRoot vfs.VirtualD
 		oldRoot.DecRef(ctx)
 	}
 }
+
+func (k *Kernel) GetUserCounters(uid auth.KUID) *userCounters {
+	k.userCountersMapMu.Lock()
+	defer k.userCountersMapMu.Unlock()
+
+	if uc, ok := k.userCountersMap[uid]; ok {
+		return uc
+	}
+
+	uc := &userCounters{}
+	k.userCountersMap[uid] = uc
+	return uc
+}

pkg/sentry/kernel/task.go

@@ -591,6 +591,12 @@ type Task struct {
 	//
 	// +checklocks:mu
 	cgroups map[Cgroup]struct{}
+
+	// userCounters is a pointer to a set of user counters.
+	//
+	// The userCounters pointer is exclusive to the task goroutine, but the
+	// userCounters instance must be atomically accessed.
+	userCounters *userCounters
 }
 
 func (t *Task) savePtraceTracer() *Task {

pkg/sentry/kernel/task_clone.go

@@ -186,6 +186,11 @@ func (t *Task) Clone(args *linux.CloneArgs) (ThreadID, *SyscallControl, error) {
 		rseqSignature = t.rseqSignature
 	}
 
+	uc := t.userCounters
+	if uc.uid != creds.RealKUID {
+		uc = t.k.GetUserCounters(creds.RealKUID)
+	}
+
 	cfg := &TaskConfig{
 		Kernel:                  t.k,
 		ThreadGroup:             tg,
@@ -204,6 +209,7 @@ func (t *Task) Clone(args *linux.CloneArgs) (ThreadID, *SyscallControl, error) {
 		RSeqAddr:                rseqAddr,
 		RSeqSignature:           rseqSignature,
 		ContainerID:             t.ContainerID(),
+		UserCounters:            uc,
 	}
 	if args.Flags&linux.CLONE_THREAD == 0 {
 		cfg.Parent = t

pkg/sentry/kernel/task_exit.go

@@ -653,6 +653,7 @@ func (t *Task) exitNotifyLocked(fromPtraceDetach bool) {
 				delete(ns.tgids, t.tg)
 			}
 		}
+		t.userCounters.decRLimitNProc()
 		t.tg.exitedCPUStats.Accumulate(t.CPUStats())
 		t.tg.ioUsage.Accumulate(t.ioUsage)
 		t.tg.signalHandlers.mu.Lock()

pkg/sentry/kernel/task_start.go

@@ -93,6 +93,9 @@ type TaskConfig struct {
 
 	// ContainerID is the container the new task belongs to.
 	ContainerID string
+
+	// UserCounters is user resource counters.
+	UserCounters *userCounters
 }
 
 // NewTask creates a new task defined by cfg.
@@ -102,8 +105,8 @@ type TaskConfig struct {
 // If successful, NewTask transfers references held by cfg to the new task.
 // Otherwise, NewTask releases them.
 func (ts *TaskSet) NewTask(ctx context.Context, cfg *TaskConfig) (*Task, error) {
-	t, err := ts.newTask(cfg)
-	if err != nil {
+	var err error
+	cleanup := func() {
 		cfg.TaskImage.release()
 		cfg.FSContext.DecRef(ctx)
 		cfg.FDTable.DecRef(ctx)
@@ -111,6 +114,15 @@ func (ts *TaskSet) NewTask(ctx context.Context, cfg *TaskConfig) (*Task, error)
 		if cfg.MountNamespaceVFS2 != nil {
 			cfg.MountNamespaceVFS2.DecRef(ctx)
 		}
+	}
+	if err := cfg.UserCounters.incRLimitNProc(ctx); err != nil {
+		cleanup()
+		return nil, err
+	}
+	t, err := ts.newTask(cfg)
+	if err != nil {
+		cfg.UserCounters.decRLimitNProc()
+		cleanup()
 		return nil, err
 	}
 	return t, nil
@@ -150,6 +162,7 @@ func (ts *TaskSet) newTask(cfg *TaskConfig) (*Task, error) {
 		futexWaiter:        futex.NewWaiter(),
 		containerID:        cfg.ContainerID,
 		cgroups:            make(map[Cgroup]struct{}),
+		userCounters:       cfg.UserCounters,
 	}
 	t.netns.Store(cfg.NetworkNamespace)
 	t.creds.Store(cfg.Credentials)

test/syscalls/linux/BUILD

@@ -2147,6 +2147,7 @@ cc_binary(
         "//test/util:capability_util",
         "//test/util:test_main",
         "//test/util:test_util",
+        "//test/util:thread_util",
     ],
 )
 

test/syscalls/linux/rlimits.cc

@@ -12,11 +12,19 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <errno.h>
+#include <stdlib.h>
 #include <sys/resource.h>
 #include <sys/time.h>
+#include <sys/types.h>
+#include <sys/wait.h>
+#include <unistd.h>
+
+#include <climits>
 
 #include "test/util/capability_util.h"
 #include "test/util/test_util.h"
+#include "test/util/thread_util.h"
 
 namespace gvisor {
 namespace testing {
@@ -67,6 +75,59 @@ TEST(RlimitTest, SetSoftRlimitAboveHard) {
   EXPECT_THAT(setrlimit(RLIMIT_NOFILE, &rl), SyscallFailsWithErrno(EINVAL));
 }
 
+TEST(RlimitTest, RlimitNProc) {
+  SKIP_IF(!ASSERT_NO_ERRNO_AND_VALUE(HaveCapability(CAP_SETUID)));
+
+  // The native test can be run in a user namespace without a mapping for
+  // kNobody or there can be other processes that are running from the kNobody
+  // user.
+  SKIP_IF(!IsRunningOnGvisor());
+
+  // Run the test in a sub-thread to avoid changing UID of the current thread.
+  ScopedThread([&] {
+    constexpr int kNobody = 65534;
+    EXPECT_THAT(syscall(SYS_setuid, kNobody), SyscallSucceeds());
+
+    struct rlimit rl = {};
+    EXPECT_THAT(getrlimit(RLIMIT_NPROC, &rl), SyscallSucceeds());
+
+    constexpr int kNProc = 10;
+    rl.rlim_cur = kNProc;
+    EXPECT_THAT(setrlimit(RLIMIT_NPROC, &rl), SyscallSucceeds());
+
+    constexpr int kIterations = 2;
+    // Run test actions a few times to check that processes are not leaked.
+    for (int iter = 0; iter < kIterations; iter++) {
+      pid_t pids[kNProc];
+      for (int i = 0; i < kNProc; i++) {
+        pid_t pid = fork();
+        if (pid == 0) {
+          while (1) {
+            sleep(1);
+          }
+          _exit(1);
+        }
+        EXPECT_THAT(pid, SyscallSucceeds());
+        pids[i] = pid;
+      }
+      auto cleanup = Cleanup([pids] {
+        for (int i = 0; i < kNProc; i++) {
+          if (pids[i] < 0) {
+            continue;
+          }
+          EXPECT_THAT(kill(pids[i], SIGKILL), SyscallSucceeds());
+          EXPECT_THAT(waitpid(pids[i], nullptr, 0), SyscallSucceeds());
+        }
+      });
+      pid_t pid = fork();
+      if (pid == 0) {
+        _exit(1);
+      }
+      EXPECT_THAT(pid, SyscallFailsWithErrno(EAGAIN));
+    }
+  }).Join();
+}
+
 }  // namespace
 
 }  // namespace testing