zircon/system/uapp/lz4/main.c

// Copyright 2016 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <errno.h>
#include <fcntl.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include <lz4/lz4frame.h>

#define BLOCK_SIZE 65536

#define WR_NEWFILE O_WRONLY | O_CREAT | O_TRUNC
#define PERM_644 S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH

static void usage(const char* arg0) {
  printf("usage: %s [-1|-9] [-d] <input file> <output file>\n", arg0);
  printf("   -1  fast compression (default)\n");
  printf("   -9  high compression (slower)\n");
  printf("   -d  decompress\n");
}

static int do_decompress(const char* infile, const char* outfile) {
  int infd, outfd;

  infd = open(infile, O_RDONLY);
  if (infd < 0) {
    fprintf(stderr, "could not open %s: %s\n", infile, strerror(errno));
    return -1;
  }

  outfd = open(outfile, WR_NEWFILE, PERM_644);
  if (outfd < 0) {
    fprintf(stderr, "could not open %s: %s\n", outfile, strerror(errno));
    close(infd);
    return -1;
  }

  LZ4F_decompressionContext_t dctx;
  LZ4F_errorCode_t errc = LZ4F_createDecompressionContext(&dctx, LZ4F_VERSION);
  if (LZ4F_isError(errc)) {
    fprintf(stderr, "could not initialize decompression: %s\n", LZ4F_getErrorName(errc));
    close(outfd);
    close(infd);
    return -1;
  }

  uint8_t inbuf[BLOCK_SIZE];
  uint8_t outbuf[BLOCK_SIZE];

  // Read first 4 bytes to let LZ4 tell us how much it expects in the first
  // pass.
  size_t src_sz = 4;
  size_t dst_sz = 0;
  ssize_t nr = read(infd, inbuf, src_sz);
  if (nr < (ssize_t)src_sz) {
    fprintf(stderr, "could not read from %s", infile);
    if (nr < 0) {
      fprintf(stderr, ": %s", strerror(errno));
    }
    fprintf(stderr, "\n");
    goto done;
  }
  size_t to_read = LZ4F_decompress(dctx, outbuf, &dst_sz, inbuf, &src_sz, NULL);
  if (LZ4F_isError(to_read)) {
    fprintf(stderr, "could not decompress %s: %s\n", infile, LZ4F_getErrorName(to_read));
    goto done;
  }
  if (to_read > BLOCK_SIZE) {
    to_read = BLOCK_SIZE;
  }

  while ((nr = read(infd, inbuf, to_read)) > 0) {
    src_sz = nr;
    ssize_t pos = 0;
    size_t next = 0;

    while (pos < nr) {
      dst_sz = BLOCK_SIZE;
      next = LZ4F_decompress(dctx, outbuf, &dst_sz, inbuf + pos, &src_sz, NULL);
      if (LZ4F_isError(next)) {
        fprintf(stderr, "could not decompress %s: %s\n", infile, LZ4F_getErrorName(to_read));
        goto done;
      }

      if (dst_sz) {
        ssize_t nw = write(outfd, outbuf, dst_sz);
        if (nw != (ssize_t)dst_sz) {
          fprintf(stderr, "could not write to %s", outfile);
          if (nw < 0) {
            fprintf(stderr, ": %s", strerror(errno));
          }
          fprintf(stderr, "\n");
          goto done;
        }
      }
      pos += src_sz;
      src_sz = nr - pos;
    }

    to_read = next;
    if (to_read > BLOCK_SIZE || to_read == 0) {
      to_read = BLOCK_SIZE;
    }
  }

  if (nr < 0) {
    fprintf(stderr, "error reading %s: %s\n", infile, strerror(errno));
    goto done;
  }

done:
  LZ4F_freeDecompressionContext(dctx);
  close(outfd);
  close(infd);
  return 0;
}

static int do_compress(const char* infile, const char* outfile, int clevel) {
  int infd, outfd;

  infd = open(infile, O_RDONLY);
  if (infd < 0) {
    fprintf(stderr, "could not open %s: %s\n", infile, strerror(errno));
    return -1;
  }

  outfd = open(outfile, WR_NEWFILE, PERM_644);
  if (outfd < 0) {
    fprintf(stderr, "could not open %s: %s\n", outfile, strerror(errno));
    close(infd);
    return -1;
  }

  LZ4F_preferences_t prefs;
  memset(&prefs, 0, sizeof(prefs));
  prefs.compressionLevel = clevel;

  uint8_t inbuf[BLOCK_SIZE];
  size_t outsize = LZ4F_compressFrameBound(BLOCK_SIZE, &prefs);
  uint8_t* outbuf = malloc(outsize);
  if (!outbuf) {
    fprintf(stderr, "out of memory\n");
    close(outfd);
    close(infd);
    return ENOMEM;
  }

  // TODO: do the whole file in one frame, using the LZ4F begin/update/end
  // functions.
  ssize_t nr = 0;
  while ((nr = read(infd, inbuf, BLOCK_SIZE)) > 0) {
    ssize_t csz = LZ4F_compressFrame(outbuf, outsize, inbuf, nr, &prefs);
    if (LZ4F_isError(csz)) {
      fprintf(stderr, "error compressing %s: %s\n", infile, LZ4F_getErrorName(csz));
      goto done;
    }

    ssize_t nw = write(outfd, outbuf, csz);
    if (nw != csz) {
      fprintf(stderr, "could not write to %s", outfile);
      if (nw < 0) {
        fprintf(stderr, ": %s", strerror(errno));
      }
      fprintf(stderr, "\n");
      goto done;
    }
  }

  if (nr < 0) {
    fprintf(stderr, "error reading %s: %s\n", infile, strerror(errno));
    goto done;
  }

done:
  free(outbuf);
  close(outfd);
  close(infd);
  return 0;
}

int main(int argc, char* argv[]) {
  int clevel = 1;
  bool decompress = false;
  const char* infile = NULL;
  const char* outfile = NULL;

  for (int i = 1; i < argc; i++) {
    if (!strcmp("-d", argv[i])) {
      decompress = true;
      continue;
    }
    if (!strcmp("-9", argv[i])) {
      clevel = 9;
      continue;
    }
    if (!strcmp("-h", argv[i])) {
      usage(argv[0]);
      return 0;
    }

    if (!infile) {
      infile = argv[i];
      continue;
    }
    if (!outfile) {
      outfile = argv[i];
      continue;
    }

    fprintf(stderr, "Unknown argument: %s\n", argv[i]);
    usage(argv[0]);
    return -1;
  }

  if (!infile || !outfile) {
    usage(argv[0]);
    return 0;
  }

  printf("%scompressing %s into %s", decompress ? "de" : "", infile, outfile);
  if (!decompress) {
    printf(" at level %d", clevel);
  }
  printf("\n");

  if (decompress) {
    return do_decompress(infile, outfile);
  } else {
    return do_compress(infile, outfile, clevel);
  }
}