utils-std

commit: 0c5c782733c7ac1cf4bd19177d2e5ec7ef269846
parent 7a561b087817f7495ad13a840e57eba79eff67c3
Author: Haelwenn (lanodan) Monnier <contact@hacktivis.me>
Date:   Mon, 27 May 2024 09:32:07 +0200

cmd/head: Refactor for read+write and sane buffer size

Diffstat:
M
cmd/head.c
145
++++++++++++++++++++++++++++++++++++++++++++++++++++---------------------------

1 file changed, 95 insertions(+), 50 deletions(-)
diff --git a/cmd/head.c b/cmd/head.c
@@ -8,44 +8,109 @@
 
 #include <assert.h>
 #include <errno.h>
-#include <stdio.h>  // fprintf, fopen
+#include <fcntl.h>  // open
+#include <stdio.h>  // fprintf, fopen, fread
 #include <stdlib.h> // strtoul
 #include <string.h> // strerror
 #include <unistd.h> // read
 
+#define MIN(a, b) (((a) < (b)) ? (a) : (b))
+
 static char *header_fmt = "==> %s <==\n";
 
 char *argv0 = "head";
 
 size_t lines = 10;
+size_t bytes = 0;
 
 char *buf = NULL;
 size_t buflen = 0;
 
 static int
-copy_bytes(FILE *in, char *filename)
+copy_bytes(char *filename)
 {
-	assert(errno == 0);
-	size_t eread = fread(buf, sizeof(char), buflen, in);
-	if(errno != 0)
+	if(buflen == 0)
+	{
+		// Note: de-allocated in main()
+		buf = malloc(BUFSIZ);
+
+		if(buf == NULL)
+		{
+			fprintf(stderr, "head: Failed to allocate buffer: %s\n", strerror(errno));
+			return 1;
+		}
+
+		buflen = BUFSIZ;
+	}
+
+	int fd = -1;
+	if(filename[0] == '-' && filename[1] == 0)
 	{
-		fprintf(stderr, "head: Failed reading file '%s': %s\n", filename, strerror(errno));
-		return 1;
+		fd = STDIN_FILENO;
+	}
+	else
+	{
+		fd = open(filename, O_RDONLY | O_CLOEXEC | O_NOCTTY);
+		if(fd < 0)
+		{
+			fprintf(stderr, "head: Failed opening file '%s': %s\n", filename, strerror(errno));
+			return 1;
+		}
 	}
 
-	if(write(STDOUT_FILENO, buf, eread) < 0)
+	int err = 0;
+
+	for(size_t i = 0; i < bytes;)
 	{
-		fprintf(
-		    stderr, "head: Failed writing line from '%s' to stdout: %s\n", filename, strerror(errno));
-		return 1;
+		ssize_t nread = read(fd, buf, MIN(buflen, bytes - i));
+		if(nread < 0)
+		{
+			fprintf(stderr, "head: Failed reading file '%s': %s\n", filename, strerror(errno));
+			err = 1;
+			break;
+		}
+
+		if(write(STDOUT_FILENO, buf, nread) < 0)
+		{
+			fprintf(
+			    stderr, "head: Failed writing line from '%s' to stdout: %s\n", filename, strerror(errno));
+			err = 1;
+			break;
+		}
+
+		i += nread;
 	}
 
-	return 0;
+	if(fd != STDIN_FILENO)
+		if(close(fd) != 0)
+		{
+			fprintf(stderr, "head: Failed closing file '%s': %s\n", filename, strerror(errno));
+			return 1;
+		}
+
+	return err;
 }
 
 static int
-copy_lines(FILE *in, char *filename)
+copy_lines(char *filename)
 {
+	int err = 0;
+	FILE *in = NULL;
+
+	if(filename[0] == '-' && filename[1] == 0)
+	{
+		in = stdin;
+	}
+	else
+	{
+		in = fopen(filename, "r");
+		if(in == NULL)
+		{
+			fprintf(stderr, "head: Failed opening file '%s': %s\n", filename, strerror(errno));
+			return 1;
+		}
+	}
+
 	for(size_t i = 0; i < lines; i++)
 	{
 		assert(errno == 0);
@@ -55,18 +120,27 @@ copy_lines(FILE *in, char *filename)
 			if(errno == 0) break;
 
 			fprintf(stderr, "head: Failed reading line from '%s': %s\n", filename, strerror(errno));
-			return 1;
+			err = 1;
+			break;
 		}
 
 		if(write(STDOUT_FILENO, buf, nread) < 0)
 		{
 			fprintf(
 			    stderr, "head: Failed writing line from '%s' to stdout: %s\n", filename, strerror(errno));
-			return 1;
+			err = 1;
+			break;
 		}
 	}
 
-	return 0;
+	if(in != stdin)
+		if(fclose(in) != 0)
+		{
+			fprintf(stderr, "head: Failed closing file '%s': %s\n", filename, strerror(errno));
+			return 1;
+		}
+
+	return err;
 }
 
 static void
@@ -78,7 +152,7 @@ usage()
 int
 main(int argc, char *argv[])
 {
-	int (*copy_action)(FILE *in, char *filename) = copy_lines;
+	int (*copy_action)(char *filename) = copy_lines;
 
 	int print_header = 0;
 	int c = -1;
@@ -107,7 +181,7 @@ main(int argc, char *argv[])
 			if(endptr != NULL && *endptr != 0)
 				if(apply_size_suffix(&size, endptr) != 0) return 1;
 
-			buflen = size;
+			bytes = (size_t)size;
 			copy_action = &copy_bytes;
 			break;
 		}
@@ -140,20 +214,9 @@ main(int argc, char *argv[])
 	argc -= optind;
 	argv += optind;
 
-	if(buflen != 0)
-	{
-		buf = malloc(buflen);
-
-		if(buf == NULL)
-		{
-			fprintf(stderr, "head: Failed to allocate buffer: %s\n", strerror(errno));
-			return 1;
-		}
-	}
-
 	if(argc <= 0)
 	{
-		char *filename = "<stdin>";
+		char *filename = "-";
 
 		if(print_header == 1)
 		{
@@ -161,7 +224,7 @@ main(int argc, char *argv[])
 			fflush(stdout);
 		}
 
-		int err = copy_action(stdin, filename);
+		int err = copy_action(filename);
 
 		if(buflen != 0) free(buf);
 
@@ -183,26 +246,8 @@ main(int argc, char *argv[])
 			fflush(stdout);
 		}
 
-		if(filename[0] == '-' && filename[1] == 0)
+		if(copy_action(filename) != 0)
 		{
-			if(copy_action(stdin, "<stdin>") != 0) break;
-
-			continue;
-		}
-
-		FILE *in = fopen(filename, "r");
-		if(in == NULL)
-		{
-			fprintf(stderr, "head: Failed opening file '%s': %s\n", filename, strerror(errno));
-			err = 1;
-			break;
-		}
-
-		if(copy_action(in, filename) != 0) return 1;
-
-		if(fclose(in) != 0)
-		{
-			fprintf(stderr, "head: Failed closing file '%s': %s\n", filename, strerror(errno));
 			err = 1;
 			break;
 		}