From 2ae51ed5e20ee045251d9d6ddf12755bca7c9c8c Mon Sep 17 00:00:00 2001
From: Nick Mathewson <nickm@torproject.org>
Date: Fri, 14 Jul 2017 16:21:51 -0400
Subject: [PATCH] Fix zstd 1.3.0 trouble: Be more respectful of its state
 machine

In zstd 1.3.0, once you have called ZSTD_endStream and been told
that your putput buffer is full, it really doesn't want you to call
ZSTD_compressStream again.  ZSTD 1.2.0 didn't seem to mind about
this.

This patch fixes the issue by making sure never to call
ZSTD_endStream if there's any more data on the input buffer to
process, by flushing even when we're about to call "endStream", and
by never calling "compress" or "flush" after "endStream".
---
 changes/bug22927           |  6 ++++++
 src/common/compress_zstd.c | 29 +++++++++++++++++++++++------
 2 files changed, 29 insertions(+), 6 deletions(-)
 create mode 100644 changes/bug22927

diff --git a/changes/bug22927 b/changes/bug22927
new file mode 100644
index 000000000..6e68e6ff0
--- /dev/null
+++ b/changes/bug22927
@@ -0,0 +1,6 @@
+  o Minor bugfixes (compatibility, zstd):
+    - Write zstd epilogues correctly when the epilogue requires reallocation
+      of the output buffer, even with zstd 1.3.0. (Previously,
+      we worked on 1.2.0 and failed with 1.3.0).  Fixes bug 22927; bugfix on
+      0.3.1.1-alpha.
+
diff --git a/src/common/compress_zstd.c b/src/common/compress_zstd.c
index a136db48b..94974dec0 100644
--- a/src/common/compress_zstd.c
+++ b/src/common/compress_zstd.c
@@ -98,6 +98,8 @@ struct tor_zstd_compress_state_t {
 #endif // HAVE_ZSTD.
 
   int compress; /**< True if we are compressing; false if we are inflating */
+  int have_called_end; /**< True if we are compressing and we've called
+                        * ZSTD_endStream */
 
   /** Number of bytes read so far.  Used to detect compression bombs. */
   size_t input_so_far;
@@ -270,9 +272,16 @@ tor_zstd_compress_process(tor_zstd_compress_state_t *state,
   ZSTD_inBuffer input = { *in, *in_len, 0 };
   ZSTD_outBuffer output = { *out, *out_len, 0 };
 
+  if (BUG(finish == 0 && state->have_called_end)) {
+    finish = 1;
+  }
+
   if (state->compress) {
-    retval = ZSTD_compressStream(state->u.compress_stream,
-                                 &output, &input);
+    if (! state->have_called_end)
+      retval = ZSTD_compressStream(state->u.compress_stream,
+                                   &output, &input);
+    else
+      retval = 0;
   } else {
     retval = ZSTD_decompressStream(state->u.decompress_stream,
                                    &output, &input);
@@ -300,7 +309,7 @@ tor_zstd_compress_process(tor_zstd_compress_state_t *state,
     return TOR_COMPRESS_ERROR;
   }
 
-  if (state->compress && !finish) {
+  if (state->compress && !state->have_called_end) {
     retval = ZSTD_flushStream(state->u.compress_stream, &output);
 
     *out = (char *)output.dst + output.pos;
@@ -314,16 +323,24 @@ tor_zstd_compress_process(tor_zstd_compress_state_t *state,
 
     // ZSTD_flushStream returns 0 if the frame is done, or >0 if it
     // is incomplete.
-    if (retval > 0)
+    if (retval > 0) {
       return TOR_COMPRESS_BUFFER_FULL;
+    }
   }
 
   if (!finish) {
-    // We're not done with the input, so no need to flush.
+    // The caller says we're not done with the input, so no need to write an
+    // epilogue.
     return TOR_COMPRESS_OK;
   } else if (state->compress) {
-    retval = ZSTD_endStream(state->u.compress_stream, &output);
+    if (*in_len) {
+      // We say that we're not done with the input, so we can't write an
+      // epilogue.
+      return TOR_COMPRESS_OK;
+    }
 
+    retval = ZSTD_endStream(state->u.compress_stream, &output);
+    state->have_called_end = 1;
     *out = (char *)output.dst + output.pos;
     *out_len = output.size - output.pos;