usr.bin/mkuzip/mkuzip.c

   1 /*-
   2  * SPDX-License-Identifier: BSD-2-Clause
   3  *
   4  * Copyright (c) 2004-2016 Maxim Sobolev <sobomax@FreeBSD.org>
   5  * All rights reserved.
   6  *
   7  * Redistribution and use in source and binary forms, with or without
   8  * modification, are permitted provided that the following conditions
   9  * are met:
  10  * 1. Redistributions of source code must retain the above copyright
  11  *    notice, this list of conditions and the following disclaimer.
  12  * 2. Redistributions in binary form must reproduce the above copyright
  13  *    notice, this list of conditions and the following disclaimer in the
  14  *    documentation and/or other materials provided with the distribution.
  15  *
  16  * THIS SOFTWARE IS PROVIDED BY THE AUTHOR AND CONTRIBUTORS ``AS IS'' AND
  17  * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  18  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  19  * ARE DISCLAIMED.  IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE
  20  * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
  21  * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
  22  * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
  23  * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
  24  * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
  25  * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
  26  * SUCH DAMAGE.
  27  */
  28
  29 #include <sys/types.h>
  30 #include <sys/endian.h>
  31 #include <sys/param.h>
  32 #include <sys/sysctl.h>
  33 #include <sys/stat.h>
  34 #include <sys/uio.h>
  35 #include <netinet/in.h>
  36 #include <assert.h>
  37 #include <ctype.h>
  38 #include <err.h>
  39 #include <fcntl.h>
  40 #include <pthread.h>
  41 #include <signal.h>
  42 #include <stdint.h>
  43 #include <stdio.h>
  44 #include <stdlib.h>
  45 #include <string.h>
  46 #include <unistd.h>
  47
  48 #include "mkuzip.h"
  49 #include "mkuz_cloop.h"
  50 #include "mkuz_blockcache.h"
  51 #include "mkuz_lzma.h"
  52 #include "mkuz_zlib.h"
  53 #include "mkuz_zstd.h"
  54 #include "mkuz_blk.h"
  55 #include "mkuz_cfg.h"
  56 #include "mkuz_conveyor.h"
  57 #include "mkuz_format.h"
  58 #include "mkuz_fqueue.h"
  59 #include "mkuz_time.h"
  60 #include "mkuz_insize.h"
  61
  62 #define DEFAULT_CLSTSIZE        16384
  63
  64 enum UZ_ALGORITHM {
  65         UZ_ZLIB = 0,
  66         UZ_LZMA,
  67         UZ_ZSTD,
  68         UZ_INVALID
  69 };
  70
  71 static const struct mkuz_format uzip_fmts[] = {
  72         [UZ_ZLIB] = {
  73                 .option = "zlib",
  74                 .magic = CLOOP_MAGIC_ZLIB,
  75                 .default_sufx = DEFAULT_SUFX_ZLIB,
  76                 .f_compress_bound = mkuz_zlib_cbound,
  77                 .f_init = mkuz_zlib_init,
  78                 .f_compress = mkuz_zlib_compress,
  79         },
  80         [UZ_LZMA] = {
  81                 .option = "lzma",
  82                 .magic = CLOOP_MAGIC_LZMA,
  83                 .default_sufx = DEFAULT_SUFX_LZMA,
  84                 .f_compress_bound = mkuz_lzma_cbound,
  85                 .f_init = mkuz_lzma_init,
  86                 .f_compress = mkuz_lzma_compress,
  87         },
  88         [UZ_ZSTD] = {
  89                 .option = "zstd",
  90                 .magic = CLOOP_MAGIC_ZSTD,
  91                 .default_sufx = DEFAULT_SUFX_ZSTD,
  92                 .f_compress_bound = mkuz_zstd_cbound,
  93                 .f_init = mkuz_zstd_init,
  94                 .f_compress = mkuz_zstd_compress,
  95         },
  96 };
  97
  98 static struct mkuz_blk *readblock(int, u_int32_t);
  99 static void usage(void) __dead2;
 100 static void cleanup(void);
 101
 102 static char *cleanfile = NULL;
 103
 104 static int
 105 cmp_blkno(const struct mkuz_blk *bp, void *p)
 106 {
 107         uint32_t *ap;
 108
 109         ap = (uint32_t *)p;
 110
 111         return (bp->info.blkno == *ap);
 112 }
 113
 114 int main(int argc, char **argv)
 115 {
 116         struct mkuz_cfg cfs;
 117         char *oname;
 118         uint64_t *toc;
 119         int i, io, opt, tmp;
 120         struct {
 121                 int en;
 122                 FILE *f;
 123         } summary;
 124         struct iovec iov[2];
 125         uint64_t offset, last_offset;
 126         struct cloop_header hdr;
 127         struct mkuz_conveyor *cvp;
 128         struct mkuz_blk_info *chit;
 129         size_t ncpusz, ncpu, magiclen;
 130         double st, et;
 131         enum UZ_ALGORITHM comp_alg;
 132         int comp_level;
 133
 134         st = getdtime();
 135
 136         ncpusz = sizeof(size_t);
 137         if (sysctlbyname("hw.ncpu", &ncpu, &ncpusz, NULL, 0) < 0) {
 138                 ncpu = 1;
 139         } else if (ncpu > MAX_WORKERS_AUTO) {
 140                 ncpu = MAX_WORKERS_AUTO;
 141         }
 142
 143         memset(&hdr, 0, sizeof(hdr));
 144         cfs.blksz = DEFAULT_CLSTSIZE;
 145         oname = NULL;
 146         cfs.verbose = 0;
 147         cfs.no_zcomp = 0;
 148         cfs.en_dedup = 0;
 149         summary.en = 0;
 150         summary.f = stderr;
 151         comp_alg = UZ_ZLIB;
 152         comp_level = USE_DEFAULT_LEVEL;
 153         cfs.nworkers = ncpu;
 154         struct mkuz_blk *iblk, *oblk;
 155
 156         while((opt = getopt(argc, argv, "A:C:o:s:vZdLSj:")) != -1) {
 157                 switch(opt) {
 158                 case 'A':
 159                         for (tmp = UZ_ZLIB; tmp < UZ_INVALID; tmp++) {
 160                                 if (strcmp(uzip_fmts[tmp].option, optarg) == 0)
 161                                         break;
 162                         }
 163                         if (tmp == UZ_INVALID)
 164                                 errx(1, "invalid algorithm specified: %s",
 165                                     optarg);
 166                                 /* Not reached */
 167                         comp_alg = tmp;
 168                         break;
 169                 case 'C':
 170                         comp_level = atoi(optarg);
 171                         break;
 172                 case 'o':
 173                         oname = optarg;
 174                         break;
 175
 176                 case 's':
 177                         tmp = atoi(optarg);
 178                         if (tmp <= 0) {
 179                                 errx(1, "invalid cluster size specified: %s",
 180                                     optarg);
 181                                 /* Not reached */
 182                         }
 183                         cfs.blksz = tmp;
 184                         break;
 185
 186                 case 'v':
 187                         cfs.verbose = 1;
 188                         break;
 189
 190                 case 'Z':
 191                         cfs.no_zcomp = 1;
 192                         break;
 193
 194                 case 'd':
 195                         cfs.en_dedup = 1;
 196                         break;
 197
 198                 case 'L':
 199                         comp_alg = UZ_LZMA;
 200                         break;
 201
 202                 case 'S':
 203                         summary.en = 1;
 204                         summary.f = stdout;
 205                         break;
 206
 207                 case 'j':
 208                         tmp = atoi(optarg);
 209                         if (tmp <= 0) {
 210                                 errx(1, "invalid number of compression threads"
 211                                     " specified: %s", optarg);
 212                                 /* Not reached */
 213                         }
 214                         cfs.nworkers = tmp;
 215                         break;
 216
 217                 default:
 218                         usage();
 219                         /* Not reached */
 220                 }
 221         }
 222         argc -= optind;
 223         argv += optind;
 224
 225         if (argc != 1) {
 226                 usage();
 227                 /* Not reached */
 228         }
 229
 230         cfs.handler = &uzip_fmts[comp_alg];
 231
 232         magiclen = strlcpy(hdr.magic, cfs.handler->magic, sizeof(hdr.magic));
 233         assert(magiclen < sizeof(hdr.magic));
 234
 235         if (cfs.en_dedup != 0) {
 236                 /*
 237                  * Dedupe requires a version 3 format.  Don't downgrade newer
 238                  * formats.
 239                  */
 240                 if (hdr.magic[CLOOP_OFS_VERSN] == CLOOP_MAJVER_2)
 241                         hdr.magic[CLOOP_OFS_VERSN] = CLOOP_MAJVER_3;
 242                 hdr.magic[CLOOP_OFS_COMPR] =
 243                     tolower(hdr.magic[CLOOP_OFS_COMPR]);
 244         }
 245
 246         if (cfs.blksz % DEV_BSIZE != 0)
 247                 errx(1, "cluster size should be multiple of %d", DEV_BSIZE);
 248
 249         cfs.cbound_blksz = cfs.handler->f_compress_bound(cfs.blksz);
 250         if (cfs.cbound_blksz > MAXPHYS)
 251                 errx(1, "maximal compressed cluster size %zu greater than MAXPHYS %zu",
 252                     cfs.cbound_blksz, (size_t)MAXPHYS);
 253
 254         cfs.handler->f_init(&comp_level);
 255         cfs.comp_level = comp_level;
 256
 257         cfs.iname = argv[0];
 258         if (oname == NULL) {
 259                 asprintf(&oname, "%s%s", cfs.iname, cfs.handler->default_sufx);
 260                 if (oname == NULL) {
 261                         err(1, "can't allocate memory");
 262                         /* Not reached */
 263                 }
 264         }
 265
 266         signal(SIGHUP, exit);
 267         signal(SIGINT, exit);
 268         signal(SIGTERM, exit);
 269         signal(SIGXCPU, exit);
 270         signal(SIGXFSZ, exit);
 271         atexit(cleanup);
 272
 273         cfs.fdr = open(cfs.iname, O_RDONLY);
 274         if (cfs.fdr < 0) {
 275                 err(1, "open(%s)", cfs.iname);
 276                 /* Not reached */
 277         }
 278         cfs.isize = mkuz_get_insize(&cfs);
 279         if (cfs.isize < 0) {
 280                 errx(1, "can't determine input image size");
 281                 /* Not reached */
 282         }
 283         hdr.nblocks = cfs.isize / cfs.blksz;
 284         if ((cfs.isize % cfs.blksz) != 0) {
 285                 if (cfs.verbose != 0)
 286                         fprintf(stderr, "file size is not multiple "
 287                         "of %d, padding data\n", cfs.blksz);
 288                 hdr.nblocks++;
 289         }
 290         toc = mkuz_safe_malloc((hdr.nblocks + 1) * sizeof(*toc));
 291
 292         /*
 293          * Initialize last+1 entry with non-heap trash.  If final padding is
 294          * added later, it may or may not be overwritten with an offset
 295          * representing the length of the final compressed block.  If not,
 296          * initialize to a defined value.
 297          */
 298         toc[hdr.nblocks] = 0;
 299
 300         cfs.fdw = open(oname, (cfs.en_dedup ? O_RDWR : O_WRONLY) | O_TRUNC | O_CREAT,
 301                    S_IRWXU | S_IRGRP | S_IXGRP | S_IROTH | S_IXOTH);
 302         if (cfs.fdw < 0) {
 303                 err(1, "open(%s)", oname);
 304                 /* Not reached */
 305         }
 306         cleanfile = oname;
 307
 308         /* Prepare header that we will write later when we have index ready. */
 309         iov[0].iov_base = (char *)&hdr;
 310         iov[0].iov_len = sizeof(hdr);
 311         iov[1].iov_base = (char *)toc;
 312         iov[1].iov_len = (hdr.nblocks + 1) * sizeof(*toc);
 313         offset = iov[0].iov_len + iov[1].iov_len;
 314
 315         /* Reserve space for header */
 316         lseek(cfs.fdw, offset, SEEK_SET);
 317
 318         if (cfs.verbose != 0) {
 319                 fprintf(stderr, "data size %ju bytes, number of clusters "
 320                     "%u, index length %zu bytes\n", cfs.isize,
 321                     hdr.nblocks, iov[1].iov_len);
 322         }
 323
 324         cvp = mkuz_conveyor_ctor(&cfs);
 325
 326         last_offset = 0;
 327         iblk = oblk = NULL;
 328         for(i = io = 0; iblk != MKUZ_BLK_EOF; i++) {
 329                 iblk = readblock(cfs.fdr, cfs.blksz);
 330                 mkuz_fqueue_enq(cvp->wrk_queue, iblk);
 331                 if (iblk != MKUZ_BLK_EOF &&
 332                     (i < (cfs.nworkers * ITEMS_PER_WORKER))) {
 333                         continue;
 334                 }
 335 drain:
 336                 oblk = mkuz_fqueue_deq_when(cvp->results, cmp_blkno, &io);
 337                 assert(oblk->info.blkno == (unsigned)io);
 338                 oblk->info.offset = offset;
 339                 chit = NULL;
 340                 if (cfs.en_dedup != 0 && oblk->info.len > 0) {
 341                         chit = mkuz_blkcache_regblock(cfs.fdw, oblk);
 342                         /*
 343                          * There should be at least one non-empty block
 344                          * between us and the backref'ed offset, otherwise
 345                          * we won't be able to parse that sequence correctly
 346                          * as it would be indistinguishible from another
 347                          * empty block.
 348                          */
 349                         if (chit != NULL && chit->offset == last_offset) {
 350                                 chit = NULL;
 351                         }
 352                 }
 353                 if (chit != NULL) {
 354                         toc[io] = htobe64(chit->offset);
 355                         oblk->info.len = 0;
 356                 } else {
 357                         if (oblk->info.len > 0 && write(cfs.fdw, oblk->data,
 358                             oblk->info.len) < 0) {
 359                                 err(1, "write(%s)", oname);
 360                                 /* Not reached */
 361                         }
 362                         toc[io] = htobe64(offset);
 363                         last_offset = offset;
 364                         offset += oblk->info.len;
 365                 }
 366                 if (cfs.verbose != 0) {
 367                         fprintf(stderr, "cluster #%d, in %u bytes, "
 368                             "out len=%lu offset=%lu", io, cfs.blksz,
 369                             (u_long)oblk->info.len, (u_long)be64toh(toc[io]));
 370                         if (chit != NULL) {
 371                                 fprintf(stderr, " (backref'ed to #%d)",
 372                                     chit->blkno);
 373                         }
 374                         fprintf(stderr, "\n");
 375                 }
 376                 free(oblk);
 377                 io += 1;
 378                 if (iblk == MKUZ_BLK_EOF) {
 379                         if (io < i)
 380                                 goto drain;
 381                         /* Last block, see if we need to add some padding */
 382                         if ((offset % DEV_BSIZE) == 0)
 383                                 continue;
 384                         oblk = mkuz_blk_ctor(DEV_BSIZE - (offset % DEV_BSIZE));
 385                         oblk->info.blkno = io;
 386                         oblk->info.len = oblk->alen;
 387                         if (cfs.verbose != 0) {
 388                                 fprintf(stderr, "padding data with %lu bytes "
 389                                     "so that file size is multiple of %d\n",
 390                                     (u_long)oblk->alen, DEV_BSIZE);
 391                         }
 392                         mkuz_fqueue_enq(cvp->results, oblk);
 393                         goto drain;
 394                 }
 395         }
 396
 397         close(cfs.fdr);
 398
 399         if (cfs.verbose != 0 || summary.en != 0) {
 400                 et = getdtime();
 401                 fprintf(summary.f, "compressed data to %ju bytes, saved %lld "
 402                     "bytes, %.2f%% decrease, %.2f bytes/sec.\n", offset,
 403                     (long long)(cfs.isize - offset),
 404                     100.0 * (long long)(cfs.isize - offset) /
 405                     (float)cfs.isize, (float)cfs.isize / (et - st));
 406         }
 407
 408         /* Convert to big endian */
 409         hdr.blksz = htonl(cfs.blksz);
 410         hdr.nblocks = htonl(hdr.nblocks);
 411         /* Write headers into pre-allocated space */
 412         lseek(cfs.fdw, 0, SEEK_SET);
 413         if (writev(cfs.fdw, iov, 2) < 0) {
 414                 err(1, "writev(%s)", oname);
 415                 /* Not reached */
 416         }
 417         cleanfile = NULL;
 418         close(cfs.fdw);
 419
 420         exit(0);
 421 }
 422
 423 static struct mkuz_blk *
 424 readblock(int fd, u_int32_t clstsize)
 425 {
 426         int numread;
 427         struct mkuz_blk *rval;
 428         static int blockcnt;
 429         off_t cpos;
 430
 431         rval = mkuz_blk_ctor(clstsize);
 432
 433         rval->info.blkno = blockcnt;
 434         blockcnt += 1;
 435         cpos = lseek(fd, 0, SEEK_CUR);
 436         if (cpos < 0) {
 437                 err(1, "readblock: lseek() failed");
 438                 /* Not reached */
 439         }
 440         rval->info.offset = cpos;
 441
 442         numread = read(fd, rval->data, clstsize);
 443         if (numread < 0) {
 444                 err(1, "readblock: read() failed");
 445                 /* Not reached */
 446         }
 447         if (numread == 0) {
 448                 free(rval);
 449                 return MKUZ_BLK_EOF;
 450         }
 451         rval->info.len = numread;
 452         return rval;
 453 }
 454
 455 static void
 456 usage(void)
 457 {
 458
 459         fprintf(stderr, "usage: mkuzip [-vZdLS] [-o outfile] [-s cluster_size] "
 460             "[-j ncompr] infile\n");
 461         exit(1);
 462 }
 463
 464 void *
 465 mkuz_safe_malloc(size_t size)
 466 {
 467         void *retval;
 468
 469         retval = malloc(size);
 470         if (retval == NULL) {
 471                 err(1, "can't allocate memory");
 472                 /* Not reached */
 473         }
 474         return retval;
 475 }
 476
 477 void *
 478 mkuz_safe_zmalloc(size_t size)
 479 {
 480         void *retval;
 481
 482         retval = mkuz_safe_malloc(size);
 483         bzero(retval, size);
 484         return retval;
 485 }
 486
 487 static void
 488 cleanup(void)
 489 {
 490
 491         if (cleanfile != NULL)
 492                 unlink(cleanfile);
 493 }
 494
 495 int
 496 mkuz_memvcmp(const void *memory, unsigned char val, size_t size)
 497 {
 498     const u_char *mm;
 499
 500     mm = (const u_char *)memory;
 501     return (*mm == val) && memcmp(mm, mm + 1, size - 1) == 0;
 502 }