Fix batchd subdir and more

1 files changed, 21 insertions, 5 deletions

diff --git a/runfd.c b/runfd.c
index 59b9d8c..bebe0ec 100644
--- a/runfd.c
+++ b/runfd.c
@@ -11,6 +11,7 @@
  */
 
 #define DEFAULT_FD 3
+#define BIN_HDR_LEN 2
 
 char *args[] = { "batchrun", NULL };
 
@@ -32,10 +33,17 @@ set_async(int fd, int on)
 		err(1, "fcntl()");
 }
 
+void
+set_cloexec(int fd, int on)
+{
+	if (fcntl(fd, F_SETFD, on ? FD_CLOEXEC : 0) == -1)
+		err(1, "fcntl()");
+}
+
 int
 main(int argc, char **argv)
 {
-	char buf[2], *end = NULL;
+	char buf[BIN_HDR_LEN], *end = NULL;
 	ssize_t ret;
 	int tries = 0, fd = DEFAULT_FD;
 
@@ -46,6 +54,7 @@ main(int argc, char **argv)
 			errx(1, "Filedescriptor supplied (\"%s\") is not a number", argv[1]);
 	}
 
+	set_cloexec(fd, 1);
 	set_async(fd, 1);
 
 retry:
@@ -53,9 +62,9 @@ retry:
 		errx(1, "Failed to read file header");
 
 	ret = pread(fd, buf, sizeof(buf), 0);
-	if (ret == 2) {
-		if (strncmp(buf, "#!", 2) != 0) {
-			fcntl(fd, F_SETFD, FD_CLOEXEC);
+	if (ret == sizeof(buf)) {
+		if (strncmp(buf, "#!", sizeof(buf)) == 0) {
+			set_cloexec(fd, 0);
 		}
 	} else if (ret == -1) {
 		if (errno == EINTR)
@@ -63,7 +72,14 @@ retry:
 		else if (errno != EAGAIN)
 			err(1, "pread()");
 	} else if (ret != 0) {
-		goto retry;
+		/**
+		 * We take a lazy approach here and only want to know if the
+		 * file has more contents (and don't care about errors).
+		 */
+		ret = pread(fd, buf + ret, sizeof(buf) - ret, ret);
+		if (ret != 0) {
+			goto retry;
+		}
 	}
 
 	set_async(fd, 0);