1
#include "source/extensions/compression/zstd/decompressor/zstd_decompressor_impl.h"
2

            
3
#include "source/common/runtime/runtime_features.h"
4

            
5
namespace Envoy {
6
namespace Extensions {
7
namespace Compression {
8
namespace Zstd {
9
namespace Decompressor {
10

            
11
namespace {
12

            
13
// How many times the output buffer is allowed to be bigger than the size of
14
// accumulated input. This value is used to detect compression bombs.
15
// TODO(rojkov): Re-design the Decompressor interface to handle compression
16
// bombs gracefully instead of this quick solution.
17
constexpr uint64_t MaxInflateRatio = 100;
18

            
19
} // namespace
20

            
21
ZstdDecompressorImpl::ZstdDecompressorImpl(Stats::Scope& scope, const std::string& stats_prefix,
22
                                           const ZstdDDictManagerPtr& ddict_manager,
23
                                           uint32_t chunk_size)
24
342
    : Envoy::Compression::Zstd::Common::Base(chunk_size), dctx_(ZSTD_createDCtx(), &ZSTD_freeDCtx),
25
342
      ddict_manager_(ddict_manager), stats_(generateStats(stats_prefix, scope)) {}
26

            
27
void ZstdDecompressorImpl::decompress(const Buffer::Instance& input_buffer,
28
341
                                      Buffer::Instance& output_buffer) {
29
341
  uint64_t limit = MaxInflateRatio * input_buffer.length();
30

            
31
4937
  for (const Buffer::RawSlice& input_slice : input_buffer.getRawSlices()) {
32
4937
    if (input_slice.len_ > 0) {
33
4937
      if (ddict_manager_ && !is_dictionary_set_) {
34
174
        is_dictionary_set_ = true;
35
        // If id == 0, it means that dictionary id could not be decoded.
36
174
        dictionary_id_ =
37
174
            ZSTD_getDictID_fromFrame(static_cast<uint8_t*>(input_slice.mem_), input_slice.len_);
38
174
        if (dictionary_id_ != 0) {
39
174
          auto dictionary = ddict_manager_->getDictionaryById(dictionary_id_);
40
174
          if (!dictionary) {
41
3
            stats_.zstd_dictionary_error_.inc();
42
3
            return;
43
3
          }
44
171
          const size_t result = ZSTD_DCtx_refDDict(dctx_.get(), dictionary);
45
171
          if (isError(result)) {
46
            return;
47
          }
48
171
        }
49
174
      }
50

            
51
4934
      setInput(input_slice);
52
4934
      if (!process(output_buffer)) {
53
1
        return;
54
1
      }
55
4933
      if (Runtime::runtimeFeatureEnabled(
56
4933
              "envoy.reloadable_features.enable_compression_bomb_protection") &&
57
4933
          (output_buffer.length() > limit)) {
58
1
        stats_.zstd_generic_error_.inc();
59
1
        ENVOY_LOG(trace,
60
1
                  "excessive decompression ratio detected: output "
61
1
                  "size {} for input size {}",
62
1
                  output_buffer.length(), input_buffer.length());
63
1
        return;
64
1
      }
65
4933
    }
66
4937
  }
67
341
}
68

            
69
4934
bool ZstdDecompressorImpl::process(Buffer::Instance& output_buffer) {
70
28189
  while (input_.pos < input_.size) {
71
23256
    const size_t result = ZSTD_decompressStream(dctx_.get(), &output_, &input_);
72
23256
    if (isError(result)) {
73
1
      return false;
74
1
    }
75

            
76
23255
    getOutput(output_buffer);
77
23255
  }
78

            
79
4933
  return true;
80
4934
}
81

            
82
23433
bool ZstdDecompressorImpl::isError(size_t result) {
83
23433
  switch (ZSTD_getErrorCode(result)) {
84
23427
  case ZSTD_error_no_error:
85
23427
    return false;
86
1
  case ZSTD_error_memory_allocation:
87
1
    stats_.zstd_memory_error_.inc();
88
1
    break;
89
1
  case ZSTD_error_dictionary_corrupted:
90
2
  case ZSTD_error_dictionary_wrong:
91
2
    stats_.zstd_dictionary_error_.inc();
92
2
    break;
93
1
  case ZSTD_error_checksum_wrong:
94
1
    stats_.zstd_checksum_wrong_error_.inc();
95
1
    break;
96
2
  default:
97
2
    stats_.zstd_generic_error_.inc();
98
2
    break;
99
23433
  }
100
6
  return true;
101
23433
}
102

            
103
} // namespace Decompressor
104
} // namespace Zstd
105
} // namespace Compression
106
} // namespace Extensions
107
} // namespace Envoy