utils-std

head.c (5551B)

// utils-std: Collection of commonly available Unix tools
// SPDX-FileCopyrightText: 2017 Haelwenn (lanodan) Monnier <contact+utils@hacktivis.me>
// SPDX-License-Identifier: MPL-2.0
#define _POSIX_C_SOURCE 200809L
#include "../lib/truncation.h" // apply_size_suffix
#include <assert.h>
#include <ctype.h> // isdigit
#include <errno.h>
#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 const char *header_fmt = "==> %s <==\n";
const char *argv0 = "head";
size_t lines = 10;
size_t bytes = 0;
char *buf = NULL;
size_t buflen = 0;
int delim = '\n';
static int
copy_bytes(const char *filename)
{
	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)
	{
		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;
		}
	}
	int err = 0;
	for(size_t i = 0; i < bytes;)
	{
		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(nread == 0) 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;
	}
	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(const 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);
		ssize_t nread = getdelim(&buf, &buflen, delim, in);
		if(nread < 0)
		{
			if(errno == 0) break;
			fprintf(stderr, "head: Failed reading line from '%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;
		}
	}
	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
usage(void)
{
	fprintf(stderr, "Usage: head [-qvz] [-c size | -n num | -num] [file...]\n");
}
int
main(int argc, char *argv[])
{
	int (*copy_action)(const char *filename) = copy_lines;
	int print_header = 0;
	while(optind < argc)
	{
		/* Skip getopt(3) for parsing -num option due to how it works:
		 * - optstring = "12"; args="head -12" considered the same as -1 -2
		 * - optstring = "1:"; args="head -12" considered the same as -1 2
		 *
		 * And we need args="head -12" as "12"
		 */
		if(argv[optind] && argv[optind][0] == '-' && isdigit(argv[optind][1]))
		{
			assert(errno == 0);
			char *arg = argv[optind] + 1;
			char *endptr = NULL;
			lines = strtoul(arg, &endptr, 0);
			if(!(errno == 0 && endptr != NULL && *endptr == '\0'))
			{
				fprintf(stderr, "head: Error while parsing number for `-%s`: %s\n", arg, strerror(errno));
				return 1;
			}
			optind++;
			continue;
		}
		int c = getopt(argc, argv, ":qvc:n:z");
		if(c == -1) break;
		switch(c)
		{
		case 'q':
			print_header = -1;
			break;
		case 'v':
			print_header = 1;
			break;
		case 'c':
		{
			assert(errno == 0);
			char *endptr = NULL;
			unsigned long size = strtoul(optarg, &endptr, 0);
			if(errno != 0)
			{
				fprintf(
				    stderr, "head: Error while parsing number for `-n %s`: %s\n", optarg, strerror(errno));
				return 1;
			}
			if(endptr != NULL && *endptr != 0)
				if(apply_size_suffix(&size, endptr) != 0) return 1;
			bytes = size;
			copy_action = &copy_bytes;
			break;
		}
		case 'n':
		{
			assert(errno == 0);
			char *endptr = NULL;
			lines = strtoul(optarg, &endptr, 0);
			if(!(errno == 0 && endptr != NULL && *endptr == '\0'))
			{
				fprintf(
				    stderr, "head: Error while parsing number for `-n %s`: %s\n", optarg, strerror(errno));
				return 1;
			}
			break;
		}
		case 'z':
			delim = '\0';
			break;
		case ':':
			fprintf(stderr, "head: Error: Missing operand for option: '-%c'\n", optopt);
			usage();
			return 1;
		case '?':
			fprintf(stderr, "head: Error: Unrecognised option: '-%c'\n", optopt);
			usage();
			return 1;
		default:
			abort();
		}
	}
	argc -= optind;
	argv += optind;
	if(argc <= 0)
	{
		const char *filename = "-";
		if(print_header == 1)
		{
			printf(header_fmt, filename);
			fflush(stdout);
		}
		int err = copy_action(filename);
		if(buflen != 0) free(buf);
		return err;
	}
	if(print_header == 0) print_header = argc > 1;
	int err = 0;
	for(int i = 0; i < argc; i++)
	{
		char *filename = argv[i];
		if(print_header > 0)
		{
			printf(header_fmt, filename);
			header_fmt = "\n==> %s <==\n";
			fflush(stdout);
		}
		if(copy_action(filename) != 0)
		{
			err = 1;
			break;
		}
	}
	if(buflen != 0) free(buf);
	return err;
}