xref: /netbsd-src/usr.bin/progress/progress.c (revision 76c7fc5f6b13ed0b1508e6b313e88e59977ed78e)
1 /*	$NetBSD: progress.c,v 1.21 2015/01/17 10:57:51 gson Exp $ */
2 
3 /*-
4  * Copyright (c) 2003 The NetBSD Foundation, Inc.
5  * All rights reserved.
6  *
7  * This code is derived from software contributed to The NetBSD Foundation
8  * by John Hawkinson.
9  *
10  * Redistribution and use in source and binary forms, with or without
11  * modification, are permitted provided that the following conditions
12  * are met:
13  * 1. Redistributions of source code must retain the above copyright
14  *    notice, this list of conditions and the following disclaimer.
15  * 2. Redistributions in binary form must reproduce the above copyright
16  *    notice, this list of conditions and the following disclaimer in the
17  *    documentation and/or other materials provided with the distribution.
18  *
19  * THIS SOFTWARE IS PROVIDED BY THE NETBSD FOUNDATION, INC. AND CONTRIBUTORS
20  * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
21  * TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
22  * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE FOUNDATION OR CONTRIBUTORS
23  * BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
24  * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
25  * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
26  * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
27  * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29  * POSSIBILITY OF SUCH DAMAGE.
30  */
31 
32 #include <sys/cdefs.h>
33 #ifndef lint
34 __RCSID("$NetBSD: progress.c,v 1.21 2015/01/17 10:57:51 gson Exp $");
35 #endif				/* not lint */
36 
37 #include <sys/types.h>
38 #include <sys/ioctl.h>
39 #include <sys/stat.h>
40 #include <sys/wait.h>
41 
42 #include <err.h>
43 #include <errno.h>
44 #include <fcntl.h>
45 #include <inttypes.h>
46 #include <limits.h>
47 #include <signal.h>
48 #include <stdio.h>
49 #include <stdlib.h>
50 #include <string.h>
51 #include <unistd.h>
52 
53 #define GLOBAL			/* force GLOBAL decls in progressbar.h to be
54 				 * declared */
55 #include "progressbar.h"
56 
57 static void broken_pipe(int unused);
58 __dead static void usage(void);
59 
60 static void
61 broken_pipe(int unused)
62 {
63 	signal(SIGPIPE, SIG_DFL);
64 	progressmeter(1);
65 	kill(getpid(), SIGPIPE);
66 }
67 
68 static void
69 usage(void)
70 {
71 	fprintf(stderr,
72 	    "usage: %s [-ez] [-b buffersize] [-f file] [-l length]\n"
73 	    "       %*.s [-p prefix] cmd [args...]\n",
74 	    getprogname(), (int) strlen(getprogname()), "");
75 	exit(EXIT_FAILURE);
76 }
77 
78 
79 int
80 main(int argc, char *argv[])
81 {
82 	char *fb_buf;
83 	char *infile = NULL;
84 	pid_t pid = 0, gzippid = 0, deadpid;
85 	int ch, fd, outpipe[2];
86 	int ws, gzipstat, cmdstat;
87 	int eflag = 0, lflag = 0, zflag = 0;
88 	ssize_t nr, nw, off;
89 	size_t buffersize;
90 	struct stat statb;
91 	struct ttysize ts;
92 
93 	setprogname(argv[0]);
94 
95 	/* defaults: Read from stdin, 0 filesize (no completion estimate) */
96 	fd = STDIN_FILENO;
97 	filesize = 0;
98 	buffersize = 64 * 1024;
99 	prefix = NULL;
100 
101 	while ((ch = getopt(argc, argv, "b:ef:l:p:z")) != -1)
102 		switch (ch) {
103 		case 'b':
104 			buffersize = (size_t) strsuftoll("buffer size", optarg,
105 			    0, SSIZE_MAX);
106 			break;
107 		case 'e':
108 			eflag++;
109 			break;
110 		case 'f':
111 			infile = optarg;
112 			break;
113 		case 'l':
114 			lflag++;
115 			filesize = strsuftoll("input size", optarg, 0,
116 			    LLONG_MAX);
117 			break;
118 		case 'p':
119 			prefix = optarg;
120 			break;
121 		case 'z':
122 			zflag++;
123 			break;
124 		case '?':
125 		default:
126 			usage();
127 			/* NOTREACHED */
128 		}
129 	argc -= optind;
130 	argv += optind;
131 
132 	if (argc < 1)
133 		usage();
134 
135 	if (infile && (fd = open(infile, O_RDONLY, 0)) < 0)
136 		err(1, "%s", infile);
137 
138 	/* stat() to get the filesize unless overridden, or -z */
139 	if (!zflag && !lflag && (fstat(fd, &statb) == 0)) {
140 		if (S_ISFIFO(statb.st_mode)) {
141 			/* stat(2) on pipe may return only the
142 			 * first few bytes with more coming.
143 			 * Don't trust!
144 			 */
145 		} else {
146 			filesize = statb.st_size;
147 		}
148 	}
149 
150 	/* gzip -l the file if we have the name and -z is given */
151 	if (zflag && !lflag && infile != NULL) {
152 		FILE *gzipsizepipe;
153 		char buf[256], *cp, *cmd;
154 
155 		/*
156 		 * Read second word of last line of gzip -l output. Looks like:
157 		 * % gzip -l ../etc.tgz
158 		 *   compressed uncompressed  ratio uncompressed_name
159 		 * 	 119737       696320  82.8% ../etc.tar
160 		 */
161 
162 		asprintf(&cmd, "gzip -l %s", infile);
163 		if ((gzipsizepipe = popen(cmd, "r")) == NULL)
164 			err(1, "reading compressed file length");
165 		for (; fgets(buf, 256, gzipsizepipe) != NULL;)
166 		    continue;
167 		strtoimax(buf, &cp, 10);
168 		filesize = strtoimax(cp, NULL, 10);
169 		if (pclose(gzipsizepipe) < 0)
170 			err(1, "closing compressed file length pipe");
171 		free(cmd);
172 	}
173 	/* Pipe input through gzip -dc if -z is given */
174 	if (zflag) {
175 		int gzippipe[2];
176 
177 		if (pipe(gzippipe) < 0)
178 			err(1, "gzip pipe");
179 		gzippid = fork();
180 		if (gzippid < 0)
181 			err(1, "fork for gzip");
182 
183 		if (gzippid) {
184 			/* parent */
185 			dup2(gzippipe[0], fd);
186 			close(gzippipe[0]);
187 			close(gzippipe[1]);
188 		} else {
189 			dup2(gzippipe[1], STDOUT_FILENO);
190 			dup2(fd, STDIN_FILENO);
191 			close(gzippipe[0]);
192 			close(gzippipe[1]);
193 			if (execlp("gzip", "gzip", "-dc", NULL))
194 				err(1, "exec()ing gzip");
195 		}
196 	}
197 
198 	/* Initialize progressbar.c's global state */
199 	bytes = 0;
200 	progress = 1;
201 	ttyout = eflag ? stderr : stdout;
202 
203 	if (ioctl(fileno(ttyout), TIOCGSIZE, &ts) == -1)
204 		ttywidth = 80;
205 	else
206 		ttywidth = ts.ts_cols;
207 
208 	fb_buf = malloc(buffersize);
209 	if (fb_buf == NULL)
210 		err(1, "malloc for buffersize");
211 
212 	if (pipe(outpipe) < 0)
213 		err(1, "output pipe");
214 	pid = fork();
215 	if (pid < 0)
216 		err(1, "fork for output pipe");
217 
218 	if (pid == 0) {
219 		/* child */
220 		dup2(outpipe[0], STDIN_FILENO);
221 		close(outpipe[0]);
222 		close(outpipe[1]);
223 		execvp(argv[0], argv);
224 		err(1, "could not exec %s", argv[0]);
225 	}
226 	close(outpipe[0]);
227 
228 	signal(SIGPIPE, broken_pipe);
229 	progressmeter(-1);
230 
231 	while (1) {
232 		do {
233 			nr = read(fd, fb_buf, buffersize);
234 		} while (nr < 0 && errno == EINTR);
235 		if (nr <= 0)
236 			break;
237 		for (off = 0; nr; nr -= nw, off += nw, bytes += nw)
238 			if ((nw = write(outpipe[1], fb_buf + off,
239 			    (size_t) nr)) < 0) {
240 				progressmeter(1);
241 				err(1, "writing %u bytes to output pipe",
242 							(unsigned) nr);
243 			}
244 	}
245 	close(outpipe[1]);
246 
247 	gzipstat = 0;
248 	cmdstat = 0;
249 	while (pid || gzippid) {
250 		deadpid = wait(&ws);
251 		/*
252 		 * We need to exit with an error if the command (or gzip)
253 		 * exited abnormally.
254 		 * Unfortunately we can't generate a true 'exited by signal'
255 		 * error without sending the signal to ourselves :-(
256 		 */
257 		ws = WIFSIGNALED(ws) ? WTERMSIG(ws) : WEXITSTATUS(ws);
258 
259 		if (deadpid != -1 && errno == EINTR)
260 			continue;
261 		if (deadpid == pid) {
262 			pid = 0;
263 			cmdstat = ws;
264 			continue;
265 		}
266 		if (deadpid == gzippid) {
267 			gzippid = 0;
268 			gzipstat = ws;
269 			continue;
270 		}
271 		break;
272 	}
273 
274 	progressmeter(1);
275 	signal(SIGPIPE, SIG_DFL);
276 
277 	free(fb_buf);
278 
279 	exit(cmdstat ? cmdstat : gzipstat);
280 }
281