xref: /netbsd-src/usr.bin/progress/progress.c (revision ead2c0eee3abe6bcf08c63bfc78eb8a93a579b2b)
1 /*	$NetBSD: progress.c,v 1.19 2011/09/16 15:39:28 joerg 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.19 2011/09/16 15:39:28 joerg Exp $");
35 #endif				/* not lint */
36 
37 #include <sys/types.h>
38 #include <sys/param.h>
39 #include <sys/socket.h>
40 #include <sys/ioctl.h>
41 #include <sys/time.h>
42 #include <sys/wait.h>
43 #include <netinet/in.h>
44 #include <arpa/ftp.h>
45 
46 #include <err.h>
47 #include <errno.h>
48 #include <fcntl.h>
49 #include <glob.h>
50 #include <signal.h>
51 #include <inttypes.h>
52 #include <limits.h>
53 #include <netdb.h>
54 #include <stdio.h>
55 #include <stdlib.h>
56 #include <string.h>
57 #include <termios.h>
58 #include <time.h>
59 #include <tzfile.h>
60 #include <unistd.h>
61 
62 #define GLOBAL			/* force GLOBAL decls in progressbar.h to be
63 				 * declared */
64 #include "progressbar.h"
65 
66 static void broken_pipe(int unused);
67 __dead static void usage(void);
68 
69 static void
70 broken_pipe(int unused)
71 {
72 	signal(SIGPIPE, SIG_DFL);
73 	progressmeter(1);
74 	kill(getpid(), SIGPIPE);
75 }
76 
77 static void
78 usage(void)
79 {
80 	fprintf(stderr,
81 	    "usage: %s [-ez] [-b buffersize] [-f file] [-l length]\n"
82 	    "       %*.s [-p prefix] cmd [args...]\n",
83 	    getprogname(), (int) strlen(getprogname()), "");
84 	exit(EXIT_FAILURE);
85 }
86 
87 
88 int
89 main(int argc, char *argv[])
90 {
91 	char *fb_buf;
92 	char *infile = NULL;
93 	pid_t pid = 0, gzippid = 0, deadpid;
94 	int ch, fd, outpipe[2];
95 	int ws, gzipstat, cmdstat;
96 	int eflag = 0, lflag = 0, zflag = 0;
97 	ssize_t nr, nw, off;
98 	size_t buffersize;
99 	struct stat statb;
100 	struct ttysize ts;
101 
102 	setprogname(argv[0]);
103 
104 	/* defaults: Read from stdin, 0 filesize (no completion estimate) */
105 	fd = STDIN_FILENO;
106 	filesize = 0;
107 	buffersize = 64 * 1024;
108 	prefix = NULL;
109 
110 	while ((ch = getopt(argc, argv, "b:ef:l:p:z")) != -1)
111 		switch (ch) {
112 		case 'b':
113 			buffersize = (size_t) strsuftoll("buffer size", optarg,
114 			    0, SSIZE_MAX);
115 			break;
116 		case 'e':
117 			eflag++;
118 			break;
119 		case 'f':
120 			infile = optarg;
121 			break;
122 		case 'l':
123 			lflag++;
124 			filesize = strsuftoll("input size", optarg, 0,
125 			    LLONG_MAX);
126 			break;
127 		case 'p':
128 			prefix = optarg;
129 			break;
130 		case 'z':
131 			zflag++;
132 			break;
133 		case '?':
134 		default:
135 			usage();
136 			/* NOTREACHED */
137 		}
138 	argc -= optind;
139 	argv += optind;
140 
141 	if (argc < 1)
142 		usage();
143 
144 	if (infile && (fd = open(infile, O_RDONLY, 0)) < 0)
145 		err(1, "%s", infile);
146 
147 	/* stat() to get the filesize unless overridden, or -z */
148 	if (!zflag && !lflag && (fstat(fd, &statb) == 0)) {
149 		if (S_ISFIFO(statb.st_mode)) {
150 			/* stat(2) on pipe may return only the
151 			 * first few bytes with more coming.
152 			 * Don't trust!
153 			 */
154 		} else {
155 			filesize = statb.st_size;
156 		}
157 	}
158 
159 	/* gzip -l the file if we have the name and -z is given */
160 	if (zflag && !lflag && infile != NULL) {
161 		FILE *gzipsizepipe;
162 		char buf[256], *cp, *cmd;
163 
164 		/*
165 		 * Read second word of last line of gzip -l output. Looks like:
166 		 * % gzip -l ../etc.tgz
167 		 *   compressed uncompressed  ratio uncompressed_name
168 		 * 	 119737       696320  82.8% ../etc.tar
169 		 */
170 
171 		asprintf(&cmd, "gzip -l %s", infile);
172 		if ((gzipsizepipe = popen(cmd, "r")) == NULL)
173 			err(1, "reading compressed file length");
174 		for (; fgets(buf, 256, gzipsizepipe) != NULL;)
175 		    continue;
176 		strtoimax(buf, &cp, 10);
177 		filesize = strtoimax(cp, NULL, 10);
178 		if (pclose(gzipsizepipe) < 0)
179 			err(1, "closing compressed file length pipe");
180 		free(cmd);
181 	}
182 	/* Pipe input through gzip -dc if -z is given */
183 	if (zflag) {
184 		int gzippipe[2];
185 
186 		if (pipe(gzippipe) < 0)
187 			err(1, "gzip pipe");
188 		gzippid = fork();
189 		if (gzippid < 0)
190 			err(1, "fork for gzip");
191 
192 		if (gzippid) {
193 			/* parent */
194 			dup2(gzippipe[0], fd);
195 			close(gzippipe[0]);
196 			close(gzippipe[1]);
197 		} else {
198 			dup2(gzippipe[1], STDOUT_FILENO);
199 			dup2(fd, STDIN_FILENO);
200 			close(gzippipe[0]);
201 			close(gzippipe[1]);
202 			if (execlp("gzip", "gzip", "-dc", NULL))
203 				err(1, "exec()ing gzip");
204 		}
205 	}
206 
207 	/* Initialize progressbar.c's global state */
208 	bytes = 0;
209 	progress = 1;
210 	ttyout = eflag ? stderr : stdout;
211 
212 	if (ioctl(fileno(ttyout), TIOCGSIZE, &ts) == -1)
213 		ttywidth = 80;
214 	else
215 		ttywidth = ts.ts_cols;
216 
217 	fb_buf = malloc(buffersize);
218 	if (fb_buf == NULL)
219 		err(1, "malloc for buffersize");
220 
221 	if (pipe(outpipe) < 0)
222 		err(1, "output pipe");
223 	pid = fork();
224 	if (pid < 0)
225 		err(1, "fork for output pipe");
226 
227 	if (pid == 0) {
228 		/* child */
229 		dup2(outpipe[0], STDIN_FILENO);
230 		close(outpipe[0]);
231 		close(outpipe[1]);
232 		execvp(argv[0], argv);
233 		err(1, "could not exec %s", argv[0]);
234 	}
235 	close(outpipe[0]);
236 
237 	signal(SIGPIPE, broken_pipe);
238 	progressmeter(-1);
239 
240 	while ((nr = read(fd, fb_buf, buffersize)) > 0)
241 		for (off = 0; nr; nr -= nw, off += nw, bytes += nw)
242 			if ((nw = write(outpipe[1], fb_buf + off,
243 			    (size_t) nr)) < 0) {
244 				progressmeter(1);
245 				err(1, "writing %u bytes to output pipe",
246 							(unsigned) nr);
247 			}
248 	close(outpipe[1]);
249 
250 	gzipstat = 0;
251 	cmdstat = 0;
252 	while (pid || gzippid) {
253 		deadpid = wait(&ws);
254 		/*
255 		 * We need to exit with an error if the command (or gzip)
256 		 * exited abnormally.
257 		 * Unfortunately we can't generate a true 'exited by signal'
258 		 * error without sending the signal to ourselves :-(
259 		 */
260 		ws = WIFSIGNALED(ws) ? WTERMSIG(ws) : WEXITSTATUS(ws);
261 
262 		if (deadpid != -1 && errno == EINTR)
263 			continue;
264 		if (deadpid == pid) {
265 			pid = 0;
266 			cmdstat = ws;
267 			continue;
268 		}
269 		if (deadpid == gzippid) {
270 			gzippid = 0;
271 			gzipstat = ws;
272 			continue;
273 		}
274 		break;
275 	}
276 
277 	progressmeter(1);
278 	signal(SIGPIPE, SIG_DFL);
279 
280 	free(fb_buf);
281 
282 	exit(cmdstat ? cmdstat : gzipstat);
283 }
284