lockf: allow locking file descriptors

This is most useful inside a shell script, allowing one to lock just
portions of a script rather than having to wrap the entire script in a
lock.

PR:		262738
Reviewed by:	0mp, allanjude (both previous versions)
Co-authored-by:	Daniel O'Connor <darius@dons.net.au>
Sponsored by:	Klara, Inc.

(cherry picked from commit 09a7fe0a5523d53ff6c26ddef9a947f293e18c22)

3 files changed, 266 insertions, 16 deletions

diff --git a/usr.bin/lockf/lockf.1 b/usr.bin/lockf/lockf.1
index 8120b2ed7630..d73033101632 100644
--- a/usr.bin/lockf/lockf.1
+++ b/usr.bin/lockf/lockf.1
@@ -22,7 +22,7 @@
 .\" OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 .\" SUCH DAMAGE.
 .\"
-.Dd August 26, 2020
+.Dd November 25, 2023
 .Dt LOCKF 1
 .Os
 .Sh NAME
@@ -35,6 +35,10 @@
 .Ar file
 .Ar command
 .Op Ar arguments
+.Nm
+.Op Fl s
+.Op Fl t Ar seconds
+.Ar fd
 .Sh DESCRIPTION
 The
 .Nm
@@ -64,6 +68,27 @@ the mere existence of the
 .Ar file
 is not considered to constitute a lock.
 .Pp
+.Nm
+may also be used to operate on a file descriptor instead of a file.
+If no
+.Ar command
+is supplied, then
+.Ar fd
+must be a file descriptor.
+The version with a
+.Ar command
+may also be used with a file descriptor by supplying it as a path
+.Pa /dev/fd/N ,
+where N is the desired file descriptor.
+The
+.Fl k
+option is implied when a file descriptor is in use, and the
+.Fl n
+and
+.Fl w
+options are silently ignored.
+This can be used to lock inside a shell script.
+.Pp
 If the
 .Nm
 utility is being used to facilitate concurrency between a number
@@ -186,6 +211,42 @@ $ lockf mylock sleep 1 & lockf -t 5 mylock echo "Success"
 Success
 [1]+  Done                    lockf mylock sleep 1
 .Ed
+Lock a file and run a script, return immediately if the lock is not
+available. Do not delete the file afterward so lock order is
+guaranteed.
+.Pp
+.Dl $ lockf -t 0 -k /tmp/my.lock myscript
+.Pp
+Protect a section of a shell script with a lock, wait up to 5 seconds
+for it to become available.
+Note that the shell script has opened the lock file
+.Fa /tmp/my.lock ,
+and
+.Nm
+is performing the lock call exclusively via the passed in file descriptor (9).
+In this case
+.Fl k
+is implied, and
+.Fl w
+has no effect because the file has already been opened by the shell.
+This example assumes that
+.Ql >
+is implemented in the shell by opening and truncating
+.Pa /tmp/my.lock ,
+rather than by replacing the lock file.
+.Bd -literal -offset indent
+(
+	lockf -s -t 5 9
+	if [ $? -ne 0 ]; then
+		echo "Failed to obtain lock"
+		exit 1
+	fi
+
+	echo Start
+	# Do some stuff
+	echo End
+) 9>/tmp/my.lock
+.Ed
 .Sh SEE ALSO
 .Xr flock 2 ,
 .Xr lockf 3 ,
diff --git a/usr.bin/lockf/lockf.c b/usr.bin/lockf/lockf.c
index 1149128021e6..49a4752a967d 100644
--- a/usr.bin/lockf/lockf.c
+++ b/usr.bin/lockf/lockf.c
@@ -29,16 +29,26 @@
 #include <sys/types.h>
 #include <sys/wait.h>
 
+#include <assert.h>
 #include <err.h>
 #include <errno.h>
 #include <fcntl.h>
+#include <limits.h>
 #include <signal.h>
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
 #include <sysexits.h>
 #include <unistd.h>
 
-static int acquire_lock(const char *name, int flags, int silent);
+#define	FDLOCK_PREFIX	"/dev/fd/"
+
+union lock_subject {
+	long		 subj_fd;
+	const char	*subj_name;
+};
+
+static int acquire_lock(union lock_subject *subj, int flags, int silent);
 static void cleanup(void);
 static void killed(int sig);
 static void timeout(int sig);
@@ -48,9 +58,35 @@ static void wait_for_lock(const char *name);
 static const char *lockname;
 static int lockfd = -1;
 static int keep;
