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