contrib/libarchive/libarchive/archive_write_add_filter_zstd.c

   1 /*-
   2  * Copyright (c) 2017 Sean Purcell
   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(S) ``AS IS'' AND ANY EXPRESS OR
  15  * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
  16  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
  17  * IN NO EVENT SHALL THE AUTHOR(S) BE LIABLE FOR ANY DIRECT, INDIRECT,
  18  * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
  19  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
  20  * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
  21  * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
  22  * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
  23  * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  24  */
  25
  26 #include "archive_platform.h"
  27
  28 __FBSDID("$FreeBSD$");
  29
  30
  31 #ifdef HAVE_ERRNO_H
  32 #include <errno.h>
  33 #endif
  34 #ifdef HAVE_STDLIB_H
  35 #include <stdlib.h>
  36 #endif
  37 #ifdef HAVE_STRING_H
  38 #include <string.h>
  39 #endif
  40 #ifdef HAVE_ZSTD_H
  41 #include <zstd.h>
  42 #endif
  43
  44 #include "archive.h"
  45 #include "archive_private.h"
  46 #include "archive_string.h"
  47 #include "archive_write_private.h"
  48
  49 /* Don't compile this if we don't have zstd.h */
  50
  51 struct private_data {
  52         int              compression_level;
  53 #if HAVE_ZSTD_H && HAVE_LIBZSTD
  54         ZSTD_CStream    *cstream;
  55         int64_t          total_in;
  56         ZSTD_outBuffer   out;
  57 #else
  58         struct archive_write_program_data *pdata;
  59 #endif
  60 };
  61
  62 static int archive_compressor_zstd_options(struct archive_write_filter *,
  63                     const char *, const char *);
  64 static int archive_compressor_zstd_open(struct archive_write_filter *);
  65 static int archive_compressor_zstd_write(struct archive_write_filter *,
  66                     const void *, size_t);
  67 static int archive_compressor_zstd_close(struct archive_write_filter *);
  68 static int archive_compressor_zstd_free(struct archive_write_filter *);
  69 #if HAVE_ZSTD_H && HAVE_LIBZSTD
  70 static int drive_compressor(struct archive_write_filter *,
  71                     struct private_data *, int, const void *, size_t);
  72 #endif
  73
  74
  75 /*
  76  * Add a zstd compression filter to this write handle.
  77  */
  78 int
  79 archive_write_add_filter_zstd(struct archive *_a)
  80 {
  81         struct archive_write *a = (struct archive_write *)_a;
  82         struct archive_write_filter *f = __archive_write_allocate_filter(_a);
  83         struct private_data *data;
  84         archive_check_magic(&a->archive, ARCHIVE_WRITE_MAGIC,
  85             ARCHIVE_STATE_NEW, "archive_write_add_filter_zstd");
  86
  87         data = calloc(1, sizeof(*data));
  88         if (data == NULL) {
  89                 archive_set_error(&a->archive, ENOMEM, "Out of memory");
  90                 return (ARCHIVE_FATAL);
  91         }
  92         f->data = data;
  93         f->open = &archive_compressor_zstd_open;
  94         f->options = &archive_compressor_zstd_options;
  95         f->close = &archive_compressor_zstd_close;
  96         f->free = &archive_compressor_zstd_free;
  97         f->code = ARCHIVE_FILTER_ZSTD;
  98         f->name = "zstd";
  99         data->compression_level = 3; /* Default level used by the zstd CLI */
 100 #if HAVE_ZSTD_H && HAVE_LIBZSTD
 101         data->cstream = ZSTD_createCStream();
 102         if (data->cstream == NULL) {
 103                 free(data);
 104                 archive_set_error(&a->archive, ENOMEM,
 105                     "Failed to allocate zstd compressor object");
 106                 return (ARCHIVE_FATAL);
 107         }
 108
 109         return (ARCHIVE_OK);
 110 #else
 111         data->pdata = __archive_write_program_allocate("zstd");
 112         if (data->pdata == NULL) {
 113                 free(data);
 114                 archive_set_error(&a->archive, ENOMEM, "Out of memory");
 115                 return (ARCHIVE_FATAL);
 116         }
 117         archive_set_error(&a->archive, ARCHIVE_ERRNO_MISC,
 118             "Using external zstd program");
 119         return (ARCHIVE_WARN);
 120 #endif
 121 }
 122
 123 static int
 124 archive_compressor_zstd_free(struct archive_write_filter *f)
 125 {
 126         struct private_data *data = (struct private_data *)f->data;
 127 #if HAVE_ZSTD_H && HAVE_LIBZSTD
 128         ZSTD_freeCStream(data->cstream);
 129         free(data->out.dst);
 130 #else
 131         __archive_write_program_free(data->pdata);
 132 #endif
 133         free(data);
 134         f->data = NULL;
 135         return (ARCHIVE_OK);
 136 }
 137
 138 /*
 139  * Set write options.
 140  */
 141 static int
 142 archive_compressor_zstd_options(struct archive_write_filter *f, const char *key,
 143     const char *value)
 144 {
 145         struct private_data *data = (struct private_data *)f->data;
 146
 147         if (strcmp(key, "compression-level") == 0) {
 148                 int level = atoi(value);
 149 #if HAVE_ZSTD_H && HAVE_LIBZSTD
 150                 if (level < 1 || level > ZSTD_maxCLevel()) {
 151 #else
 152                 /* If we don't have the library, hard-code the max level */
 153                 if (level < 1 || level > 22) {
 154 #endif
 155                         return (ARCHIVE_WARN);
 156                 }
 157                 data->compression_level = level;
 158                 return (ARCHIVE_OK);
 159         }
 160
 161         /* Note: The "warn" return is just to inform the options
 162          * supervisor that we didn't handle it.  It will generate
 163          * a suitable error if no one used this option. */
 164         return (ARCHIVE_WARN);
 165 }
 166
 167 #if HAVE_ZSTD_H && HAVE_LIBZSTD
 168 /*
 169  * Setup callback.
 170  */
 171 static int
 172 archive_compressor_zstd_open(struct archive_write_filter *f)
 173 {
 174         struct private_data *data = (struct private_data *)f->data;
 175
 176         if (data->out.dst == NULL) {
 177                 size_t bs = ZSTD_CStreamOutSize(), bpb;
 178                 if (f->archive->magic == ARCHIVE_WRITE_MAGIC) {
 179                         /* Buffer size should be a multiple number of
 180                          * the of bytes per block for performance. */
 181                         bpb = archive_write_get_bytes_per_block(f->archive);
 182                         if (bpb > bs)
 183                                 bs = bpb;
 184                         else if (bpb != 0)
 185                                 bs -= bs % bpb;
 186                 }
 187                 data->out.size = bs;
 188                 data->out.pos = 0;
 189                 data->out.dst
 190                     = (unsigned char *)malloc(data->out.size);
 191                 if (data->out.dst == NULL) {
 192                         archive_set_error(f->archive, ENOMEM,
 193                             "Can't allocate data for compression buffer");
 194                         return (ARCHIVE_FATAL);
 195                 }
 196         }
 197
 198         f->write = archive_compressor_zstd_write;
 199
 200         if (ZSTD_isError(ZSTD_initCStream(data->cstream,
 201             data->compression_level))) {
 202                 archive_set_error(f->archive, ARCHIVE_ERRNO_MISC,
 203                     "Internal error initializing zstd compressor object");
 204                 return (ARCHIVE_FATAL);
 205         }
 206
 207         return (ARCHIVE_OK);
 208 }
 209
 210 /*
 211  * Write data to the compressed stream.
 212  */
 213 static int
 214 archive_compressor_zstd_write(struct archive_write_filter *f, const void *buff,
 215     size_t length)
 216 {
 217         struct private_data *data = (struct private_data *)f->data;
 218         int ret;
 219
 220         /* Update statistics */
 221         data->total_in += length;
 222
 223         if ((ret = drive_compressor(f, data, 0, buff, length)) != ARCHIVE_OK)
 224                 return (ret);
 225
 226         return (ARCHIVE_OK);
 227 }
 228
 229 /*
 230  * Finish the compression...
 231  */
 232 static int
 233 archive_compressor_zstd_close(struct archive_write_filter *f)
 234 {
 235         struct private_data *data = (struct private_data *)f->data;
 236
 237         /* Finish zstd frame */
 238         return drive_compressor(f, data, 1, NULL, 0);
 239 }
 240
 241 /*
 242  * Utility function to push input data through compressor,
 243  * writing full output blocks as necessary.
 244  *
 245  * Note that this handles both the regular write case (finishing ==
 246  * false) and the end-of-archive case (finishing == true).
 247  */
 248 static int
 249 drive_compressor(struct archive_write_filter *f,
 250     struct private_data *data, int finishing, const void *src, size_t length)
 251 {
 252         ZSTD_inBuffer in = (ZSTD_inBuffer) { src, length, 0 };
 253
 254         for (;;) {
 255                 if (data->out.pos == data->out.size) {
 256                         const int ret = __archive_write_filter(f->next_filter,
 257                             data->out.dst, data->out.size);
 258                         if (ret != ARCHIVE_OK)
 259                                 return (ARCHIVE_FATAL);
 260                         data->out.pos = 0;
 261                 }
 262
 263                 /* If there's nothing to do, we're done. */
 264                 if (!finishing && in.pos == in.size)
 265                         return (ARCHIVE_OK);
 266
 267                 {
 268                         const size_t zstdret = !finishing ?
 269                             ZSTD_compressStream(data->cstream, &data->out, &in)
 270                             : ZSTD_endStream(data->cstream, &data->out);
 271
 272                         if (ZSTD_isError(zstdret)) {
 273                                 archive_set_error(f->archive,
 274                                     ARCHIVE_ERRNO_MISC,
 275                                     "Zstd compression failed: %s",
 276                                     ZSTD_getErrorName(zstdret));
 277                                 return (ARCHIVE_FATAL);
 278                         }
 279
 280                         /* If we're finishing, 0 means nothing left to flush */
 281                         if (finishing && zstdret == 0) {
 282                                 const int ret = __archive_write_filter(f->next_filter,
 283                                     data->out.dst, data->out.pos);
 284                                 return (ret);
 285                         }
 286                 }
 287         }
 288 }
 289
 290 #else /* HAVE_ZSTD_H && HAVE_LIBZSTD */
 291
 292 static int
 293 archive_compressor_zstd_open(struct archive_write_filter *f)
 294 {
 295         struct private_data *data = (struct private_data *)f->data;
 296         struct archive_string as;
 297         int r;
 298
 299         archive_string_init(&as);
 300         archive_string_sprintf(&as, "zstd -%d", data->compression_level);
 301
 302         f->write = archive_compressor_zstd_write;
 303         r = __archive_write_program_open(f, data->pdata, as.s);
 304         archive_string_free(&as);
 305         return (r);
 306 }
 307
 308 static int
 309 archive_compressor_zstd_write(struct archive_write_filter *f, const void *buff,
 310     size_t length)
 311 {
 312         struct private_data *data = (struct private_data *)f->data;
 313
 314         return __archive_write_program_write(f, data->pdata, buff, length);
 315 }
 316
 317 static int
 318 archive_compressor_zstd_close(struct archive_write_filter *f)
 319 {
 320         struct private_data *data = (struct private_data *)f->data;
 321
 322         return __archive_write_program_close(f, data->pdata);
 323 }
 324
 325 #endif /* HAVE_ZSTD_H && HAVE_LIBZSTD */