epan/tvbuff_zstd.c

   1 /* tvbuff_zstd.c
   2  * Copyright 2022, Kevin Albertson <kevin.eric.albertson [AT] gmail.com>
   3  *
   4  * Wireshark - Network traffic analyzer
   5  * By Gerald Combs <gerald@wireshark.org>
   6  * Copyright 1998 Gerald Combs
   7  *
   8  * SPDX-License-Identifier: GPL-2.0-or-later
   9  */
  10
  11 /*
  12  * Decompress ZSTD: http://facebook.github.io/zstd/
  13  */
  14
  15 #include "config.h"
  16
  17 #ifdef HAVE_ZSTD
  18 #include <zstd.h>
  19 #endif
  20
  21 #include "proto.h" // DISSECTOR_ASSERT_HINT
  22 #include "tvbuff.h"
  23
  24 #include "tvbuff-int.h" // tvb_add_to_chain
  25
  26 #define MAX_LOOP_ITERATIONS 100
  27
  28 tvbuff_t *tvb_uncompress_zstd(tvbuff_t *tvb, const int offset, int comprlen)
  29 {
  30 #ifndef HAVE_ZSTD
  31     // Cast to void to silence unused warnings.
  32     (void)tvb;
  33     (void)offset;
  34     (void)comprlen;
  35     return NULL;
  36 #else
  37     ZSTD_inBuffer input = {tvb_memdup(NULL, tvb, offset, comprlen), comprlen, 0};
  38     ZSTD_DStream *zds = ZSTD_createDStream();
  39     size_t rc = 0;
  40     uint8_t *uncompr = NULL;
  41     size_t uncompr_len = 0;
  42     bool ok = false;
  43     int count = 0;
  44
  45     // ZSTD does not consume the last byte of the frame until it has flushed all of the decompressed data of the frame.
  46     // Therefore, loop while there is more input.
  47     ZSTD_outBuffer output = {g_malloc(ZSTD_DStreamOutSize()), ZSTD_DStreamOutSize(), 0};
  48     while (input.pos < input.size && count < MAX_LOOP_ITERATIONS)
  49     {
  50         rc = ZSTD_decompressStream(zds, &output, &input);
  51         if (ZSTD_isError(rc))
  52         {
  53             goto end;
  54         }
  55
  56         if (output.pos > 0)
  57         {
  58             if (!uncompr)
  59             {
  60                 DISSECTOR_ASSERT (uncompr_len == 0);
  61                 uncompr = g_malloc(output.pos);
  62             } else {
  63                 uncompr = g_realloc(uncompr, uncompr_len + output.pos);
  64             }
  65             memcpy (uncompr + uncompr_len, output.dst, output.pos);
  66             uncompr_len += output.pos;
  67             // Reset the output buffer.
  68             output.pos = 0;
  69         }
  70         count++;
  71         DISSECTOR_ASSERT_HINT(count < MAX_LOOP_ITERATIONS, "MAX_LOOP_ITERATIONS exceeded");
  72     }
  73     if (rc > 0)
  74     {
  75         // There is extra data that was not decompressed.
  76         goto end;
  77     }
  78
  79     ok = true;
  80 end:
  81     g_free((void *)output.dst);
  82     wmem_free(NULL, (void *)input.src);
  83     ZSTD_freeDStream(zds);
  84     if (ok)
  85     {
  86         tvbuff_t *uncompr_tvb;
  87         uncompr_tvb = tvb_new_real_data (uncompr, (unsigned)uncompr_len, (unsigned)uncompr_len);
  88         tvb_set_free_cb (uncompr_tvb, g_free);
  89         return uncompr_tvb;
  90     }
  91
  92     if (uncompr)
  93     {
  94         g_free (uncompr);
  95     }
  96
  97     return NULL;
  98 #endif /* HAVE_ZSTD */
  99 }
 100
 101 tvbuff_t *tvb_child_uncompress_zstd(tvbuff_t *parent, tvbuff_t *tvb, const int offset, int comprlen)
 102 {
 103     tvbuff_t *uncompressed = tvb_uncompress_zstd(tvb, offset, comprlen);
 104     if (!uncompressed)
 105     {
 106         return uncompressed;
 107     }
 108     tvb_add_to_chain(parent, uncompressed);
 109     return uncompressed;
 110 }