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