xref: /freebsd/tests/sys/kern/tty/readsz.c (revision 7937bfbc0ca53fe7cdd0d54414f9296e273a518e)
1 /*-
2  * Copyright (c) 2024 Kyle Evans <kevans@FreeBSD.org>
3  *
4  * SPDX-License-Identifier: BSD-2-Clause
5  */
6 
7 #include <sys/param.h>
8 
9 #include <err.h>
10 #include <errno.h>
11 #include <limits.h>
12 #include <stdbool.h>
13 #include <stdint.h>
14 #include <stdio.h>
15 #include <stdlib.h>
16 #include <unistd.h>
17 
18 static void
19 usage(void)
20 {
21 
22 	fprintf(stderr, "usage: %s [-b bytes | -c lines | -e] [-s buffer-size]\n",
23 	    getprogname());
24 	exit(1);
25 }
26 
27 int
28 main(int argc, char *argv[])
29 {
30 	char *buf;
31 	const char *errstr;
32 	size_t bufsz = 0, reps;
33 	ssize_t ret;
34 	enum { MODE_BYTES, MODE_COUNT, MODE_EOF } mode;
35 	int ch;
36 
37 	/*
38 	 * -b specifies number of bytes.
39 	 * -c specifies number of read() calls.
40 	 * -e specifies eof (default)
41 	 * -s to pass a buffer size
42 	 *
43 	 * Reading N lines is the same as -c with a high buffer size.
44 	 */
45 	mode = MODE_EOF;
46 	while ((ch = getopt(argc, argv, "b:c:es:")) != -1) {
47 		switch (ch) {
48 		case 'b':
49 			mode = MODE_BYTES;
50 			reps = strtonum(optarg, 0, SSIZE_MAX, &errstr);
51 			if (errstr != NULL)
52 				errx(1, "strtonum: %s", errstr);
53 			break;
54 		case 'c':
55 			mode = MODE_COUNT;
56 			reps = strtonum(optarg, 1, SSIZE_MAX, &errstr);
57 			if (errstr != NULL)
58 				errx(1, "strtonum: %s", errstr);
59 			break;
60 		case 'e':
61 			mode = MODE_EOF;
62 			break;
63 		case 's':
64 			bufsz = strtonum(optarg, 1, SSIZE_MAX, &errstr);
65 			if (errstr != NULL)
66 				errx(1, "strtonum: %s", errstr);
67 			break;
68 		default:
69 			usage();
70 		}
71 	}
72 
73 	if (bufsz == 0) {
74 		if (mode == MODE_BYTES)
75 			bufsz = reps;
76 		else
77 			bufsz = LINE_MAX;
78 	}
79 
80 	buf = malloc(bufsz);
81 	if (buf == NULL)
82 		err(1, "malloc");
83 
84 	for (;;) {
85 		size_t readsz;
86 
87 		/*
88 		 * Be careful not to over-read if we're in byte-mode.  In every other
89 		 * mode, we'll read as much as we can.
90 		 */
91 		if (mode == MODE_BYTES)
92 			readsz = MIN(bufsz, reps);
93 		else
94 			readsz = bufsz;
95 
96 		ret = read(STDIN_FILENO, buf, readsz);
97 		if (ret == -1 && errno == EINTR)
98 			continue;
99 		if (ret == -1)
100 			err(1, "read");
101 		if (ret == 0) {
102 			if (mode == MODE_EOF)
103 				return (0);
104 			errx(1, "premature EOF");
105 		}
106 
107 		/* Write out what we've got */
108 		write(STDOUT_FILENO, buf, ret);
109 
110 		/*
111 		 * Bail out if we've hit our metric (byte mode / count mode).
112 		 */
113 		switch (mode) {
114 		case MODE_BYTES:
115 			reps -= ret;
116 			if (reps == 0)
117 				return (0);
118 			break;
119 		case MODE_COUNT:
120 			reps--;
121 			if (reps == 0)
122 				return (0);
123 			break;
124 		default:
125 			break;
126 		}
127 	}
128 
129 	return (0);
130 }
131