1 /*
2  * prog_util.c - utility functions for programs
3  *
4  * Copyright 2016 Eric Biggers
5  *
6  * Permission is hereby granted, free of charge, to any person
7  * obtaining a copy of this software and associated documentation
8  * files (the "Software"), to deal in the Software without
9  * restriction, including without limitation the rights to use,
10  * copy, modify, merge, publish, distribute, sublicense, and/or sell
11  * copies of the Software, and to permit persons to whom the
12  * Software is furnished to do so, subject to the following
13  * conditions:
14  *
15  * The above copyright notice and this permission notice shall be
16  * included in all copies or substantial portions of the Software.
17  *
18  * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
19  * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
20  * OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
21  * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
22  * HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
23  * WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
24  * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
25  * OTHER DEALINGS IN THE SOFTWARE.
26  */
27 
28 #include "prog_util.h"
29 
30 #include <errno.h>
31 #include <fcntl.h>
32 #include <stdarg.h>
33 #ifdef _WIN32
34 #  include <windows.h>
35 #else
36 #  include <unistd.h>
37 #  include <sys/mman.h>
38 #endif
39 
40 #ifndef O_BINARY
41 #  define O_BINARY 0
42 #endif
43 #ifndef O_SEQUENTIAL
44 #  define O_SEQUENTIAL 0
45 #endif
46 #ifndef O_NOFOLLOW
47 #  define O_NOFOLLOW 0
48 #endif
49 #ifndef O_NONBLOCK
50 #  define O_NONBLOCK 0
51 #endif
52 #ifndef O_NOCTTY
53 #  define O_NOCTTY 0
54 #endif
55 
56 /* The invocation name of the program (filename component only) */
57 const tchar *prog_invocation_name;
58 
59 static void
do_msg(const char * format,bool with_errno,va_list va)60 do_msg(const char *format, bool with_errno, va_list va)
61 {
62 	int saved_errno = errno;
63 
64 	fprintf(stderr, "%"TS": ", prog_invocation_name);
65 	vfprintf(stderr, format, va);
66 	if (with_errno)
67 		fprintf(stderr, ": %s\n", strerror(saved_errno));
68 	else
69 		fprintf(stderr, "\n");
70 
71 	errno = saved_errno;
72 }
73 
74 /* Print a message to standard error */
75 void
msg(const char * format,...)76 msg(const char *format, ...)
77 {
78 	va_list va;
79 
80 	va_start(va, format);
81 	do_msg(format, false, va);
82 	va_end(va);
83 }
84 
85 /* Print a message to standard error, including a description of errno */
86 void
msg_errno(const char * format,...)87 msg_errno(const char *format, ...)
88 {
89 	va_list va;
90 
91 	va_start(va, format);
92 	do_msg(format, true, va);
93 	va_end(va);
94 }
95 
96 /* malloc() wrapper */
97 void *
xmalloc(size_t size)98 xmalloc(size_t size)
99 {
100 	void *p = malloc(size);
101 	if (p == NULL && size == 0)
102 		p = malloc(1);
103 	if (p == NULL)
104 		msg("Out of memory");
105 	return p;
106 }
107 
108 /*
109  * Retrieve a pointer to the filename component of the specified path.
110  *
111  * Note: this does not modify the path.  Therefore, it is not guaranteed to work
112  * properly for directories, since a path to a directory might have trailing
113  * slashes.
114  */
115 static const tchar *
get_filename(const tchar * path)116 get_filename(const tchar *path)
117 {
118 	const tchar *slash = tstrrchr(path, '/');
119 #ifdef _WIN32
120 	const tchar *backslash = tstrrchr(path, '\\');
121 	if (backslash != NULL && (slash == NULL || backslash > slash))
122 		slash = backslash;
123 #endif
124 	if (slash != NULL)
125 		return slash + 1;
126 	return path;
127 }
128 
129 void
begin_program(tchar * argv[])130 begin_program(tchar *argv[])
131 {
132 	prog_invocation_name = get_filename(argv[0]);
133 
134 #ifdef FREESTANDING
135 	/* This allows testing freestanding library builds. */
136 	libdeflate_set_memory_allocator(malloc, free);
137 #endif
138 }
139 
140 /* Create a copy of 'path' surrounded by double quotes */
141 static tchar *
quote_path(const tchar * path)142 quote_path(const tchar *path)
143 {
144 	size_t len = tstrlen(path);
145 	tchar *result;
146 
147 	result = xmalloc((1 + len + 1 + 1) * sizeof(tchar));
148 	if (result == NULL)
149 		return NULL;
150 	result[0] = '"';
151 	tmemcpy(&result[1], path, len);
152 	result[1 + len] = '"';
153 	result[1 + len + 1] = '\0';
154 	return result;
155 }
156 
157 /* Open a file for reading, or set up standard input for reading */
158 int
xopen_for_read(const tchar * path,bool symlink_ok,struct file_stream * strm)159 xopen_for_read(const tchar *path, bool symlink_ok, struct file_stream *strm)
160 {
161 	strm->mmap_token = NULL;
162 	strm->mmap_mem = NULL;
163 
164 	if (path == NULL) {
165 		strm->is_standard_stream = true;
166 		strm->name = T("standard input");
167 		strm->fd = STDIN_FILENO;
168 	#ifdef _WIN32
169 		_setmode(strm->fd, O_BINARY);
170 	#endif
171 		return 0;
172 	}
173 
174 	strm->is_standard_stream = false;
175 
176 	strm->name = quote_path(path);
177 	if (strm->name == NULL)
178 		return -1;
179 
180 	strm->fd = topen(path, O_RDONLY | O_BINARY | O_NONBLOCK | O_NOCTTY |
181 			 (symlink_ok ? 0 : O_NOFOLLOW) | O_SEQUENTIAL);
182 	if (strm->fd < 0) {
183 		msg_errno("Can't open %"TS" for reading", strm->name);
184 		free(strm->name);
185 		return -1;
186 	}
187 
188 #if defined(HAVE_POSIX_FADVISE) && (O_SEQUENTIAL == 0)
189 	(void)posix_fadvise(strm->fd, 0, 0, POSIX_FADV_SEQUENTIAL);
190 #endif
191 
192 	return 0;
193 }
194 
195 /* Open a file for writing, or set up standard output for writing */
196 int
xopen_for_write(const tchar * path,bool overwrite,struct file_stream * strm)197 xopen_for_write(const tchar *path, bool overwrite, struct file_stream *strm)
198 {
199 	int ret = -1;
200 
201 	strm->mmap_token = NULL;
202 	strm->mmap_mem = NULL;
203 
204 	if (path == NULL) {
205 		strm->is_standard_stream = true;
206 		strm->name = T("standard output");
207 		strm->fd = STDOUT_FILENO;
208 	#ifdef _WIN32
209 		_setmode(strm->fd, O_BINARY);
210 	#endif
211 		return 0;
212 	}
213 
214 	strm->is_standard_stream = false;
215 
216 	strm->name = quote_path(path);
217 	if (strm->name == NULL)
218 		goto err;
219 retry:
220 	strm->fd = topen(path, O_WRONLY | O_BINARY | O_NOFOLLOW |
221 				O_CREAT | O_EXCL, 0644);
222 	if (strm->fd < 0) {
223 		if (errno != EEXIST) {
224 			msg_errno("Can't open %"TS" for writing", strm->name);
225 			goto err;
226 		}
227 		if (!overwrite) {
228 			if (!isatty(STDERR_FILENO) || !isatty(STDIN_FILENO)) {
229 				msg("%"TS" already exists; use -f to overwrite",
230 				    strm->name);
231 				ret = -2; /* warning only */
232 				goto err;
233 			}
234 			fprintf(stderr, "%"TS": %"TS" already exists; "
235 				"overwrite? (y/n) ",
236 				prog_invocation_name, strm->name);
237 			if (getchar() != 'y') {
238 				msg("Not overwriting.");
239 				goto err;
240 			}
241 		}
242 		if (tunlink(path) != 0) {
243 			msg_errno("Unable to delete %"TS, strm->name);
244 			goto err;
245 		}
246 		goto retry;
247 	}
248 
249 	return 0;
250 
251 err:
252 	free(strm->name);
253 	return ret;
254 }
255 
256 /* Read the full contents of a file into memory */
257 static int
read_full_contents(struct file_stream * strm)258 read_full_contents(struct file_stream *strm)
259 {
260 	size_t filled = 0;
261 	size_t capacity = 4096;
262 	char *buf;
263 	int ret;
264 
265 	buf = xmalloc(capacity);
266 	if (buf == NULL)
267 		return -1;
268 	do {
269 		if (filled == capacity) {
270 			char *newbuf;
271 
272 			if (capacity == SIZE_MAX)
273 				goto oom;
274 			capacity += MIN(SIZE_MAX - capacity, capacity);
275 			newbuf = realloc(buf, capacity);
276 			if (newbuf == NULL)
277 				goto oom;
278 			buf = newbuf;
279 		}
280 		ret = xread(strm, &buf[filled], capacity - filled);
281 		if (ret < 0)
282 			goto err;
283 		filled += ret;
284 	} while (ret != 0);
285 
286 	strm->mmap_mem = buf;
287 	strm->mmap_size = filled;
288 	return 0;
289 
290 err:
291 	free(buf);
292 	return ret;
293 oom:
294 	msg("Out of memory!  %"TS" is too large to be processed by "
295 	    "this program as currently implemented.", strm->name);
296 	ret = -1;
297 	goto err;
298 }
299 
300 /* Map the contents of a file into memory */
301 int
map_file_contents(struct file_stream * strm,u64 size)302 map_file_contents(struct file_stream *strm, u64 size)
303 {
304 	if (size == 0) /* mmap isn't supported on empty files */
305 		return read_full_contents(strm);
306 
307 	if (size > SIZE_MAX) {
308 		msg("%"TS" is too large to be processed by this program",
309 		    strm->name);
310 		return -1;
311 	}
312 #ifdef _WIN32
313 	strm->mmap_token = CreateFileMapping(
314 				(HANDLE)(intptr_t)_get_osfhandle(strm->fd),
315 				NULL, PAGE_READONLY, 0, 0, NULL);
316 	if (strm->mmap_token == NULL) {
317 		DWORD err = GetLastError();
318 		if (err == ERROR_BAD_EXE_FORMAT) /* mmap unsupported */
319 			return read_full_contents(strm);
320 		msg("Unable create file mapping for %"TS": Windows error %u",
321 		    strm->name, (unsigned int)err);
322 		return -1;
323 	}
324 
325 	strm->mmap_mem = MapViewOfFile((HANDLE)strm->mmap_token,
326 				       FILE_MAP_READ, 0, 0, size);
327 	if (strm->mmap_mem == NULL) {
328 		msg("Unable to map %"TS" into memory: Windows error %u",
329 		    strm->name, (unsigned int)GetLastError());
330 		CloseHandle((HANDLE)strm->mmap_token);
331 		return -1;
332 	}
333 #else /* _WIN32 */
334 	strm->mmap_mem = mmap(NULL, size, PROT_READ, MAP_SHARED, strm->fd, 0);
335 	if (strm->mmap_mem == MAP_FAILED) {
336 		strm->mmap_mem = NULL;
337 		if (errno == ENODEV) /* mmap isn't supported on this file */
338 			return read_full_contents(strm);
339 		if (errno == ENOMEM) {
340 			msg("%"TS" is too large to be processed by this "
341 			    "program", strm->name);
342 		} else {
343 			msg_errno("Unable to map %"TS" into memory",
344 				  strm->name);
345 		}
346 		return -1;
347 	}
348 
349 #ifdef HAVE_POSIX_MADVISE
350 	(void)posix_madvise(strm->mmap_mem, size, POSIX_MADV_SEQUENTIAL);
351 #endif
352 	strm->mmap_token = strm; /* anything that's not NULL */
353 
354 #endif /* !_WIN32 */
355 	strm->mmap_size = size;
356 	return 0;
357 }
358 
359 /*
360  * Read from a file, returning the full count to indicate all bytes were read, a
361  * short count (possibly 0) to indicate EOF, or -1 to indicate error.
362  */
363 ssize_t
xread(struct file_stream * strm,void * buf,size_t count)364 xread(struct file_stream *strm, void *buf, size_t count)
365 {
366 	char *p = buf;
367 	size_t orig_count = count;
368 
369 	while (count != 0) {
370 		ssize_t res = read(strm->fd, p, MIN(count, INT_MAX));
371 		if (res == 0)
372 			break;
373 		if (res < 0) {
374 			if (errno == EAGAIN || errno == EINTR)
375 				continue;
376 			msg_errno("Error reading from %"TS, strm->name);
377 			return -1;
378 		}
379 		p += res;
380 		count -= res;
381 	}
382 	return orig_count - count;
383 }
384 
385 /* Write to a file, returning 0 if all bytes were written or -1 on error */
386 int
full_write(struct file_stream * strm,const void * buf,size_t count)387 full_write(struct file_stream *strm, const void *buf, size_t count)
388 {
389 	const char *p = buf;
390 
391 	while (count != 0) {
392 		ssize_t res = write(strm->fd, p, MIN(count, INT_MAX));
393 		if (res <= 0) {
394 			msg_errno("Error writing to %"TS, strm->name);
395 			return -1;
396 		}
397 		p += res;
398 		count -= res;
399 	}
400 	return 0;
401 }
402 
403 /* Close a file, returning 0 on success or -1 on error */
404 int
xclose(struct file_stream * strm)405 xclose(struct file_stream *strm)
406 {
407 	int ret = 0;
408 
409 	if (!strm->is_standard_stream) {
410 		if (close(strm->fd) != 0) {
411 			msg_errno("Error closing %"TS, strm->name);
412 			ret = -1;
413 		}
414 		free(strm->name);
415 	}
416 
417 	if (strm->mmap_token != NULL) {
418 #ifdef _WIN32
419 		UnmapViewOfFile(strm->mmap_mem);
420 		CloseHandle((HANDLE)strm->mmap_token);
421 #else
422 		munmap(strm->mmap_mem, strm->mmap_size);
423 #endif
424 		strm->mmap_token = NULL;
425 	} else {
426 		free(strm->mmap_mem);
427 	}
428 	strm->mmap_mem = NULL;
429 	strm->fd = -1;
430 	strm->name = NULL;
431 	return ret;
432 }
433 
434 /*
435  * Parse the compression level given on the command line, returning the
436  * compression level on success or -1 on error
437  */
438 int
parse_compression_level(tchar opt_char,const tchar * arg)439 parse_compression_level(tchar opt_char, const tchar *arg)
440 {
441 	int level;
442 
443 	if (arg == NULL)
444 		arg = T("");
445 
446 	if (opt_char < '0' || opt_char > '9')
447 		goto invalid;
448 	level = opt_char - '0';
449 
450 	if (arg[0] != '\0') {
451 		if (arg[0] < '0' || arg[0] > '9')
452 			goto invalid;
453 		if (arg[1] != '\0')	/* Levels are at most 2 digits */
454 			goto invalid;
455 		if (level == 0)		/* Don't allow arguments like "-01" */
456 			goto invalid;
457 		level = (level * 10) + (arg[0] - '0');
458 	}
459 
460 	if (level < 0 || level > 12)
461 		goto invalid;
462 
463 	return level;
464 
465 invalid:
466 	msg("Invalid compression level: \"%"TC"%"TS"\".  "
467 	    "Must be an integer in the range [0, 12].", opt_char, arg);
468 	return -1;
469 }
470 
471 /* Allocate a new DEFLATE compressor */
472 struct libdeflate_compressor *
alloc_compressor(int level)473 alloc_compressor(int level)
474 {
475 	struct libdeflate_compressor *c;
476 
477 	c = libdeflate_alloc_compressor(level);
478 	if (c == NULL) {
479 		msg_errno("Unable to allocate compressor with "
480 			  "compression level %d", level);
481 	}
482 	return c;
483 }
484 
485 /* Allocate a new DEFLATE decompressor */
486 struct libdeflate_decompressor *
alloc_decompressor(void)487 alloc_decompressor(void)
488 {
489 	struct libdeflate_decompressor *d;
490 
491 	d = libdeflate_alloc_decompressor();
492 	if (d == NULL)
493 		msg_errno("Unable to allocate decompressor");
494 
495 	return d;
496 }
497