+static int fdlock;
 static volatile sig_atomic_t timed_out;
 
 /*
+ * Check if fdlock is implied by the given `lockname`.  We'll write the fd that
+ * is represented by it out to ofd, and the caller is expected to do any
+ * necessary validation on it.
+ */
+static int
+fdlock_implied(const char *name, long *ofd)
+{
+	char *endp;
+	long fd;
+
+	if (strncmp(name, FDLOCK_PREFIX, sizeof(FDLOCK_PREFIX) - 1) != 0)
+		return (0);
+
+	/* Skip past the prefix. */
+	name += sizeof(FDLOCK_PREFIX) - 1;
+	errno = 0;
+	fd = strtol(name, &endp, 10);
+	if (errno != 0 || *endp != '\0')
+		return (0);
+
+	*ofd = fd;
+	return (1);
+}
+
+/*
  * Execute an arbitrary command while holding a file lock.
  */
 int
@@ -58,6 +94,7 @@ main(int argc, char **argv)
 {
 	int ch, flags, silent, status, waitsec;
 	pid_t child;
+	union lock_subject subj;
 
 	silent = keep = 0;
 	flags = O_CREAT | O_RDONLY;
@@ -89,11 +126,54 @@ main(int argc, char **argv)
 			usage();
 		}
 	}
-	if (argc - optind < 2)
-		usage();
-	lockname = argv[optind++];
+
 	argc -= optind;
 	argv += optind;
+
+	if (argc == 0)
+		usage();
+
+	lockname = argv[0];
+
+	argc--;
+	argv++;
+
+	/*
+	 * If there aren't any arguments left, then we must be in fdlock mode.
+	 */
+	if (argc == 0 && *lockname != '/') {
+		fdlock = 1;
+		subj.subj_fd = -1;
+	} else {
+		fdlock = fdlock_implied(lockname, &subj.subj_fd);
+		if (argc == 0 && !fdlock) {
+			fprintf(stderr, "Expected fd, got '%s'\n", lockname);
+			usage();
+		}
+	}
+
+	if (fdlock) {
+		if (subj.subj_fd < 0) {
+			char *endp;
+
+			errno = 0;
+			subj.subj_fd = strtol(lockname, &endp, 10);
+			if (errno != 0 || *endp != '\0') {
+				fprintf(stderr, "Expected fd, got '%s'\n",
+				    lockname);
+				usage();
+			}
+		}
+
+		if (subj.subj_fd < 0 || subj.subj_fd > INT_MAX) {
+			fprintf(stderr, "fd '%ld' out of range\n",
+			    subj.subj_fd);
+			usage();
+		}
+	} else {
+		subj.subj_name = lockname;
+	}
+
 	if (waitsec > 0) {		/* Set up a timeout. */
 		struct sigaction act;
 
@@ -125,13 +205,13 @@ main(int argc, char **argv)
 	 * avoiding the separate step of waiting for the lock.  This
 	 * yields fairness and improved performance.
 	 */
-	lockfd = acquire_lock(lockname, flags | O_NONBLOCK, silent);
+	lockfd = acquire_lock(&subj, flags | O_NONBLOCK, silent);
 	while (lockfd == -1 && !timed_out && waitsec != 0) {
-		if (keep)
-			lockfd = acquire_lock(lockname, flags, silent);
+		if (keep || fdlock)
+			lockfd = acquire_lock(&subj, flags, silent);
 		else {
 			wait_for_lock(lockname);
-			lockfd = acquire_lock(lockname, flags | O_NONBLOCK,
+			lockfd = acquire_lock(&subj, flags | O_NONBLOCK,
 			    silent);
 		}
 	}
@@ -142,7 +222,15 @@ main(int argc, char **argv)
 			exit(EX_TEMPFAIL);
 		errx(EX_TEMPFAIL, "%s: already locked", lockname);
 	}
+
 	/* At this point, we own the lock. */
+
+	/* Nothing else to do for FD lock, just exit */
+	if (argc == 0) {
+		assert(fdlock);
+		return 0;
+	}
+
 	if (atexit(cleanup) == -1)
 		err(EX_OSERR, "atexit failed");
 	if ((child = fork()) == -1)
@@ -166,25 +254,34 @@ main(int argc, char **argv)
 }
 
 /*
- * Try to acquire a lock on the given file, creating the file if
+ * Try to acquire a lock on the given file/fd, creating the file if
  * necessary.  The flags argument is O_NONBLOCK or 0, depending on
  * whether we should wait for the lock.  Returns an open file descriptor
  * on success, or -1 on failure.
  */
 static int
