From 152b55879c6ac7723a49d51ea48bc20dd05959f9 Mon Sep 17 00:00:00 2001
From: senhuang42 <senhuang96@fb.com>
Date: Wed, 2 Dec 2020 15:27:39 -0500
Subject: [PATCH] Add unit tests to seekable

---
 contrib/seekable_format/tests/.gitignore      |  1 +
 contrib/seekable_format/tests/Makefile        | 38 +++++++++++
 .../seekable_format/tests/seekable_tests.c    | 63 +++++++++++++++++++
 contrib/seekable_format/zstdseek_decompress.c |  6 +-
 4 files changed, 106 insertions(+), 2 deletions(-)
 create mode 100644 contrib/seekable_format/tests/.gitignore
 create mode 100644 contrib/seekable_format/tests/Makefile
 create mode 100644 contrib/seekable_format/tests/seekable_tests.c

diff --git a/contrib/seekable_format/tests/.gitignore b/contrib/seekable_format/tests/.gitignore
new file mode 100644
index 00000000..f831eaf3
--- /dev/null
+++ b/contrib/seekable_format/tests/.gitignore
@@ -0,0 +1 @@
+seekable_tests
diff --git a/contrib/seekable_format/tests/Makefile b/contrib/seekable_format/tests/Makefile
new file mode 100644
index 00000000..e72469d4
--- /dev/null
+++ b/contrib/seekable_format/tests/Makefile
@@ -0,0 +1,38 @@
+# ################################################################
+# Copyright (c) 2017-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under both the BSD-style license (found in the
+# LICENSE file in the root directory of this source tree) and the GPLv2 (found
+# in the COPYING file in the root directory of this source tree).
+# ################################################################
+
+# This Makefile presumes libzstd is built, using `make` in / or /lib/
+
+ZSTDLIB_PATH = ../../../lib
+ZSTDLIB_NAME = libzstd.a
+ZSTDLIB = $(ZSTDLIB_PATH)/$(ZSTDLIB_NAME)
+
+CPPFLAGS += -I../ -I../../../lib -I../../../lib/common
+
+CFLAGS ?= -O3
+CFLAGS += -g
+
+SEEKABLE_OBJS = ../zstdseek_compress.c ../zstdseek_decompress.c $(ZSTDLIB)
+
+.PHONY: default clean test
+
+default: test
+
+test: seekable_tests
+
+$(ZSTDLIB):
+	make -C $(ZSTDLIB_PATH) $(ZSTDLIB_NAME)
+
+seekable_tests : seekable_tests.c $(SEEKABLE_OBJS)
+	$(CC) $(CPPFLAGS) $(CFLAGS) $^ $(LDFLAGS) -o $@
+
+clean:
+	@rm -f core *.o tmp* result* *.zst \
+		seekable_tests
+	@echo Cleaning completed
diff --git a/contrib/seekable_format/tests/seekable_tests.c b/contrib/seekable_format/tests/seekable_tests.c
new file mode 100644
index 00000000..f2556b51
--- /dev/null
+++ b/contrib/seekable_format/tests/seekable_tests.c
@@ -0,0 +1,63 @@
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+
+#include "zstd_seekable.h"
+
+/* Basic unit tests for zstd seekable format */
+int main(int argc, const char** argv)
+{
+    unsigned testNb = 1;
+    printf("Beginning zstd seekable format tests...\n");
+    printf("Test %u - check that seekable decompress does not hang: ", testNb++);
+    {   /* Github issue #2335 */
+        const size_t compressed_size = 17;
+        const uint8_t compressed_data[17] = {
+            '^',
+            '*',
+            'M',
+            '\x18',
+            '\t',
+            '\x00',
+            '\x00',
+            '\x00',
+            '\x00',
+            '\x00',
+            '\x00',
+            '\x00',
+            ';',
+            (uint8_t)('\xb1'),
+            (uint8_t)('\xea'),
+            (uint8_t)('\x92'),
+            (uint8_t)('\x8f'),
+        };
+        const size_t uncompressed_size = 32;
+        uint8_t uncompressed_data[32];
+
+        ZSTD_seekable* stream = ZSTD_seekable_create();
+        size_t status = ZSTD_seekable_initBuff(stream, compressed_data, compressed_size);
+        if (ZSTD_isError(status)) {
+            ZSTD_seekable_free(stream);
+            goto _test_error;
+        }
+
+        const size_t offset = 2;
+        /* Should return an error, but not hang */
+        status = ZSTD_seekable_decompress(stream, uncompressed_data, uncompressed_size, offset);
+        if (!ZSTD_isError(status)) {
+            ZSTD_seekable_free(stream);
+            goto _test_error;
+        }
+
+        ZSTD_seekable_free(stream);
+    }
+    printf("Success!\n");
+
+    /* TODO: Add more tests */
+    printf("Finished tests\n");
+    return 0;
+
+_test_error:
+    printf("test failed! Exiting..\n");
+    return 1;
+}
diff --git a/contrib/seekable_format/zstdseek_decompress.c b/contrib/seekable_format/zstdseek_decompress.c
index eb9fd71d..cc5c8595 100644
--- a/contrib/seekable_format/zstdseek_decompress.c
+++ b/contrib/seekable_format/zstdseek_decompress.c
@@ -420,8 +420,10 @@ size_t ZSTD_seekable_decompress(ZSTD_seekable* zs, void* dst, size_t len, unsign
                              outTmp.pos - prevOutPos);
             }
             forwardProgress = outTmp.pos - prevOutPos;
-            if (!forwardProgress && noOutputProgressCount++ > ZSTD_SEEKABLE_NO_OUTPUT_PROGRESS_MAX) {
-                return ERROR(seekableIO);
+            if (forwardProgress == 0) {
+                if (noOutputProgressCount++ > ZSTD_SEEKABLE_NO_OUTPUT_PROGRESS_MAX) {
+                    return ERROR(seekableIO);
+                }
             } else {
                 noOutputProgressCount = 0;
             }