1 files changed, 112 insertions, 17 deletions

diff --git a/src/common/compress.c b/src/common/compress.c
index 5a3359aa05..6fe4569868 100644
--- a/src/common/compress.c
+++ b/src/common/compress.c
@@ -24,9 +24,13 @@
 #include "torlog.h"
 #include "compress.h"
 #include "compress_lzma.h"
+#include "compress_none.h"
 #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
@@ -64,8 +68,12 @@ guess_compress_size(int compress, compress_method_t method,
                     size_t in_len)
 {
   // ignore these for now.
-  (void)method;
   (void)compression_level;
+  if (method == NO_METHOD) {
+    /* Guess that we'll need an extra byte, to avoid a needless realloc
+     * for nul-termination */
+    return (in_len < SIZE_MAX) ? in_len + 1 : in_len;
+  }
 
   /* Always guess a factor of 2. */
   if (compress) {
@@ -212,12 +220,17 @@ 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 any bytes are written to <b>out</b>, an extra byte NUL is always
+ * written at the end, but not counted in <b>out_len</b>.  This is a
+ * safety feature to ensure that the output can be treated as a
+ * NUL-terminated string -- though of course, callers should check
+ * out_len anyway.
  *
  * 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
@@ -247,8 +260,8 @@ detect_compression_method(const char *in, size_t in_len)
   } else if (in_len > 2 && (in[0] & 0x0f) == 8 &&
              (ntohs(get_uint16(in)) % 31) == 0) {
     return ZLIB_METHOD;
-  } else if (in_len > 3 &&
-             fast_memeq(in, "\x5d\x00\x00\x00", 4)) {
+  } else if (in_len > 2 &&
+             fast_memeq(in, "\x5d\x00\x00", 3)) {
     return LZMA_METHOD;
   } else if (in_len > 3 &&
              fast_memeq(in, "\x28\xb5\x2f\xfd", 4)) {
@@ -271,12 +284,78 @@ tor_compress_supports_method(compress_method_t method)
     case ZSTD_METHOD:
       return tor_zstd_method_supported();
     case NO_METHOD:
+      return 1;
     case UNKNOWN_METHOD:
     default:
       return 0;
   }
 }
 
+/**
+ * Return a bitmask of the supported compression types, where 1&lt;&lt;m is
+ * set in the bitmask if and only if compression with method <b>m</b> is
+ * supported.
+ */
+unsigned
+tor_compress_get_supported_method_bitmask(void)
+{
+  static unsigned supported = 0;
+  if (supported == 0) {
+    compress_method_t m;
+    for (m = NO_METHOD; m <= UNKNOWN_METHOD; ++m) {
+      if (tor_compress_supports_method(m)) {
+        supported |= (1u << m);
+      }
+    }
+  }
+  return supported;
+}
+
+/** Table of compression method names.  These should have an "x-" prefix,
+ * if they are not listed in the IANA content coding registry. */
+static const struct {
+  const char *name;
+  compress_method_t method;
+} compression_method_names[] = {
+  { "gzip", GZIP_METHOD },
+  { "deflate", ZLIB_METHOD },
+  // We call this "x-tor-lzma" rather than "x-lzma", because we impose a
+  // lower maximum memory usage on the decoding side.
+  { "x-tor-lzma", LZMA_METHOD },
+  { "x-zstd" , ZSTD_METHOD },
+  { "identity", NO_METHOD },
+
+  /* Later entries in this table are not canonical; these are recognized but
+   * not emitted. */
+  { "x-gzip", GZIP_METHOD },
+};
+
+/** Return the canonical string representation of the compression method
+ * <b>method</b>, or NULL if the method isn't recognized. */
+const char *
+compression_method_get_name(compress_method_t method)
+{
+  unsigned i;
+  for (i = 0; i < ARRAY_LENGTH(compression_method_names); ++i) {
+    if (method == compression_method_names[i].method)
+      return compression_method_names[i].name;
+  }
+  return NULL;
+}
+
+/** Return the compression method represented by the string <b>name</b>, or
+ * UNKNOWN_METHOD if the string isn't recognized. */
+compress_method_t
+compression_method_get_by_name(const char *name)
+{
+  unsigned i;
+  for (i = 0; i < ARRAY_LENGTH(compression_method_names); ++i) {
+    if (!strcmp(compression_method_names[i].name, name))
+      return compression_method_names[i].method;
+  }
+  return UNKNOWN_METHOD;
+}
+
 /** Return a string representation of the version of the library providing the
  * compression method given in <b>method</b>. Returns NULL if <b>method</b> is
  * unknown or unsupported. */
@@ -324,7 +403,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();
 }
@@ -384,11 +464,15 @@ tor_compress_new(int compress, compress_method_t method,
       state->u.zstd_state = zstd_state;
       break;
     }
-    case NO_METHOD:
+    case NO_METHOD: {
+      break;
+    }
     case UNKNOWN_METHOD:
       goto err;
   }
 
+  atomic_counter_add(&total_compress_allocation,
+                     sizeof(tor_compress_state_t));
   return state;
 
  err:
@@ -430,6 +514,8 @@ tor_compress_process(tor_compress_state_t *state,
                                        out, out_len, in, in_len,
                                        finish);
     case NO_METHOD:
+      return tor_cnone_compress_process(out, out_len, in, in_len,
+                                        finish);
     case UNKNOWN_METHOD:
       goto err;
   }
@@ -457,10 +543,13 @@ tor_compress_free(tor_compress_state_t *state)
       tor_zstd_compress_free(state->u.zstd_state);
       break;
     case NO_METHOD:
+      break;
     case UNKNOWN_METHOD:
       break;
   }
 
+  atomic_counter_sub(&total_compress_allocation,
+                     sizeof(tor_compress_state_t));
   tor_free(state);
 }
 
@@ -470,27 +559,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();