-acquire_lock(const char *name, int flags, int silent)
+acquire_lock(union lock_subject *subj, int flags, int silent)
 {
 	int fd;
 
-	if ((fd = open(name, O_EXLOCK|flags, 0666)) == -1) {
+	if (fdlock) {
+		assert(subj->subj_fd >= 0 && subj->subj_fd <= INT_MAX);
+		fd = (int)subj->subj_fd;
+
+		if (flock(fd, LOCK_EX | LOCK_NB) == -1) {
+			if (errno == EAGAIN || errno == EINTR)
+				return (-1);
+			err(EX_CANTCREAT, "cannot lock fd %d", fd);
+		}
+	} else if ((fd = open(subj->subj_name, O_EXLOCK|flags, 0666)) == -1) {
 		if (errno == EAGAIN || errno == EINTR)
 			return (-1);
 		else if (errno == ENOENT && (flags & O_CREAT) == 0) {
 			if (!silent)
-				warn("%s", name);
+				warn("%s", subj->subj_name);
 			exit(EX_UNAVAILABLE);
 		}
-		err(EX_CANTCREAT, "cannot open %s", name);
+		err(EX_CANTCREAT, "cannot open %s", subj->subj_name);
 	}
 	return (fd);
 }
@@ -196,7 +293,7 @@ static void
 cleanup(void)
 {
 
-	if (keep)
+	if (keep || fdlock)
 		flock(lockfd, LOCK_UN);
 	else
 		unlink(lockname);
@@ -231,7 +328,8 @@ usage(void)
 {
 
 	fprintf(stderr,
-	    "usage: lockf [-knsw] [-t seconds] file command [arguments]\n");
+	    "usage: lockf [-knsw] [-t seconds] file command [arguments]\n"
+	    "       lockf [-s] [-t seconds] fd\n");
 	exit(EX_USAGE);
 }
 
diff --git a/usr.bin/lockf/tests/lockf_test.sh b/usr.bin/lockf/tests/lockf_test.sh
index 8696ab82a996..cc6938d2306e 100644
--- a/usr.bin/lockf/tests/lockf_test.sh
+++ b/usr.bin/lockf/tests/lockf_test.sh
@@ -60,6 +60,96 @@ basic_body()
 	atf_check test ! -e "testlock"
 }
 
+atf_test_case fdlock
+fdlock_body()
+{
+	# First, make sure we don't get a false positive -- existing uses with
+	# numeric filenames shouldn't switch to being fdlocks automatically.
+	atf_check lockf -k "9" sleep 0
+	atf_check test -e "9"
+	rm "9"
+
+	subexit_lockfail=1
+	subexit_created=2
+	subexit_lockok=3
+	subexit_concurrent=4
+	(
+		lockf -s -t 0 9
+		if [ $? -ne 0 ]; then
+			exit "$subexit_lockfail"
+		fi
+
+		if [ -e "9" ]; then
+			exit "$subexit_created"
+		fi
+	) 9> "testlock1"
+	rc=$?
+
+	atf_check test "$rc" -eq 0
+
+	sub_delay=5
+
+	# But is it actually locking?  Child 1 will acquire the lock and then
+	# signal that it's ok for the second child to try.  The second child
+	# will try to acquire the lock and fail immediately, signal that it
+	# tried, then try again with an indefinite timeout.  On that one, we'll
+	# just check how long we ended up waiting -- it should be at least
+	# $sub_delay.
+	(
+		lockf -s -t 0 /dev/fd/9
+		if [ $? -ne 0 ]; then
+			exit "$subexit_lockfail"
+		fi
+
+		# Signal
+		touch ".lock_acquired"
+
+		while [ ! -e ".lock_attempted" ]; do
+			sleep 0.5
+		done
+
+		sleep "$sub_delay"
+
+		if [ -e ".lock_acquired_again" ]; then
+			exit "$subexit_concurrent"
+		fi
+	) 9> "testlock2" &
+	lpid1=$!
+
+	(
+		while [ ! -e ".lock_acquired" ]; do
+			sleep 0.5
+		done
+
+		# Got the signal, try
+		lockf -s -t 0 9
+		if [ $? -ne "${EX_TEMPFAIL}" ]; then
+			exit "$subexit_lockok"
+		fi
+
+		touch ".lock_attempted"
+		start=$(date +"%s")
+		lockf -s 9
+		touch ".lock_acquired_again"
+		now=$(date +"%s")
+		elapsed=$((now - start))
+
+		if [ "$elapsed" -lt "$sub_delay" ]; then
+			exit "$subexit_concurrent"
+		fi
+	) 9> "testlock2" &
+	lpid2=$!
+
+	wait "$lpid1"
+	status1=$?
+
+	wait "$lpid2"
+	status2=$?
+
+	atf_check test "$status1" -eq 0
+	atf_check test "$status2" -eq 0
+}
+
 atf_test_case keep
 keep_body()
 {
@@ -141,6 +231,7 @@ atf_init_test_cases()
 {
 	atf_add_test_case badargs
 	atf_add_test_case basic
+	atf_add_test_case fdlock
 	atf_add_test_case keep
 	atf_add_test_case needfile
 	atf_add_test_case timeout