usr.bin/mkimg/image.c

   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         int error;
  98
  99         error = image_copyout_region(fd, 0, image_size);
 100         if (!error)
 101                 error = image_copyout_done(fd);
 102         return (error);
 103 }
 104
 105 int
 106 image_copyout_done(int fd)
 107 {
 108         off_t ofs;
 109         int error;
 110
 111         ofs = lseek(fd, 0L, SEEK_CUR);
 112         if (ofs == -1)
 113                 return (0);
 114         error = (ftruncate(fd, ofs) == -1) ? errno : 0;
 115         return (error);
 116 }
 117
 118 int
 119 image_copyout_region(int fd, lba_t blk, lba_t size)
 120 {
 121         char *buffer;
 122         off_t ofs;
 123         size_t sz;
 124         ssize_t rdsz, wrsz;
 125         int error;
 126
 127         ofs = lseek(fd, 0L, SEEK_CUR);
 128
 129         blk *= secsz;
 130         if (lseek(image_fd, blk, SEEK_SET) != blk)
 131                 return (errno);
 132         buffer = malloc(BUFFER_SIZE);
 133         if (buffer == NULL)
 134                 return (errno);
 135         error = 0;
 136         size *= secsz;
 137         while (size > 0) {
 138                 sz = (BUFFER_SIZE < size) ? BUFFER_SIZE : size;
 139                 rdsz = read(image_fd, buffer, sz);
 140                 if (rdsz <= 0) {
 141                         error = (rdsz < 0) ? errno : 0;
 142                         break;
 143                 }
 144                 wrsz = (ofs == -1) ?
 145                     write(fd, buffer, rdsz) :
 146                     sparse_write(fd, buffer, rdsz);
 147                 if (wrsz < 0) {
 148                         error = errno;
 149                         break;
 150                 }
 151                 assert(wrsz == rdsz);
 152                 size -= rdsz;
 153         }
 154         free(buffer);
 155         return (error);
 156 }
 157
 158 int
 159 image_data(lba_t blk, lba_t size)
 160 {
 161         char *buffer, *p;
 162
 163         blk *= secsz;
 164         if (lseek(image_fd, blk, SEEK_SET) != blk)
 165                 return (1);
 166
 167         size *= secsz;
 168         buffer = malloc(size);
 169         if (buffer == NULL)
 170                 return (1);
 171
 172         if (read(image_fd, buffer, size) != (ssize_t)size) {
 173                 free(buffer);
 174                 return (1);
 175         }
 176
 177         p = buffer;
 178         while (size > 0 && *p == '\0')
 179                 size--, p++;
 180
 181         free(buffer);
 182         return ((size == 0) ? 0 : 1);
 183 }
 184
 185 lba_t
 186 image_get_size(void)
 187 {
 188
 189         return (image_size);
 190 }
 191
 192 int
 193 image_set_size(lba_t blk)
 194 {
 195
 196         image_size = blk;
 197         if (ftruncate(image_fd, blk * secsz) == -1)
 198                 return (errno);
 199         return (0);
 200 }
 201
 202 int
 203 image_write(lba_t blk, void *buf, ssize_t len)
 204 {
 205
 206         blk *= secsz;
 207         if (lseek(image_fd, blk, SEEK_SET) != blk)
 208                 return (errno);
 209         len *= secsz;
 210         if (sparse_write(image_fd, buf, len) != len)
 211                 return (errno);
 212         return (0);
 213 }
 214
 215 int
 216 image_init(void)
 217 {
 218         const char *tmpdir;
 219
 220         if (atexit(cleanup) == -1)
 221                 return (errno);
 222         if ((tmpdir = getenv("TMPDIR")) == NULL || *tmpdir == '\0')
 223                 tmpdir = _PATH_TMP;
 224         snprintf(image_tmpfile, sizeof(image_tmpfile), "%s/mkimg-XXXXXX",
 225             tmpdir);
 226         image_fd = mkstemp(image_tmpfile);
 227         if (image_fd == -1)
 228                 return (errno);
 229         return (0);
 230 }