From ad1b36291fe87ca9d647980e3afdc71344d66519 Mon Sep 17 00:00:00 2001
From: Tavian Barnes <tavianator@tavianator.com>
Date: Mon, 6 Nov 2023 10:06:41 -0500
Subject: ctx: Store the original and current RLIMIT_NOFILE

---
 src/ctx.c  |  6 ++----
 src/ctx.h  |  8 ++++----
 src/eval.c | 33 +++++++++++++++++++--------------
 src/eval.h |  2 +-
 src/exec.c | 26 ++++++++++++--------------
 5 files changed, 38 insertions(+), 37 deletions(-)

diff --git a/src/ctx.c b/src/ctx.c
index 561ecae..692f5c5 100644
--- a/src/ctx.c
+++ b/src/ctx.c
@@ -30,12 +30,10 @@ struct bfs_ctx *bfs_ctx_new(void) {
 
 	trie_init(&ctx->files);
 
-	struct rlimit rl;
-	if (getrlimit(RLIMIT_NOFILE, &rl) != 0) {
+	if (getrlimit(RLIMIT_NOFILE, &ctx->orig_nofile) != 0) {
 		goto fail;
 	}
-	ctx->nofile_soft = rl.rlim_cur;
-	ctx->nofile_hard = rl.rlim_max;
+	ctx->cur_nofile = ctx->orig_nofile;
 
 	ctx->users = bfs_users_new();
 	if (!ctx->users) {
diff --git a/src/ctx.h b/src/ctx.h
index 96406bd..1c7df63 100644
--- a/src/ctx.h
+++ b/src/ctx.h
@@ -87,10 +87,10 @@ struct bfs_ctx {
 	/** The number of files owned by the context. */
 	int nfiles;
 
-	/** The initial RLIMIT_NOFILE soft limit. */
-	rlim_t nofile_soft;
-	/** The initial RLIMIT_NOFILE hard limit. */
-	rlim_t nofile_hard;
+	/** The initial RLIMIT_NOFILE limits. */
+	struct rlimit orig_nofile;
+	/** The current RLIMIT_NOFILE limits. */
+	struct rlimit cur_nofile;
 
 	/** The current time. */
 	struct timespec now;
diff --git a/src/eval.c b/src/eval.c
index 5ba3de8..b511eba 100644
--- a/src/eval.c
+++ b/src/eval.c
@@ -1428,25 +1428,30 @@ done:
 }
 
 /** Raise RLIMIT_NOFILE if possible, and return the new limit. */
-static int raise_fdlimit(const struct bfs_ctx *ctx) {
+static int raise_fdlimit(struct bfs_ctx *ctx) {
+	rlim_t cur = ctx->orig_nofile.rlim_cur;
+	rlim_t max = ctx->orig_nofile.rlim_max;
+
 	rlim_t target = 64 << 10;
-	if (rlim_cmp(target, ctx->nofile_hard) > 0) {
-		target = ctx->nofile_hard;
+	if (rlim_cmp(target, max) > 0) {
+		target = max;
 	}
 
-	int ret = target;
+	if (rlim_cmp(target, cur) <= 0) {
+		return target;
+	}
 
-	if (rlim_cmp(target, ctx->nofile_soft) > 0) {
-		const struct rlimit rl = {
-			.rlim_cur = target,
-			.rlim_max = ctx->nofile_hard,
-		};
-		if (setrlimit(RLIMIT_NOFILE, &rl) != 0) {
-			ret = ctx->nofile_soft;
-		}
+	const struct rlimit rl = {
+		.rlim_cur = target,
+		.rlim_max = max,
+	};
+
+	if (setrlimit(RLIMIT_NOFILE, &rl) != 0) {
+		return cur;
 	}
 
-	return ret;
+	ctx->cur_nofile = rl;
+	return target;
 }
 
 /** Preallocate the fd table in the kernel. */
@@ -1581,7 +1586,7 @@ static bool eval_must_buffer(const struct bfs_expr *expr) {
 	return false;
 }
 
-int bfs_eval(const struct bfs_ctx *ctx) {
+int bfs_eval(struct bfs_ctx *ctx) {
 	if (!ctx->expr) {
 		return EXIT_SUCCESS;
 	}
diff --git a/src/eval.h b/src/eval.h
index bdb9440..98bbc08 100644
--- a/src/eval.h
+++ b/src/eval.h
@@ -39,7 +39,7 @@ typedef bool bfs_eval_fn(const struct bfs_expr *expr, struct bfs_eval *state);
  * @return
  *         EXIT_SUCCESS on success, otherwise on failure.
  */
-int bfs_eval(const struct bfs_ctx *ctx);
+int bfs_eval(struct bfs_ctx *ctx);
 
 // Predicate evaluation functions
 
diff --git a/src/exec.c b/src/exec.c
index 97cfafa..a3297e8 100644
--- a/src/exec.c
+++ b/src/exec.c
@@ -327,8 +327,10 @@ static void bfs_exec_closewd(struct bfs_exec *execbuf, const struct BFTW *ftwbuf
 
 /** Actually spawn the process. */
 static int bfs_exec_spawn(const struct bfs_exec *execbuf) {
+	const struct bfs_ctx *ctx = execbuf->ctx;
+
 	// Flush the context state for consistency with the external process
-	bfs_ctx_flush(execbuf->ctx);
+	bfs_ctx_flush(ctx);
 
 	if (execbuf->flags & BFS_EXEC_CONFIRM) {
 		for (size_t i = 0; i < execbuf->argc; ++i) {
@@ -356,34 +358,30 @@ static int bfs_exec_spawn(const struct bfs_exec *execbuf) {
 	pid_t pid = -1;
 	int error;
 
-	struct bfs_spawn ctx;
-	if (bfs_spawn_init(&ctx) != 0) {
+	struct bfs_spawn spawn;
+	if (bfs_spawn_init(&spawn) != 0) {
 		return -1;
 	}
 
-	if (bfs_spawn_setflags(&ctx, BFS_SPAWN_USEPATH) != 0) {
+	if (bfs_spawn_setflags(&spawn, BFS_SPAWN_USEPATH) != 0) {
 		goto fail;
 	}
 
 	// Reset RLIMIT_NOFILE, to avoid breaking applications that use select()
-	struct rlimit rl = {
-		.rlim_cur = execbuf->ctx->nofile_soft,
-		.rlim_max = execbuf->ctx->nofile_hard,
-	};
-	if (bfs_spawn_addsetrlimit(&ctx, RLIMIT_NOFILE, &rl) != 0) {
+	if (bfs_spawn_addsetrlimit(&spawn, RLIMIT_NOFILE, &ctx->orig_nofile) != 0) {
 		goto fail;
 	}
 
 	if (execbuf->wd_fd >= 0) {
-		if (bfs_spawn_addfchdir(&ctx, execbuf->wd_fd) != 0) {
+		if (bfs_spawn_addfchdir(&spawn, execbuf->wd_fd) != 0) {
 			goto fail;
 		}
 	}
 
-	pid = bfs_spawn(execbuf->argv[0], &ctx, execbuf->argv, NULL);
+	pid = bfs_spawn(execbuf->argv[0], &spawn, execbuf->argv, NULL);
 fail:
 	error = errno;
-	bfs_spawn_destroy(&ctx);
+	bfs_spawn_destroy(&spawn);
 	if (pid < 0) {
 		errno = error;
 		return -1;
@@ -409,9 +407,9 @@ fail:
 		if (!str) {
 			str = "unknown";
 		}
-		bfs_warning(execbuf->ctx, "Command '${ex}%s${rs}' terminated by signal %d (%s)\n", execbuf->argv[0], sig, str);
+		bfs_warning(ctx, "Command '${ex}%s${rs}' terminated by signal %d (%s)\n", execbuf->argv[0], sig, str);
 	} else {
-		bfs_warning(execbuf->ctx, "Command '${ex}%s${rs}' terminated abnormally\n", execbuf->argv[0]);
+		bfs_warning(ctx, "Command '${ex}%s${rs}' terminated abnormally\n", execbuf->argv[0]);
 	}
 
 	errno = 0;
-- 
cgit v1.2.3