From caab29ec889a0771191b58714c306439b2415d91 Mon Sep 17 00:00:00 2001
From: Ashutosh Gupta <ashutosh.gupta@st.niituniversity.in>
Date: Tue, 28 Dec 2021 18:14:38 +0530
Subject: [PATCH] HDFS-14099. Unknown frame descriptor when decompressing
 multiple frames (#3836)

Co-authored-by: xuzq <xuzengqiang@kuaishou.com>
Co-authored-by: Ashutosh Gupta <ashugpt@amazon.com>
Signed-off-by: Akira Ajisaka <aajisaka@apache.org>
---
 .../compress/zstd/ZStandardDecompressor.java  | 13 ++++
 .../TestZStandardCompressorDecompressor.java  | 59 +++++++++++++++++++
 2 files changed, 72 insertions(+)

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 bc9d29cb4f..adf2fe629f 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
@@ -113,6 +113,12 @@ private void setInputFromSavedData() {
     compressedDirectBuf.put(
         userBuf, userBufOff, bytesInCompressedBuffer);
 
+    // Set the finished to false when compressedDirectBuf still
+    // contains some bytes.
+    if (compressedDirectBuf.position() > 0 && finished) {
+      finished = false;
+    }
+
     userBufOff += bytesInCompressedBuffer;
     userBufferBytesToConsume -= bytesInCompressedBuffer;
   }
@@ -186,6 +192,13 @@ public int decompress(byte[] b, int off, int len)
         0,
         directBufferSize
     );
+
+    // Set the finished to false when compressedDirectBuf still
+    // contains some bytes.
+    if (remaining > 0 && finished) {
+      finished = false;
+    }
+
     uncompressedDirectBuf.limit(n);
 
     // Get at most 'len' bytes
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 f12226d897..d4c0718220 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
@@ -231,6 +231,65 @@ public void testCompressorDecompressorLogicWithCompressionStreams()
     }
   }
 
+  /**
+   * Verify decompressor logic with some finish operation in compress.
+   */
+  @Test
+  public void testCompressorDecompressorWithFinish() throws Exception {
+    DataOutputStream deflateOut = null;
+    DataInputStream inflateIn = null;
+    int byteSize = 1024 * 100;
+    byte[] bytes = generate(byteSize);
+    int firstLength = 1024 * 30;
+
+    int bufferSize = IO_FILE_BUFFER_SIZE_DEFAULT;
+    try {
+      DataOutputBuffer compressedDataBuffer = new DataOutputBuffer();
+      CompressionOutputStream deflateFilter =
+              new CompressorStream(compressedDataBuffer, new ZStandardCompressor(),
+                      bufferSize);
+
+      deflateOut =
+              new DataOutputStream(new BufferedOutputStream(deflateFilter));
+
+      // Write some data and finish.
+      deflateOut.write(bytes, 0, firstLength);
+      deflateFilter.finish();
+      deflateOut.flush();
+
+      // ResetState then write some data and finish.
+      deflateFilter.resetState();
+      deflateOut.write(bytes, firstLength, firstLength);
+      deflateFilter.finish();
+      deflateOut.flush();
+
+      // ResetState then write some data and finish.
+      deflateFilter.resetState();
+      deflateOut.write(bytes, firstLength * 2, byteSize - firstLength * 2);
+      deflateFilter.finish();
+      deflateOut.flush();
+
+      DataInputBuffer deCompressedDataBuffer = new DataInputBuffer();
+      deCompressedDataBuffer.reset(compressedDataBuffer.getData(), 0,
+              compressedDataBuffer.getLength());
+
+      CompressionInputStream inflateFilter =
+              new DecompressorStream(deCompressedDataBuffer,
+                      new ZStandardDecompressor(bufferSize), bufferSize);
+
+      inflateIn = new DataInputStream(new BufferedInputStream(inflateFilter));
+
+      byte[] result = new byte[byteSize];
+      inflateIn.read(result);
+      assertArrayEquals(
+              "original array not equals compress/decompressed array", bytes,
+              result);
+    } finally {
+      IOUtils.closeStream(deflateOut);
+      IOUtils.closeStream(inflateIn);
+    }
+  }
+
   @Test
   public void testZStandardCompressDecompressInMultiThreads() throws Exception {
     MultithreadedTestUtil.TestContext ctx =