Merge remote-tracking branch 'ahf/bugs/22066'

3 files changed, 126 insertions, 30 deletions

diff --git a/src/common/compress.c b/src/common/compress.c
index 2e7412fd0d..771f5ab7b6 100644
--- a/src/common/compress.c
+++ b/src/common/compress.c
@@ -27,6 +27,9 @@
 #include "compress_zlib.h"
 #include "compress_zstd.h"
 
+/** Total number of bytes allocated for compression state overhead. */
+static atomic_counter_t total_compress_allocation;
+
 /** @{ */
 /* These macros define the maximum allowable compression factor.  Anything of
  * size greater than CHECK_FOR_COMPRESSION_BOMB_AFTER is not allowed to
@@ -212,12 +215,11 @@ tor_compress(char **out, size_t *out_len,
                            1, LOG_WARN);
 }
 
-/** Given zero or more zlib-compressed or gzip-compressed strings of
- * total length
- * <b>in_len</b> bytes at <b>in</b>, uncompress them into a newly allocated
- * buffer, using the method described in <b>method</b>.  Store the uncompressed
- * string in *<b>out</b>, and its length in *<b>out_len</b>.  Return 0 on
- * success, -1 on failure.
+/** Given zero or more compressed strings of total length <b>in_len</b> bytes
+ * at <b>in</b>, uncompress them into a newly allocated buffer, using the
+ * method described in <b>method</b>.  Store the uncompressed string in
+ * *<b>out</b>, and its length in *<b>out_len</b>.  Return 0 on success, -1 on
+ * failure.
  *
  * If <b>complete_only</b> is true, we consider a truncated input as a
  * failure; otherwise we decompress as much as we can.  Warn about truncated
@@ -367,7 +369,8 @@ tor_compress_header_version_str(compress_method_t method)
 size_t
 tor_compress_get_total_allocation(void)
 {
-  return tor_zlib_get_total_allocation() +
+  return atomic_counter_get(&total_compress_allocation) +
+         tor_zlib_get_total_allocation() +
          tor_lzma_get_total_allocation() +
          tor_zstd_get_total_allocation();
 }
@@ -432,6 +435,8 @@ tor_compress_new(int compress, compress_method_t method,
       goto err;
   }
 
+  atomic_counter_add(&total_compress_allocation,
+                     sizeof(tor_compress_state_t));
   return state;
 
  err:
@@ -504,6 +509,8 @@ tor_compress_free(tor_compress_state_t *state)
       break;
   }
 
+  atomic_counter_sub(&total_compress_allocation,
+                     sizeof(tor_compress_state_t));
   tor_free(state);
 }
 
@@ -513,27 +520,33 @@ tor_compress_state_size(const tor_compress_state_t *state)
 {
   tor_assert(state != NULL);
 
+  size_t size = sizeof(tor_compress_state_t);
+
   switch (state->method) {
     case GZIP_METHOD:
     case ZLIB_METHOD:
-      return tor_zlib_compress_state_size(state->u.zlib_state);
+      size += tor_zlib_compress_state_size(state->u.zlib_state);
+      break;
     case LZMA_METHOD:
-      return tor_lzma_compress_state_size(state->u.lzma_state);
+      size += tor_lzma_compress_state_size(state->u.lzma_state);
+      break;
     case ZSTD_METHOD:
-      return tor_zstd_compress_state_size(state->u.zstd_state);
+      size += tor_zstd_compress_state_size(state->u.zstd_state);
+      break;
     case NO_METHOD:
     case UNKNOWN_METHOD:
-      goto err;
+      break;
   }
 
- err:
-  return 0;
+  return size;
 }
 
 /** Initialize all compression modules. */
 void
 tor_compress_init(void)
 {
+  atomic_counter_init(&total_compress_allocation);
+
   tor_zlib_init();
   tor_lzma_init();
   tor_zstd_init();
diff --git a/src/common/compress_lzma.c b/src/common/compress_lzma.c
index 953971b82d..af1ba7ec11 100644
--- a/src/common/compress_lzma.c
+++ b/src/common/compress_lzma.c
@@ -127,13 +127,44 @@ struct tor_lzma_compress_state_t {
   size_t allocation;
 };
 
+#ifdef HAVE_LZMA
+/** Return an approximate number of bytes stored in memory to hold the LZMA
+ * encoder/decoder state. */
+static size_t
+tor_lzma_state_size_precalc(int compress, compression_level_t level)
+{
+  uint64_t memory_usage;
+
+  if (compress)
+    memory_usage = lzma_easy_encoder_memusage(memory_level(level));
+  else
+    memory_usage = lzma_easy_decoder_memusage(memory_level(level));
+
+  if (memory_usage == UINT64_MAX) {
+    log_warn(LD_GENERAL, "Unsupported compression level passed to LZMA %s",
+                         compress ? "encoder" : "decoder");
+    goto err;
+  }
+
+  if (memory_usage + sizeof(tor_lzma_compress_state_t) > SIZE_MAX)
+    memory_usage = SIZE_MAX;
+  else
+    memory_usage += sizeof(tor_lzma_compress_state_t);
+
+  return (size_t)memory_usage;
+
+ err:
+  return 0;
+}
+#endif // HAVE_LZMA.
+
 /** Construct and return a tor_lzma_compress_state_t object using
  * <b>method</b>. If <b>compress</b>, it's for compression; otherwise it's for
  * decompression. */
 tor_lzma_compress_state_t *
 tor_lzma_compress_new(int compress,
                       compress_method_t method,
-                      compression_level_t compression_level)
+                      compression_level_t level)
 {
   tor_assert(method == LZMA_METHOD);
 
@@ -147,15 +178,10 @@ tor_lzma_compress_new(int compress,
   // also what `tor_malloc_zero()` does.
   result = tor_malloc_zero(sizeof(tor_lzma_compress_state_t));
   result->compress = compress;
-
-  // FIXME(ahf): We should either try to do the pre-calculation that is done
-  // with the zlib backend or use a custom allocator here where we pass our
-  // tor_lzma_compress_state_t as the opaque value.
-  result->allocation = 0;
+  result->allocation = tor_lzma_state_size_precalc(compress, level);
 
   if (compress) {
-    lzma_lzma_preset(&stream_options,
-                     memory_level(compression_level));
+    lzma_lzma_preset(&stream_options, memory_level(level));
 
     retval = lzma_alone_encoder(&result->stream, &stream_options);
 
diff --git a/src/common/compress_zstd.c b/src/common/compress_zstd.c
index deaefc1779..4a8f03e30a 100644
--- a/src/common/compress_zstd.c
+++ b/src/common/compress_zstd.c
@@ -20,7 +20,6 @@
 
 #ifdef HAVE_ZSTD
 #include <zstd.h>
-#include <zstd_errors.h>
 #endif
 
 /** Total number of bytes allocated for Zstandard state. */
@@ -109,27 +108,86 @@ struct tor_zstd_compress_state_t {
   size_t allocation;
 };
 
+#ifdef HAVE_ZSTD
+/** Return an approximate number of bytes stored in memory to hold the
+ * Zstandard compression/decompression state. */
+static size_t
+tor_zstd_state_size_precalc(int compress, int preset)
+{
+  tor_assert(preset > 0);
+
+  size_t memory_usage = sizeof(tor_zstd_compress_state_t);
+
+  // The Zstandard library provides a number of functions that would be useful
+  // here, but they are, unfortunately, still considered experimental and are
+  // thus only available in libzstd if we link against the library statically.
+  //
+  // The code in this function tries to approximate the calculations without
+  // being able to use the following:
+  //
+  // - We do not have access to neither the internal members of ZSTD_CStream
+  //   and ZSTD_DStream and their internal context objects.
+  //
+  // - We cannot use ZSTD_sizeof_CStream() and ZSTD_sizeof_DStream() since they
+  //   are unexposed.
+  //
+  // In the future it might be useful to check if libzstd have started
+  // providing these functions in a stable manner and simplify this function.
+  if (compress) {
+    // We try to approximate the ZSTD_sizeof_CStream(ZSTD_CStream *stream)
+    // function here. This function uses the following fields to make its
+    // estimate:
+
+    // - sizeof(ZSTD_CStream): Around 192 bytes on a 64-bit machine:
+    memory_usage += 192;
+
+    // - ZSTD_sizeof_CCtx(stream->cctx): This function requires access to
+    // variables that are not exposed via the public API. We use a _very_
+    // simplified function to calculate the estimated amount of bytes used in
+    // this struct.
+    memory_usage += (preset - 0.5) * 1024 * 1024;
+    // - ZSTD_sizeof_CDict(stream->cdictLocal): Unused in Tor: 0 bytes.
+    // - stream->outBuffSize: 128 KB:
+    memory_usage += 128 * 1024;
+    // - stream->inBuffSize: 2048 KB:
+    memory_usage += 2048 * 1024;
+  } else {
+    // We try to approximate the ZSTD_sizeof_DStream(ZSTD_DStream *stream)
+    // function here. This function uses the following fields to make its
+    // estimate:
+
+    // - sizeof(ZSTD_DStream): Around 208 bytes on a 64-bit machine:
+    memory_usage += 208;
+    // - ZSTD_sizeof_DCtx(stream->dctx): Around 150 KB.
+    memory_usage += 150 * 1024;
+
+    // - ZSTD_sizeof_DDict(stream->ddictLocal): Unused in Tor: 0 bytes.
+    // - stream->inBuffSize: 0 KB.
+    // - stream->outBuffSize: 0 KB.
+  }
+
+  return memory_usage;
+}
+#endif // HAVE_ZSTD.
+
 /** Construct and return a tor_zstd_compress_state_t object using
  * <b>method</b>. If <b>compress</b>, it's for compression; otherwise it's for
  * decompression. */
 tor_zstd_compress_state_t *
 tor_zstd_compress_new(int compress,
                       compress_method_t method,
-                      compression_level_t compression_level)
+                      compression_level_t level)
 {
   tor_assert(method == ZSTD_METHOD);
 
 #ifdef HAVE_ZSTD
+  const int preset = memory_level(level);
   tor_zstd_compress_state_t *result;
   size_t retval;
 
   result = tor_malloc_zero(sizeof(tor_zstd_compress_state_t));
   result->compress = compress;
-
-  // FIXME(ahf): We should either try to do the pre-calculation that is done
-  // with the zlib backend or use a custom allocator here where we pass our
-  // tor_zstd_compress_state_t as the opaque value.
-  result->allocation = 0;
+  result->allocation = tor_zstd_state_size_precalc(compress, preset);
 
   if (compress) {
     result->u.compress_stream = ZSTD_createCStream();
@@ -139,8 +197,7 @@ tor_zstd_compress_new(int compress,
       goto err;
     }
 
-    retval = ZSTD_initCStream(result->u.compress_stream,
-                              memory_level(compression_level));
+    retval = ZSTD_initCStream(result->u.compress_stream, preset);
 
     if (ZSTD_isError(retval)) {
       log_warn(LD_GENERAL, "Zstandard stream initialization error: %s",