From 8f97d6f2cdfccefba5457ae3d561e9ce0109da3f Mon Sep 17 00:00:00 2001
From: Akira Ajisaka <aajisaka@apache.org>
Date: Wed, 24 Oct 2018 13:49:28 +0900
Subject: [PATCH] HADOOP-15822. zstd compressor can fail with a small output
 buffer. Contributed by Jason Lowe.

---
 dev-support/docker/Dockerfile                         |  1 +
 .../io/compress/zstd/ZStandardDecompressor.java       |  4 ++--
 .../hadoop/io/compress/zstd/ZStandardCompressor.c     | 11 +++++++----
 .../hadoop/io/compress/zstd/ZStandardDecompressor.c   |  1 +
 .../zstd/TestZStandardCompressorDecompressor.java     | 10 ++++------
 5 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/dev-support/docker/Dockerfile b/dev-support/docker/Dockerfile
index 7679500f09..67140c9d7a 100644
--- a/dev-support/docker/Dockerfile
+++ b/dev-support/docker/Dockerfile
@@ -63,6 +63,7 @@ RUN apt-get -q update \
         libsnappy-dev \
         libssl-dev \
         libtool \
+        libzstd1-dev \
         locales \
         make \
         pinentry-curses \
diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/io/compress/zstd/ZStandardDecompressor.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/io/compress/zstd/ZStandardDecompressor.java
index 73d73e1c47..bc9d29cb4f 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/io/compress/zstd/ZStandardDecompressor.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/io/compress/zstd/ZStandardDecompressor.java
@@ -262,8 +262,8 @@ int inflateDirect(ByteBuffer src, ByteBuffer dst) throws IOException {
 
     int originalPosition = dst.position();
     int n = inflateBytesDirect(
-        src, src.position(), src.remaining(), dst, dst.position(),
-        dst.remaining()
+        src, src.position(), src.limit(), dst, dst.position(),
+        dst.limit()
     );
     dst.position(originalPosition + n);
     if (bytesInCompressedBuffer > 0) {
diff --git a/hadoop-common-project/hadoop-common/src/main/native/src/org/apache/hadoop/io/compress/zstd/ZStandardCompressor.c b/hadoop-common-project/hadoop-common/src/main/native/src/org/apache/hadoop/io/compress/zstd/ZStandardCompressor.c
index 055683ad51..289554b4cf 100644
--- a/hadoop-common-project/hadoop-common/src/main/native/src/org/apache/hadoop/io/compress/zstd/ZStandardCompressor.c
+++ b/hadoop-common-project/hadoop-common/src/main/native/src/org/apache/hadoop/io/compress/zstd/ZStandardCompressor.c
@@ -195,10 +195,13 @@ JNIEXPORT jint Java_org_apache_hadoop_io_compress_zstd_ZStandardCompressor_defla
     ZSTD_inBuffer input = { uncompressed_bytes, uncompressed_direct_buf_len, uncompressed_direct_buf_off };
     ZSTD_outBuffer output = { compressed_bytes, compressed_direct_buf_len, 0 };
 
-    size_t size = dlsym_ZSTD_compressStream(stream, &output, &input);
-    if (dlsym_ZSTD_isError(size)) {
-        THROW(env, "java/lang/InternalError", dlsym_ZSTD_getErrorName(size));
-        return (jint) 0;
+    size_t size;
+    if (uncompressed_direct_buf_len != 0) {
+        size = dlsym_ZSTD_compressStream(stream, &output, &input);
+        if (dlsym_ZSTD_isError(size)) {
+            THROW(env, "java/lang/InternalError", dlsym_ZSTD_getErrorName(size));
+            return (jint) 0;
+        }
     }
     if (finish && input.pos == input.size) {
         // end the stream, flush and  write the frame epilogue
diff --git a/hadoop-common-project/hadoop-common/src/main/native/src/org/apache/hadoop/io/compress/zstd/ZStandardDecompressor.c b/hadoop-common-project/hadoop-common/src/main/native/src/org/apache/hadoop/io/compress/zstd/ZStandardDecompressor.c
index b97dd1797a..e28359b526 100644
--- a/hadoop-common-project/hadoop-common/src/main/native/src/org/apache/hadoop/io/compress/zstd/ZStandardDecompressor.c
+++ b/hadoop-common-project/hadoop-common/src/main/native/src/org/apache/hadoop/io/compress/zstd/ZStandardDecompressor.c
@@ -178,6 +178,7 @@ JNIEXPORT jint JNICALL Java_org_apache_hadoop_io_compress_zstd_ZStandardDecompre
         return (jint) 0;
     }
     uncompressed_bytes = ((char*) uncompressed_bytes) + uncompressed_direct_buf_off;
+    uncompressed_direct_buf_len -= uncompressed_direct_buf_off;
 
     ZSTD_inBuffer input = { compressed_bytes, compressed_direct_buf_len, compressed_direct_buf_off };
     ZSTD_outBuffer output = { uncompressed_bytes, uncompressed_direct_buf_len, 0 };
diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/io/compress/zstd/TestZStandardCompressorDecompressor.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/io/compress/zstd/TestZStandardCompressorDecompressor.java
index 04def24de7..dcfb7e9e32 100644
--- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/io/compress/zstd/TestZStandardCompressorDecompressor.java
+++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/io/compress/zstd/TestZStandardCompressorDecompressor.java
@@ -414,13 +414,11 @@ private void compressDecompressLoop(int rawDataSize) throws IOException {
     outBuf.clear();
     while (!decompressor.finished()) {
       decompressor.decompress(inBuf, outBuf);
-      if (outBuf.remaining() == 0) {
-        outBuf.flip();
-        while (outBuf.remaining() > 0) {
-          assertEquals(expected.get(), outBuf.get());
-        }
-        outBuf.clear();
+      outBuf.flip();
+      while (outBuf.remaining() > 0) {
+        assertEquals(expected.get(), outBuf.get());
       }
+      outBuf.clear();
     }
     outBuf.flip();
     while (outBuf.remaining() > 0) {