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         int ret;
 176
 177         ret = __archive_write_open_filter(f->next_filter);
 178         if (ret != ARCHIVE_OK)
 179                 return (ret);
 180
 181         if (data->out.dst == NULL) {
 182                 size_t bs = ZSTD_CStreamOutSize(), bpb;
 183                 if (f->archive->magic == ARCHIVE_WRITE_MAGIC) {
 184                         /* Buffer size should be a multiple number of
 185                          * the of bytes per block for performance. */
 186                         bpb = archive_write_get_bytes_per_block(f->archive);
 187                         if (bpb > bs)
 188                                 bs = bpb;
 189                         else if (bpb != 0)
 190                                 bs -= bs % bpb;
 191                 }
 192                 data->out.size = bs;
 193                 data->out.pos = 0;
 194                 data->out.dst
 195                     = (unsigned char *)malloc(data->out.size);
 196                 if (data->out.dst == NULL) {
 197                         archive_set_error(f->archive, ENOMEM,
 198                             "Can't allocate data for compression buffer");
 199                         return (ARCHIVE_FATAL);
 200                 }
 201         }
 202
 203         f->write = archive_compressor_zstd_write;
 204
 205         if (ZSTD_isError(ZSTD_initCStream(data->cstream,
 206             data->compression_level))) {
 207                 archive_set_error(f->archive, ARCHIVE_ERRNO_MISC,
 208                     "Internal error initializing zstd compressor object");
 209                 return (ARCHIVE_FATAL);
 210         }
 211
 212         return (ARCHIVE_OK);
 213 }
 214
 215 /*
 216  * Write data to the compressed stream.
 217  */
 218 static int
 219 archive_compressor_zstd_write(struct archive_write_filter *f, const void *buff,
 220     size_t length)
 221 {
 222         struct private_data *data = (struct private_data *)f->data;
 223         int ret;
 224
 225         /* Update statistics */
 226         data->total_in += length;
 227
 228         if ((ret = drive_compressor(f, data, 0, buff, length)) != ARCHIVE_OK)
 229                 return (ret);
 230
 231         return (ARCHIVE_OK);
 232 }
 233
 234 /*
 235  * Finish the compression...
 236  */
 237 static int
 238 archive_compressor_zstd_close(struct archive_write_filter *f)
 239 {
 240         struct private_data *data = (struct private_data *)f->data;
 241         int r1, r2;
 242
 243         /* Finish zstd frame */
 244         r1 = drive_compressor(f, data, 1, NULL, 0);
 245
 246         r2 = __archive_write_close_filter(f->next_filter);
 247
 248         return r1 < r2 ? r1 : r2;
 249 }
 250
 251 /*
 252  * Utility function to push input data through compressor,
 253  * writing full output blocks as necessary.
 254  *
 255  * Note that this handles both the regular write case (finishing ==
 256  * false) and the end-of-archive case (finishing == true).
 257  */
 258 static int
 259 drive_compressor(struct archive_write_filter *f,
 260     struct private_data *data, int finishing, const void *src, size_t length)
 261 {
 262         ZSTD_inBuffer in = (ZSTD_inBuffer) { src, length, 0 };
 263
 264         for (;;) {
 265                 if (data->out.pos == data->out.size) {
 266                         const int ret = __archive_write_filter(f->next_filter,
 267                             data->out.dst, data->out.size);
 268                         if (ret != ARCHIVE_OK)
 269                                 return (ARCHIVE_FATAL);
 270                         data->out.pos = 0;
 271                 }
 272
 273                 /* If there's nothing to do, we're done. */
 274                 if (!finishing && in.pos == in.size)
 275                         return (ARCHIVE_OK);
 276
 277                 {
 278                         const size_t zstdret = !finishing ?
 279                             ZSTD_compressStream(data->cstream, &data->out, &in)
 280                             : ZSTD_endStream(data->cstream, &data->out);
 281
 282                         if (ZSTD_isError(zstdret)) {
 283                                 archive_set_error(f->archive,
 284                                     ARCHIVE_ERRNO_MISC,
 285                                     "Zstd compression failed: %s",
 286                                     ZSTD_getErrorName(zstdret));
 287                                 return (ARCHIVE_FATAL);
 288                         }
 289
 290                         /* If we're finishing, 0 means nothing left to flush */
 291                         if (finishing && zstdret == 0) {
 292                                 const int ret = __archive_write_filter(f->next_filter,
 293                                     data->out.dst, data->out.pos);
 294                                 return (ret);
 295                         }
 296                 }
 297         }
 298 }
 299
 300 #else /* HAVE_ZSTD_H && HAVE_LIBZSTD */
 301
 302 static int
 303 archive_compressor_zstd_open(struct archive_write_filter *f)
 304 {
 305         struct private_data *data = (struct private_data *)f->data;
 306         struct archive_string as;
 307         int r;
 308
 309         archive_string_init(&as);
 310         archive_string_sprintf(&as, "zstd -%d", data->compression_level);
 311
 312         f->write = archive_compressor_zstd_write;
 313         r = __archive_write_program_open(f, data->pdata, as.s);
 314         archive_string_free(&as);
 315         return (r);
 316 }
 317
 318 static int
 319 archive_compressor_zstd_write(struct archive_write_filter *f, const void *buff,
 320     size_t length)
 321 {
 322         struct private_data *data = (struct private_data *)f->data;
 323
 324         return __archive_write_program_write(f, data->pdata, buff, length);
 325 }
 326
 327 static int
 328 archive_compressor_zstd_close(struct archive_write_filter *f)
 329 {
 330         struct private_data *data = (struct private_data *)f->data;
 331
 332         return __archive_write_program_close(f, data->pdata);
 333 }
 334
 335 #endif /* HAVE_ZSTD_H && HAVE_LIBZSTD */