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/mman.h>
  31 #include <sys/queue.h>
  32 #include <sys/stat.h>
  33 #include <sys/types.h>
  34 #include <assert.h>
  35 #include <errno.h>
  36 #include <limits.h>
  37 #include <paths.h>
  38 #include <stdint.h>
  39 #include <stdio.h>
  40 #include <stdlib.h>
  41 #include <string.h>
  42 #include <unistd.h>
  43
  44 #include "image.h"
  45 #include "mkimg.h"
  46
  47 struct chunk {
  48         STAILQ_ENTRY(chunk) ch_list;
  49         size_t  ch_size;                /* Size of chunk in bytes. */
  50         lba_t   ch_block;               /* Block address in image. */
  51         union {
  52                 struct {
  53                         off_t   ofs;    /* Offset in backing file. */
  54                         int     fd;     /* FD of backing file. */
  55                 } file;
  56                 struct {
  57                         void    *ptr;   /* Pointer to data in memory */
  58                 } mem;
  59         } ch_u;
  60         u_int   ch_type;
  61 #define CH_TYPE_ZEROES          0       /* Chunk is a gap (no data). */
  62 #define CH_TYPE_FILE            1       /* File-backed chunk. */
  63 #define CH_TYPE_MEMORY          2       /* Memory-backed chunk */
  64 };
  65
  66 static STAILQ_HEAD(chunk_head, chunk) image_chunks;
  67 static u_int image_nchunks;
  68
  69 static char image_swap_file[PATH_MAX];
  70 static int image_swap_fd = -1;
  71 static u_int image_swap_pgsz;
  72 static off_t image_swap_size;
  73
  74 static lba_t image_size;
  75
  76 static int
  77 is_empty_sector(void *buf)
  78 {
  79         uint64_t *p = buf;
  80         size_t n, max;
  81
  82         assert(((uintptr_t)p & 3) == 0);
  83
  84         max = secsz / sizeof(uint64_t);
  85         for (n = 0; n < max; n++) {
  86                 if (p[n] != 0UL)
  87                         return (0);
  88         }
  89         return (1);
  90 }
  91
  92 /*
  93  * Swap file handlng.
  94  */
  95
  96 static off_t
  97 image_swap_alloc(size_t size)
  98 {
  99         off_t ofs;
 100         size_t unit;
 101
 102         unit = (secsz > image_swap_pgsz) ? secsz : image_swap_pgsz;
 103         assert((unit & (unit - 1)) == 0);
 104
 105         size = (size + unit - 1) & ~(unit - 1);
 106
 107         ofs = image_swap_size;
 108         image_swap_size += size;
 109         if (ftruncate(image_swap_fd, image_swap_size) == -1) {
 110                 image_swap_size = ofs;
 111                 ofs = -1LL;
 112         }
 113         return (ofs);
 114 }
 115
 116 /*
 117  * Image chunk handling.
 118  */
 119
 120 static struct chunk *
 121 image_chunk_find(lba_t blk)
 122 {
 123         static struct chunk *last = NULL;
 124         struct chunk *ch;
 125
 126         ch = (last != NULL && last->ch_block <= blk)
 127             ? last : STAILQ_FIRST(&image_chunks);
 128         while (ch != NULL) {
 129                 if (ch->ch_block <= blk &&
 130                     (lba_t)(ch->ch_block + (ch->ch_size / secsz)) > blk) {
 131                         last = ch;
 132                         break;
 133                 }
 134                 ch = STAILQ_NEXT(ch, ch_list);
 135         }
 136         return (ch);
 137 }
 138
 139 static size_t
 140 image_chunk_grow(struct chunk *ch, size_t sz)
 141 {
 142         size_t dsz, newsz;
 143
 144         newsz = ch->ch_size + sz;
 145         if (newsz > ch->ch_size) {
 146                 ch->ch_size = newsz;
 147                 return (0);
 148         }
 149         /* We would overflow -- create new chunk for remainder. */
 150         dsz = SIZE_MAX - ch->ch_size;
 151         assert(dsz < sz);
 152         ch->ch_size = SIZE_MAX;
 153         return (sz - dsz);
 154 }
 155
 156 static struct chunk *
 157 image_chunk_memory(struct chunk *ch, lba_t blk)
 158 {
 159         struct chunk *new;
 160         void *ptr;
 161
 162         ptr = calloc(1, secsz);
 163         if (ptr == NULL)
 164                 return (NULL);
 165
 166         if (ch->ch_block < blk) {
 167                 new = malloc(sizeof(*new));
 168                 if (new == NULL) {
 169                         free(ptr);
 170                         return (NULL);
 171                 }
 172                 memcpy(new, ch, sizeof(*new));
 173                 ch->ch_size = (blk - ch->ch_block) * secsz;
 174                 new->ch_block = blk;
 175                 new->ch_size -= ch->ch_size;
 176                 STAILQ_INSERT_AFTER(&image_chunks, ch, new, ch_list);
 177                 image_nchunks++;
 178                 ch = new;
 179         }
 180
 181         if (ch->ch_size > secsz) {
 182                 new = malloc(sizeof(*new));
 183                 if (new == NULL) {
 184                         free(ptr);
 185                         return (NULL);
 186                 }
 187                 memcpy(new, ch, sizeof(*new));
 188                 ch->ch_size = secsz;
 189                 new->ch_block++;
 190                 new->ch_size -= secsz;
 191                 STAILQ_INSERT_AFTER(&image_chunks, ch, new, ch_list);
 192                 image_nchunks++;
 193         }
 194
 195         ch->ch_type = CH_TYPE_MEMORY;
 196         ch->ch_u.mem.ptr = ptr;
 197         return (ch);
 198 }
 199
 200 static int
 201 image_chunk_skipto(lba_t to)
 202 {
 203         struct chunk *ch;
 204         lba_t from;
 205         size_t sz;
 206
 207         ch = STAILQ_LAST(&image_chunks, chunk, ch_list);
 208         from = (ch != NULL) ? ch->ch_block + (ch->ch_size / secsz) : 0LL;
 209
 210         assert(from <= to);
 211
 212         /* Nothing to do? */
 213         if (from == to)
 214                 return (0);
 215         /* Avoid bugs due to overflows. */
 216         if ((uintmax_t)(to - from) > (uintmax_t)(SIZE_MAX / secsz))
 217                 return (EFBIG);
 218         sz = (to - from) * secsz;
 219         if (ch != NULL && ch->ch_type == CH_TYPE_ZEROES) {
 220                 sz = image_chunk_grow(ch, sz);
 221                 if (sz == 0)
 222                         return (0);
 223                 from = ch->ch_block + (ch->ch_size / secsz);
 224         }
 225         ch = malloc(sizeof(*ch));
 226         if (ch == NULL)
 227                 return (ENOMEM);
 228         memset(ch, 0, sizeof(*ch));
 229         ch->ch_block = from;
 230         ch->ch_size = sz;
 231         ch->ch_type = CH_TYPE_ZEROES;
 232         STAILQ_INSERT_TAIL(&image_chunks, ch, ch_list);
 233         image_nchunks++;
 234         return (0);
 235 }
 236
 237 static int
 238 image_chunk_append(lba_t blk, size_t sz, off_t ofs, int fd)
 239 {
 240         struct chunk *ch;
 241
 242         ch = STAILQ_LAST(&image_chunks, chunk, ch_list);
 243         if (ch != NULL && ch->ch_type == CH_TYPE_FILE) {
 244                 if (fd == ch->ch_u.file.fd &&
 245                     blk == (lba_t)(ch->ch_block + (ch->ch_size / secsz)) &&
 246                     ofs == (off_t)(ch->ch_u.file.ofs + ch->ch_size)) {
 247                         sz = image_chunk_grow(ch, sz);
 248                         if (sz == 0)
 249                                 return (0);
 250                         blk = ch->ch_block + (ch->ch_size / secsz);
 251                         ofs = ch->ch_u.file.ofs + ch->ch_size;
 252                 }
 253         }
 254         ch = malloc(sizeof(*ch));
 255         if (ch == NULL)
 256                 return (ENOMEM);
 257         memset(ch, 0, sizeof(*ch));
 258         ch->ch_block = blk;
 259         ch->ch_size = sz;
 260         ch->ch_type = CH_TYPE_FILE;
 261         ch->ch_u.file.ofs = ofs;
 262         ch->ch_u.file.fd = fd;
 263         STAILQ_INSERT_TAIL(&image_chunks, ch, ch_list);
 264         image_nchunks++;
 265         return (0);
 266 }
 267
 268 static int
 269 image_chunk_copyin(lba_t blk, void *buf, size_t sz, off_t ofs, int fd)
 270 {
 271         uint8_t *p = buf;
 272         int error;
 273
 274         error = 0;
 275         sz = (sz + secsz - 1) & ~(secsz - 1);
 276         while (!error && sz > 0) {
 277                 if (is_empty_sector(p))
 278                         error = image_chunk_skipto(blk + 1);
 279                 else
 280                         error = image_chunk_append(blk, secsz, ofs, fd);
 281                 blk++;
 282                 p += secsz;
 283                 sz -= secsz;
 284                 ofs += secsz;
 285         }
 286         return (error);
 287 }
 288
 289 /*
 290  * File mapping support.
 291  */
 292
 293 static void *
 294 image_file_map(int fd, off_t ofs, size_t sz)
 295 {
 296         void *ptr;
 297         size_t unit;
 298         int flags, prot;
 299
 300         unit = (secsz > image_swap_pgsz) ? secsz : image_swap_pgsz;
 301         assert((unit & (unit - 1)) == 0);
 302
 303         flags = MAP_NOCORE | MAP_NOSYNC | MAP_SHARED;
 304         /* Allow writing to our swap file only. */
 305         prot = PROT_READ | ((fd == image_swap_fd) ? PROT_WRITE : 0);
 306         sz = (sz + unit - 1) & ~(unit - 1);
 307         ptr = mmap(NULL, sz, prot, flags, fd, ofs);
 308         return ((ptr == MAP_FAILED) ? NULL : ptr);
 309 }
 310
 311 static int
 312 image_file_unmap(void *buffer, size_t sz)
 313 {
 314         size_t unit;
 315
 316         unit = (secsz > image_swap_pgsz) ? secsz : image_swap_pgsz;
 317         sz = (sz + unit - 1) & ~(unit - 1);
 318         munmap(buffer, sz);
 319         return (0);
 320 }
 321
 322 /*
 323  * Input/source file handling.
 324  */
 325
 326 static int
 327 image_copyin_stream(lba_t blk, int fd, uint64_t *sizep)
 328 {
 329         char *buffer;
 330         uint64_t bytesize;
 331         off_t swofs;
 332         size_t iosz;
 333         ssize_t rdsz;
 334         int error;
 335
 336         /*
 337          * This makes sure we're doing I/O in multiples of the page
 338          * size as well as of the sector size. 2MB is the minimum
 339          * by virtue of secsz at least 512 bytes and the page size
 340          * at least 4K bytes.
 341          */
 342         iosz = secsz * image_swap_pgsz;
 343
 344         bytesize = 0;
 345         do {
 346                 swofs = image_swap_alloc(iosz);
 347                 if (swofs == -1LL)
 348                         return (errno);
 349                 buffer = image_file_map(image_swap_fd, swofs, iosz);
 350                 if (buffer == NULL)
 351                         return (errno);
 352                 rdsz = read(fd, buffer, iosz);
 353                 if (rdsz > 0)
 354                         error = image_chunk_copyin(blk, buffer, rdsz, swofs,
 355                             image_swap_fd);
 356                 else if (rdsz < 0)
 357                         error = errno;
 358                 else
 359                         error = 0;
 360                 image_file_unmap(buffer, iosz);
 361                 /* XXX should we relinguish unused swap space? */
 362                 if (error)
 363                         return (error);
 364
 365                 bytesize += rdsz;
 366                 blk += (rdsz + secsz - 1) / secsz;
 367         } while (rdsz > 0);
 368
 369         if (sizep != NULL)
 370                 *sizep = bytesize;
 371         return (0);
 372 }
 373
 374 static int
 375 image_copyin_mapped(lba_t blk, int fd, uint64_t *sizep)
 376 {
 377         off_t cur, data, end, hole, pos;
 378         void *buf;
 379         uint64_t bytesize;
 380         size_t iosz, sz;
 381         int error;
 382
 383         /*
 384          * We'd like to know the size of the file and we must
 385          * be able to seek in order to mmap(2). If this isn't
 386          * possible, then treat the file as a stream/pipe.
 387          */
 388         end = lseek(fd, 0L, SEEK_END);
 389         if (end == -1L)
 390                 return (image_copyin_stream(blk, fd, sizep));
 391
 392         /*
 393          * We need the file opened for the duration and our
 394          * caller is going to close the file. Make a dup(2)
 395          * so that control the faith of the descriptor.
 396          */
 397         fd = dup(fd);
 398         if (fd == -1)
 399                 return (errno);
 400
 401         iosz = secsz * image_swap_pgsz;
 402
 403         bytesize = 0;
 404         cur = pos = 0;
 405         error = 0;
 406         while (!error && cur < end) {
 407                 hole = lseek(fd, cur, SEEK_HOLE);
 408                 if (hole == -1)
 409                         hole = end;
 410                 data = lseek(fd, cur, SEEK_DATA);
 411                 if (data == -1)
 412                         data = end;
 413
 414                 /*
 415                  * Treat the entire file as data if sparse files
 416                  * are not supported by the underlying file system.
 417                  */
 418                 if (hole == end && data == end)
 419                         data = cur;
 420
 421                 if (cur == hole && data > hole) {
 422                         hole = pos;
 423                         pos = data & ~((uint64_t)secsz - 1);
 424
 425                         blk += (pos - hole) / secsz;
 426                         error = image_chunk_skipto(blk);
 427
 428                         bytesize += pos - hole;
 429                         cur = data;
 430                 } else if (cur == data && hole > data) {
 431                         data = pos;
 432                         pos = (hole + secsz - 1) & ~((uint64_t)secsz - 1);
 433
 434                         while (data < pos) {
 435                                 sz = (pos - data > (off_t)iosz)
 436                                     ? iosz : (size_t)(pos - data);
 437
 438                                 buf = image_file_map(fd, data, sz);
 439                                 if (buf != NULL) {
 440                                         error = image_chunk_copyin(blk, buf,
 441                                             sz, data, fd);
 442                                         image_file_unmap(buf, sz);
 443                                 } else
 444                                         error = errno;
 445
 446                                 blk += sz / secsz;
 447                                 bytesize += sz;
 448                                 data += sz;
 449                         }
 450                         cur = hole;
 451                 } else {
 452                         /*
 453                          * I don't know what this means or whether it
 454                          * can happen at all...
 455                          */
 456                         error = EDOOFUS;
 457                         break;
 458                 }
 459         }
 460         if (error)
 461                 close(fd);
 462         if (!error && sizep != NULL)
 463                 *sizep = bytesize;
 464         return (error);
 465 }
 466
 467 int
 468 image_copyin(lba_t blk, int fd, uint64_t *sizep)
 469 {
 470         struct stat sb;
 471         int error;
 472
 473         error = image_chunk_skipto(blk);
 474         if (!error) {
 475                 if (fstat(fd, &sb) == -1 || !S_ISREG(sb.st_mode))
 476                         error = image_copyin_stream(blk, fd, sizep);
 477                 else
 478                         error = image_copyin_mapped(blk, fd, sizep);
 479         }
 480         return (error);
 481 }
 482
 483 /*
 484  * Output/sink file handling.
 485  */
 486
 487 int
 488 image_copyout(int fd)
 489 {
 490         int error;
 491
 492         error = image_copyout_region(fd, 0, image_size);
 493         if (!error)
 494                 error = image_copyout_done(fd);
 495         return (error);
 496 }
 497
 498 int
 499 image_copyout_done(int fd)
 500 {
 501         off_t ofs;
 502         int error;
 503
 504         ofs = lseek(fd, 0L, SEEK_CUR);
 505         if (ofs == -1)
 506                 return (0);
 507         error = (ftruncate(fd, ofs) == -1) ? errno : 0;
 508         return (error);
 509 }
 510
 511 static int
 512 image_copyout_memory(int fd, size_t size, void *ptr)
 513 {
 514
 515         if (write(fd, ptr, size) == -1)
 516                 return (errno);
 517         return (0);
 518 }
 519
 520 static int
 521 image_copyout_zeroes(int fd, size_t size)
 522 {
 523         static uint8_t *zeroes = NULL;
 524         size_t sz;
 525         int error;
 526
 527         if (lseek(fd, (off_t)size, SEEK_CUR) != -1)
 528                 return (0);
 529
 530         /*
 531          * If we can't seek, we must write.
 532          */
 533
 534         if (zeroes == NULL) {
 535                 zeroes = calloc(1, secsz);
 536                 if (zeroes == NULL)
 537                         return (ENOMEM);
 538         }
 539
 540         while (size > 0) {
 541                 sz = (size > secsz) ? secsz : size;
 542                 error = image_copyout_memory(fd, sz, zeroes);
 543                 if (error)
 544                         return (error);
 545                 size -= sz;
 546         }
 547         return (0);
 548 }
 549
 550 static int
 551 image_copyout_file(int fd, size_t size, int ifd, off_t iofs)
 552 {
 553         void *buf;
 554         size_t iosz, sz;
 555         int error;
 556
 557         iosz = secsz * image_swap_pgsz;
 558
 559         while (size > 0) {
 560                 sz = (size > iosz) ? iosz : size;
 561                 buf = image_file_map(ifd, iofs, sz);
 562                 if (buf == NULL)
 563                         return (errno);
 564                 error = image_copyout_memory(fd, sz, buf);
 565                 image_file_unmap(buf, sz);
 566                 if (error)
 567                         return (error);
 568                 size -= sz;
 569                 iofs += sz;
 570         }
 571         return (0);
 572 }
 573
 574 int
 575 image_copyout_region(int fd, lba_t blk, lba_t size)
 576 {
 577         struct chunk *ch;
 578         size_t ofs, sz;
 579         int error;
 580
 581         size *= secsz;
 582
 583         while (size > 0) {
 584                 ch = image_chunk_find(blk);
 585                 if (ch == NULL)
 586                         return (EINVAL);
 587                 ofs = (blk - ch->ch_block) * secsz;
 588                 sz = ch->ch_size - ofs;
 589                 sz = ((lba_t)sz < size) ? sz : (size_t)size;
 590                 switch (ch->ch_type) {
 591                 case CH_TYPE_ZEROES:
 592                         error = image_copyout_zeroes(fd, sz);
 593                         break;
 594                 case CH_TYPE_FILE:
 595                         error = image_copyout_file(fd, sz, ch->ch_u.file.fd,
 596                             ch->ch_u.file.ofs + ofs);
 597                         break;
 598                 case CH_TYPE_MEMORY:
 599                         error = image_copyout_memory(fd, sz, ch->ch_u.mem.ptr);
 600                         break;
 601                 default:
 602                         return (EDOOFUS);
 603                 }
 604                 size -= sz;
 605                 blk += sz / secsz;
 606         }
 607         return (0);
 608 }
 609
 610 int
 611 image_data(lba_t blk, lba_t size)
 612 {
 613         struct chunk *ch;
 614         lba_t lim;
 615
 616         while (1) {
 617                 ch = image_chunk_find(blk);
 618                 if (ch == NULL)
 619                         return (0);
 620                 if (ch->ch_type != CH_TYPE_ZEROES)
 621                         return (1);
 622                 lim = ch->ch_block + (ch->ch_size / secsz);
 623                 if (lim >= blk + size)
 624                         return (0);
 625                 size -= lim - blk;
 626                 blk = lim;
 627         }
 628         /*NOTREACHED*/
 629 }
 630
 631 lba_t
 632 image_get_size(void)
 633 {
 634
 635         return (image_size);
 636 }
 637
 638 int
 639 image_set_size(lba_t blk)
 640 {
 641         int error;
 642
 643         error = image_chunk_skipto(blk);
 644         if (!error)
 645                 image_size = blk;
 646         return (error);
 647 }
 648
 649 int
 650 image_write(lba_t blk, void *buf, ssize_t len)
 651 {
 652         struct chunk *ch;
 653
 654         while (len > 0) {
 655                 if (!is_empty_sector(buf)) {
 656                         ch = image_chunk_find(blk);
 657                         if (ch == NULL)
 658                                 return (ENXIO);
 659                         /* We may not be able to write to files. */
 660                         if (ch->ch_type == CH_TYPE_FILE)
 661                                 return (EINVAL);
 662                         if (ch->ch_type == CH_TYPE_ZEROES) {
 663                                 ch = image_chunk_memory(ch, blk);
 664                                 if (ch == NULL)
 665                                         return (ENOMEM);
 666                         }
 667                         assert(ch->ch_type == CH_TYPE_MEMORY);
 668                         memcpy(ch->ch_u.mem.ptr, buf, secsz);
 669                 }
 670                 blk++;
 671                 buf = (char *)buf + secsz;
 672                 len--;
 673         }
 674         return (0);
 675 }
 676
 677 static void
 678 image_cleanup(void)
 679 {
 680         struct chunk *ch;
 681
 682         while ((ch = STAILQ_FIRST(&image_chunks)) != NULL) {
 683                 switch (ch->ch_type) {
 684                 case CH_TYPE_FILE:
 685                         /* We may be closing the same file multiple times. */
 686                         if (ch->ch_u.file.fd != -1)
 687                                 close(ch->ch_u.file.fd);
 688                         break;
 689                 case CH_TYPE_MEMORY:
 690                         free(ch->ch_u.mem.ptr);
 691                         break;
 692                 default:
 693                         break;
 694                 }
 695                 STAILQ_REMOVE_HEAD(&image_chunks, ch_list);
 696                 free(ch);
 697         }
 698         if (image_swap_fd != -1)
 699                 close(image_swap_fd);
 700         unlink(image_swap_file);
 701 }
 702
 703 int
 704 image_init(void)
 705 {
 706         const char *tmpdir;
 707
 708         STAILQ_INIT(&image_chunks);
 709         image_nchunks = 0;
 710
 711         image_swap_size = 0;
 712         image_swap_pgsz = getpagesize();
 713
 714         if (atexit(image_cleanup) == -1)
 715                 return (errno);
 716         if ((tmpdir = getenv("TMPDIR")) == NULL || *tmpdir == '\0')
 717                 tmpdir = _PATH_TMP;
 718         snprintf(image_swap_file, sizeof(image_swap_file), "%s/mkimg-XXXXXX",
 719             tmpdir);
 720         image_swap_fd = mkstemp(image_swap_file);
 721         if (image_swap_fd == -1)
 722                 return (errno);
 723         return (0);
 724 }