compress_zstd.c 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448
  1. /* Copyright (c) 2004, Roger Dingledine.
  2. * Copyright (c) 2004-2006, Roger Dingledine, Nick Mathewson.
  3. * Copyright (c) 2007-2017, The Tor Project, Inc. */
  4. /* See LICENSE for licensing information */
  5. /**
  6. * \file compress_zstd.c
  7. * \brief Compression backend for Zstandard.
  8. *
  9. * This module should never be invoked directly. Use the compress module
  10. * instead.
  11. **/
  12. #include "orconfig.h"
  13. #include "util.h"
  14. #include "torlog.h"
  15. #include "compress.h"
  16. #include "compress_zstd.h"
  17. #ifdef HAVE_ZSTD
  18. #ifdef HAVE_CFLAG_WUNUSED_CONST_VARIABLE
  19. DISABLE_GCC_WARNING(unused-const-variable)
  20. #endif
  21. #include <zstd.h>
  22. #ifdef HAVE_CFLAG_WUNUSED_CONST_VARIABLE
  23. ENABLE_GCC_WARNING(unused-const-variable)
  24. #endif
  25. #endif
  26. /** Total number of bytes allocated for Zstandard state. */
  27. static atomic_counter_t total_zstd_allocation;
  28. #ifdef HAVE_ZSTD
  29. /** Given <b>level</b> return the memory level. */
  30. static int
  31. memory_level(compression_level_t level)
  32. {
  33. switch (level) {
  34. default:
  35. case BEST_COMPRESSION:
  36. case HIGH_COMPRESSION: return 9;
  37. case MEDIUM_COMPRESSION: return 8;
  38. case LOW_COMPRESSION: return 7;
  39. }
  40. }
  41. #endif /* defined(HAVE_ZSTD) */
  42. /** Return 1 if Zstandard compression is supported; otherwise 0. */
  43. int
  44. tor_zstd_method_supported(void)
  45. {
  46. #ifdef HAVE_ZSTD
  47. return 1;
  48. #else
  49. return 0;
  50. #endif
  51. }
  52. /** Return a string representation of the version of the currently running
  53. * version of libzstd. Returns NULL if Zstandard is unsupported. */
  54. const char *
  55. tor_zstd_get_version_str(void)
  56. {
  57. #ifdef HAVE_ZSTD
  58. static char version_str[16];
  59. size_t version_number;
  60. version_number = ZSTD_versionNumber();
  61. tor_snprintf(version_str, sizeof(version_str),
  62. "%d.%d.%d",
  63. (int) version_number / 10000 % 100,
  64. (int) version_number / 100 % 100,
  65. (int) version_number % 100);
  66. return version_str;
  67. #else /* !(defined(HAVE_ZSTD)) */
  68. return NULL;
  69. #endif /* defined(HAVE_ZSTD) */
  70. }
  71. /** Return a string representation of the version of the version of libzstd
  72. * used at compilation time. Returns NULL if Zstandard is unsupported. */
  73. const char *
  74. tor_zstd_get_header_version_str(void)
  75. {
  76. #ifdef HAVE_ZSTD
  77. return ZSTD_VERSION_STRING;
  78. #else
  79. return NULL;
  80. #endif
  81. }
  82. /** Internal Zstandard state for incremental compression/decompression.
  83. * The body of this struct is not exposed. */
  84. struct tor_zstd_compress_state_t {
  85. #ifdef HAVE_ZSTD
  86. union {
  87. /** Compression stream. Used when <b>compress</b> is true. */
  88. ZSTD_CStream *compress_stream;
  89. /** Decompression stream. Used when <b>compress</b> is false. */
  90. ZSTD_DStream *decompress_stream;
  91. } u; /**< Zstandard stream objects. */
  92. #endif /* defined(HAVE_ZSTD) */
  93. int compress; /**< True if we are compressing; false if we are inflating */
  94. int have_called_end; /**< True if we are compressing and we've called
  95. * ZSTD_endStream */
  96. /** Number of bytes read so far. Used to detect compression bombs. */
  97. size_t input_so_far;
  98. /** Number of bytes written so far. Used to detect compression bombs. */
  99. size_t output_so_far;
  100. /** Approximate number of bytes allocated for this object. */
  101. size_t allocation;
  102. };
  103. #ifdef HAVE_ZSTD
  104. /** Return an approximate number of bytes stored in memory to hold the
  105. * Zstandard compression/decompression state. */
  106. static size_t
  107. tor_zstd_state_size_precalc(int compress, int preset)
  108. {
  109. tor_assert(preset > 0);
  110. size_t memory_usage = sizeof(tor_zstd_compress_state_t);
  111. // The Zstandard library provides a number of functions that would be useful
  112. // here, but they are, unfortunately, still considered experimental and are
  113. // thus only available in libzstd if we link against the library statically.
  114. //
  115. // The code in this function tries to approximate the calculations without
  116. // being able to use the following:
  117. //
  118. // - We do not have access to neither the internal members of ZSTD_CStream
  119. // and ZSTD_DStream and their internal context objects.
  120. //
  121. // - We cannot use ZSTD_sizeof_CStream() and ZSTD_sizeof_DStream() since they
  122. // are unexposed.
  123. //
  124. // In the future it might be useful to check if libzstd have started
  125. // providing these functions in a stable manner and simplify this function.
  126. if (compress) {
  127. // We try to approximate the ZSTD_sizeof_CStream(ZSTD_CStream *stream)
  128. // function here. This function uses the following fields to make its
  129. // estimate:
  130. // - sizeof(ZSTD_CStream): Around 192 bytes on a 64-bit machine:
  131. memory_usage += 192;
  132. // - ZSTD_sizeof_CCtx(stream->cctx): This function requires access to
  133. // variables that are not exposed via the public API. We use a _very_
  134. // simplified function to calculate the estimated amount of bytes used in
  135. // this struct.
  136. // memory_usage += (preset - 0.5) * 1024 * 1024;
  137. memory_usage += (preset * 1024 * 1024) - (512 * 1024);
  138. // - ZSTD_sizeof_CDict(stream->cdictLocal): Unused in Tor: 0 bytes.
  139. // - stream->outBuffSize: 128 KB:
  140. memory_usage += 128 * 1024;
  141. // - stream->inBuffSize: 2048 KB:
  142. memory_usage += 2048 * 1024;
  143. } else {
  144. // We try to approximate the ZSTD_sizeof_DStream(ZSTD_DStream *stream)
  145. // function here. This function uses the following fields to make its
  146. // estimate:
  147. // - sizeof(ZSTD_DStream): Around 208 bytes on a 64-bit machine:
  148. memory_usage += 208;
  149. // - ZSTD_sizeof_DCtx(stream->dctx): Around 150 KB.
  150. memory_usage += 150 * 1024;
  151. // - ZSTD_sizeof_DDict(stream->ddictLocal): Unused in Tor: 0 bytes.
  152. // - stream->inBuffSize: 0 KB.
  153. // - stream->outBuffSize: 0 KB.
  154. }
  155. return memory_usage;
  156. }
  157. #endif /* defined(HAVE_ZSTD) */
  158. /** Construct and return a tor_zstd_compress_state_t object using
  159. * <b>method</b>. If <b>compress</b>, it's for compression; otherwise it's for
  160. * decompression. */
  161. tor_zstd_compress_state_t *
  162. tor_zstd_compress_new(int compress,
  163. compress_method_t method,
  164. compression_level_t level)
  165. {
  166. tor_assert(method == ZSTD_METHOD);
  167. #ifdef HAVE_ZSTD
  168. const int preset = memory_level(level);
  169. tor_zstd_compress_state_t *result;
  170. size_t retval;
  171. result = tor_malloc_zero(sizeof(tor_zstd_compress_state_t));
  172. result->compress = compress;
  173. result->allocation = tor_zstd_state_size_precalc(compress, preset);
  174. if (compress) {
  175. result->u.compress_stream = ZSTD_createCStream();
  176. if (result->u.compress_stream == NULL) {
  177. // LCOV_EXCL_START
  178. log_warn(LD_GENERAL, "Error while creating Zstandard compression "
  179. "stream");
  180. goto err;
  181. // LCOV_EXCL_STOP
  182. }
  183. retval = ZSTD_initCStream(result->u.compress_stream, preset);
  184. if (ZSTD_isError(retval)) {
  185. // LCOV_EXCL_START
  186. log_warn(LD_GENERAL, "Zstandard stream initialization error: %s",
  187. ZSTD_getErrorName(retval));
  188. goto err;
  189. // LCOV_EXCL_STOP
  190. }
  191. } else {
  192. result->u.decompress_stream = ZSTD_createDStream();
  193. if (result->u.decompress_stream == NULL) {
  194. // LCOV_EXCL_START
  195. log_warn(LD_GENERAL, "Error while creating Zstandard decompression "
  196. "stream");
  197. goto err;
  198. // LCOV_EXCL_STOP
  199. }
  200. retval = ZSTD_initDStream(result->u.decompress_stream);
  201. if (ZSTD_isError(retval)) {
  202. // LCOV_EXCL_START
  203. log_warn(LD_GENERAL, "Zstandard stream initialization error: %s",
  204. ZSTD_getErrorName(retval));
  205. goto err;
  206. // LCOV_EXCL_STOP
  207. }
  208. }
  209. atomic_counter_add(&total_zstd_allocation, result->allocation);
  210. return result;
  211. err:
  212. // LCOV_EXCL_START
  213. if (compress) {
  214. ZSTD_freeCStream(result->u.compress_stream);
  215. } else {
  216. ZSTD_freeDStream(result->u.decompress_stream);
  217. }
  218. tor_free(result);
  219. return NULL;
  220. // LCOV_EXCL_STOP
  221. #else /* !(defined(HAVE_ZSTD)) */
  222. (void)compress;
  223. (void)method;
  224. (void)level;
  225. return NULL;
  226. #endif /* defined(HAVE_ZSTD) */
  227. }
  228. /** Compress/decompress some bytes using <b>state</b>. Read up to
  229. * *<b>in_len</b> bytes from *<b>in</b>, and write up to *<b>out_len</b> bytes
  230. * to *<b>out</b>, adjusting the values as we go. If <b>finish</b> is true,
  231. * we've reached the end of the input.
  232. *
  233. * Return TOR_COMPRESS_DONE if we've finished the entire
  234. * compression/decompression.
  235. * Return TOR_COMPRESS_OK if we're processed everything from the input.
  236. * Return TOR_COMPRESS_BUFFER_FULL if we're out of space on <b>out</b>.
  237. * Return TOR_COMPRESS_ERROR if the stream is corrupt.
  238. */
  239. tor_compress_output_t
  240. tor_zstd_compress_process(tor_zstd_compress_state_t *state,
  241. char **out, size_t *out_len,
  242. const char **in, size_t *in_len,
  243. int finish)
  244. {
  245. #ifdef HAVE_ZSTD
  246. size_t retval;
  247. tor_assert(state != NULL);
  248. tor_assert(*in_len <= UINT_MAX);
  249. tor_assert(*out_len <= UINT_MAX);
  250. ZSTD_inBuffer input = { *in, *in_len, 0 };
  251. ZSTD_outBuffer output = { *out, *out_len, 0 };
  252. if (BUG(finish == 0 && state->have_called_end)) {
  253. finish = 1;
  254. }
  255. if (state->compress) {
  256. if (! state->have_called_end)
  257. retval = ZSTD_compressStream(state->u.compress_stream,
  258. &output, &input);
  259. else
  260. retval = 0;
  261. } else {
  262. retval = ZSTD_decompressStream(state->u.decompress_stream,
  263. &output, &input);
  264. }
  265. state->input_so_far += input.pos;
  266. state->output_so_far += output.pos;
  267. *out = (char *)output.dst + output.pos;
  268. *out_len = output.size - output.pos;
  269. *in = (char *)input.src + input.pos;
  270. *in_len = input.size - input.pos;
  271. if (! state->compress &&
  272. tor_compress_is_compression_bomb(state->input_so_far,
  273. state->output_so_far)) {
  274. log_warn(LD_DIR, "Possible compression bomb; abandoning stream.");
  275. return TOR_COMPRESS_ERROR;
  276. }
  277. if (ZSTD_isError(retval)) {
  278. log_warn(LD_GENERAL, "Zstandard %s didn't finish: %s.",
  279. state->compress ? "compression" : "decompression",
  280. ZSTD_getErrorName(retval));
  281. return TOR_COMPRESS_ERROR;
  282. }
  283. if (state->compress && !state->have_called_end) {
  284. retval = ZSTD_flushStream(state->u.compress_stream, &output);
  285. *out = (char *)output.dst + output.pos;
  286. *out_len = output.size - output.pos;
  287. if (ZSTD_isError(retval)) {
  288. log_warn(LD_GENERAL, "Zstandard compression unable to flush: %s.",
  289. ZSTD_getErrorName(retval));
  290. return TOR_COMPRESS_ERROR;
  291. }
  292. // ZSTD_flushStream returns 0 if the frame is done, or >0 if it
  293. // is incomplete.
  294. if (retval > 0) {
  295. return TOR_COMPRESS_BUFFER_FULL;
  296. }
  297. }
  298. if (!finish) {
  299. // The caller says we're not done with the input, so no need to write an
  300. // epilogue.
  301. return TOR_COMPRESS_OK;
  302. } else if (state->compress) {
  303. if (*in_len) {
  304. // We say that we're not done with the input, so we can't write an
  305. // epilogue.
  306. return TOR_COMPRESS_OK;
  307. }
  308. retval = ZSTD_endStream(state->u.compress_stream, &output);
  309. state->have_called_end = 1;
  310. *out = (char *)output.dst + output.pos;
  311. *out_len = output.size - output.pos;
  312. if (ZSTD_isError(retval)) {
  313. log_warn(LD_GENERAL, "Zstandard compression unable to write "
  314. "epilogue: %s.",
  315. ZSTD_getErrorName(retval));
  316. return TOR_COMPRESS_ERROR;
  317. }
  318. // endStream returns the number of bytes that is needed to write the
  319. // epilogue.
  320. if (retval > 0)
  321. return TOR_COMPRESS_BUFFER_FULL;
  322. return TOR_COMPRESS_DONE;
  323. } else /* if (!state->compress) */ {
  324. // ZSTD_decompressStream returns 0 if the frame is done, or >0 if it
  325. // is incomplete.
  326. // We check this above.
  327. tor_assert_nonfatal(!ZSTD_isError(retval));
  328. // Start a new frame if this frame is done
  329. if (retval == 0)
  330. return TOR_COMPRESS_DONE;
  331. // Don't check out_len, it might have some space left if the next output
  332. // chunk is larger than the remaining space
  333. else if (*in_len > 0)
  334. return TOR_COMPRESS_BUFFER_FULL;
  335. else
  336. return TOR_COMPRESS_OK;
  337. }
  338. #else /* !(defined(HAVE_ZSTD)) */
  339. (void)state;
  340. (void)out;
  341. (void)out_len;
  342. (void)in;
  343. (void)in_len;
  344. (void)finish;
  345. return TOR_COMPRESS_ERROR;
  346. #endif /* defined(HAVE_ZSTD) */
  347. }
  348. /** Deallocate <b>state</b>. */
  349. void
  350. tor_zstd_compress_free_(tor_zstd_compress_state_t *state)
  351. {
  352. if (state == NULL)
  353. return;
  354. atomic_counter_sub(&total_zstd_allocation, state->allocation);
  355. #ifdef HAVE_ZSTD
  356. if (state->compress) {
  357. ZSTD_freeCStream(state->u.compress_stream);
  358. } else {
  359. ZSTD_freeDStream(state->u.decompress_stream);
  360. }
  361. #endif /* defined(HAVE_ZSTD) */
  362. tor_free(state);
  363. }
  364. /** Return the approximate number of bytes allocated for <b>state</b>. */
  365. size_t
  366. tor_zstd_compress_state_size(const tor_zstd_compress_state_t *state)
  367. {
  368. tor_assert(state != NULL);
  369. return state->allocation;
  370. }
  371. /** Return the approximate number of bytes allocated for all Zstandard
  372. * states. */
  373. size_t
  374. tor_zstd_get_total_allocation(void)
  375. {
  376. return atomic_counter_get(&total_zstd_allocation);
  377. }
  378. /** Initialize the zstd module */
  379. void
  380. tor_zstd_init(void)
  381. {
  382. atomic_counter_init(&total_zstd_allocation);
  383. }