xref: /freebsd/usr.bin/mkimg/image.c (revision a223d3ed90bfe313ce5987d468a25a915d7d1254)
1 /*-
2  * Copyright (c) 2014 Juniper Networks, Inc.
3  * All rights reserved.
4  *
5  * Redistribution and use in source and binary forms, with or without
6  * modification, are permitted provided that the following conditions
7  * are met:
8  * 1. Redistributions of source code must retain the above copyright
9  *    notice, this list of conditions and the following disclaimer.
10  * 2. Redistributions in binary form must reproduce the above copyright
11  *    notice, this list of conditions and the following disclaimer in the
12  *    documentation and/or other materials provided with the distribution.
13  *
14  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
15  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
16  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
17  * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
18  * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
19  * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
20  * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
21  * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
22  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
23  * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
24  * SUCH DAMAGE.
25  */
26 
27 #include <sys/cdefs.h>
28 __FBSDID("$FreeBSD$");
29 
30 #include <sys/types.h>
31 #include <assert.h>
32 #include <errno.h>
33 #include <limits.h>
34 #include <paths.h>
35 #include <stdio.h>
36 #include <stdlib.h>
37 #include <unistd.h>
38 
39 #include "image.h"
40 #include "mkimg.h"
41 
42 #define	BUFFER_SIZE	(1024*1024)
43 
44 static char image_tmpfile[PATH_MAX];
45 static int image_fd = -1;
46 static lba_t image_size;
47 
48 static void
49 cleanup(void)
50 {
51 
52 	if (image_fd != -1)
53 		close(image_fd);
54 	unlink(image_tmpfile);
55 }
56 
57 int
58 image_copyin(lba_t blk, int fd, uint64_t *sizep)
59 {
60 	char *buffer;
61 	uint64_t bytesize;
62 	ssize_t bcnt, rdsz;
63 	int error, partial;
64 
65 	assert(BUFFER_SIZE % secsz == 0);
66 
67 	buffer = malloc(BUFFER_SIZE);
68 	if (buffer == NULL)
69 		return (ENOMEM);
70 	bytesize = 0;
71 	partial = 0;
72 	while (1) {
73 		rdsz = read(fd, buffer, BUFFER_SIZE);
74 		if (rdsz <= 0) {
75 			error = (rdsz < 0) ? errno : 0;
76 			break;
77 		}
78 		if (partial)
79 			abort();
80 		bytesize += rdsz;
81 		bcnt = (rdsz + secsz - 1) / secsz;
82 		error = image_write(blk, buffer, bcnt);
83 		if (error)
84 			break;
85 		blk += bcnt;
86 		partial = ((ssize_t)(bcnt * secsz) != rdsz) ? 1 : 0;
87 	}
88 	free(buffer);
89 	if (sizep != NULL)
90 		*sizep = bytesize;
91 	return (error);
92 }
93 
94 int
95 image_copyout(int fd)
96 {
97 	char *buffer;
98 	off_t ofs;
99 	ssize_t rdsz, wrsz;
100 	int error;
101 
102 	ofs = lseek(fd, 0L, SEEK_CUR);
103 
104 	if (lseek(image_fd, 0, SEEK_SET) != 0)
105 		return (errno);
106 	buffer = malloc(BUFFER_SIZE);
107 	if (buffer == NULL)
108 		return (errno);
109 	error = 0;
110 	while (1) {
111 		rdsz = read(image_fd, buffer, BUFFER_SIZE);
112 		if (rdsz <= 0) {
113 			error = (rdsz < 0) ? errno : 0;
114 			break;
115 		}
116 		wrsz = (ofs == -1) ?
117 		    write(fd, buffer, rdsz) :
118 		    sparse_write(fd, buffer, rdsz);
119 		if (wrsz < 0) {
120 			error = errno;
121 			break;
122 		}
123 	}
124 	free(buffer);
125 	if (error)
126 		return (error);
127 	ofs = lseek(fd, 0L, SEEK_CUR);
128 	if (ofs == -1)
129 		return (errno);
130 	error = (ftruncate(fd, ofs) == -1) ? errno : 0;
131 	return (error);
132 }
133 
134 lba_t
135 image_get_size(void)
136 {
137 
138 	return (image_size);
139 }
140 
141 int
142 image_set_size(lba_t blk)
143 {
144 
145 	image_size = blk;
146 	if (ftruncate(image_fd, blk * secsz) == -1)
147 		return (errno);
148 	return (0);
149 }
150 
151 int
152 image_write(lba_t blk, void *buf, ssize_t len)
153 {
154 
155 	blk *= secsz;
156 	if (lseek(image_fd, blk, SEEK_SET) != blk)
157 		return (errno);
158 	len *= secsz;
159 	if (sparse_write(image_fd, buf, len) != len)
160 		return (errno);
161 	return (0);
162 }
163 
164 int
165 image_init(void)
166 {
167 	const char *tmpdir;
168 
169 	if (atexit(cleanup) == -1)
170 		return (errno);
171 	if ((tmpdir = getenv("TMPDIR")) == NULL || *tmpdir == '\0')
172 		tmpdir = _PATH_TMP;
173 	snprintf(image_tmpfile, sizeof(image_tmpfile), "%s/mkimg-XXXXXX",
174 	    tmpdir);
175 	image_fd = mkstemp(image_tmpfile);
176 	if (image_fd == -1)
177 		return (errno);
178 	return (0);
179 }
180