Fix zlib.

PiperOrigin-RevId: 367282082

Author: yimingz-a

tensorflow_serving/util/net_http/compression/BUILD

@@ -35,5 +35,7 @@ cc_test(
     deps = [
         ":gzip_zlib",
         "//tensorflow_serving/core/test_util:test_main",
+        "@com_google_absl//absl/strings",
+        "@com_google_googletest//:gtest",
     ],
 )

tensorflow_serving/util/net_http/compression/gzip_zlib.cc

@@ -695,7 +695,9 @@ int ZLib::UncompressAtMost(Bytef *dest, uLongf *destLen, const Bytef *source,
 // mode, we also check the gzip footer to make sure we pass the gzip
 // consistency checks.  We RETURN true iff both types of checks pass.
 bool ZLib::UncompressChunkDone() {
-  assert(!first_chunk_ && uncomp_init_);
+  if (first_chunk_ || !uncomp_init_) {
+    return false;
+  }
   // Make sure we're at the end-of-compressed-data point.  This means
   // if we call inflate with Z_FINISH we won't consume any input or
   // write any output
@@ -771,7 +773,7 @@ int ZLib::Uncompress(Bytef *dest, uLongf *destLen, const Bytef *source,
 
 // read uncompress length from gzip footer
 uLongf ZLib::GzipUncompressedLength(const Bytef *source, uLong len) {
-  assert(len > 4);
+  if (len <= 4) return 0;  // malformed data.
   return (static_cast<uLongf>(source[len - 1]) << 24) +
          (static_cast<uLongf>(source[len - 2]) << 16) +
          (static_cast<uLongf>(source[len - 3]) << 8) +

tensorflow_serving/util/net_http/compression/gzip_zlib.h

@@ -173,7 +173,8 @@ class ZLib {
   int Uncompress(Bytef *dest, uLongf *destLen, const Bytef *source,
                  uLong sourceLen);
 
-  // Get the uncompressed size from the gzip footer.
+  // Get the uncompressed size from the gzip footer. Returns 0 if source is too
+  // short (len < 5).
   uLongf GzipUncompressedLength(const Bytef *source, uLong len);
 
   // Special helper function to help uncompress gzipped documents:

tensorflow_serving/util/net_http/compression/gzip_zlib_test.cc

@@ -20,7 +20,9 @@ limitations under the License.
 #include <string>
 #include <vector>
 
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
+#include "absl/strings/str_cat.h"
 #include "absl/strings/string_view.h"
 
 namespace tensorflow {
@@ -602,6 +604,98 @@ TEST(ZLibTest, BytewiseRead) {
   ASSERT_EQ(truncated_output, text);
 }
 
+TEST(ZLibTest, TruncatedData) {
+  const int kBufferLen = 64;
+  std::string uncompressed = "Hello, World!";
+  std::string compressed(
+      "\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03\xf3\x48\xcd\xc9\xc9"
+      "\xd7\x51\x08\xcf\x2f\xca\x49\x51\x04\x00\xd0\xc3\x4a\xec\x0d"
+      "\x00\x00\x00",
+      33);
+
+  // Verify that "compressed" contains valid gzip data.
+  {
+    ZLib zlib;
+    // zlib.SetGzipHeaderMode();
+    char uncompbuf[kBufferLen];
+    bzero(uncompbuf, kBufferLen);
+    uLongf uncomplen = kBufferLen;
+    int err = zlib.Uncompress(
+        reinterpret_cast<Bytef*>(uncompbuf), &uncomplen,
+        reinterpret_cast<const Bytef*>(compressed.c_str()), compressed.size());
+    ASSERT_EQ(err, Z_OK);
+    ASSERT_EQ(uncompressed, absl::string_view(uncompbuf, uncomplen));
+  }
+
+  // Test truncated data with ZLib::Uncompress().
+  for (int len = compressed.size() - 1; len > 0; len--) {
+    SCOPED_TRACE(absl::StrCat("Decompressing first ", len, " out of ",
+                              compressed.size(), " bytes"));
+    ZLib zlib;
+    // zlib.SetGzipHeaderMode();
+    char uncompbuf[kBufferLen];
+    bzero(uncompbuf, kBufferLen);
+    uLongf uncomplen = kBufferLen;
+    int err = zlib.Uncompress(
+        reinterpret_cast<Bytef*>(uncompbuf), &uncomplen,
+        reinterpret_cast<const Bytef*>(compressed.c_str()), len);
+    ASSERT_NE(err, Z_OK);
+  }
+
+  // Test truncated data with ZLib::UncompressAtMost() and
+  // ZLib::UncompressDone().
+  for (int len = compressed.size() - 1; len > 0; len--) {
+    SCOPED_TRACE(absl::StrCat("Decompressing first ", len, " out of ",
+                              compressed.size(), " bytes"));
+    ZLib zlib;
+    // zlib.SetGzipHeaderMode();
+    char uncompbuf[kBufferLen];
+    bzero(uncompbuf, kBufferLen);
+    uLongf uncomplen = kBufferLen;
+    uLongf complen = len;
+    int err = zlib.UncompressAtMost(
+        reinterpret_cast<Bytef*>(uncompbuf), &uncomplen,
+        reinterpret_cast<const Bytef*>(compressed.c_str()), &complen);
+    ASSERT_EQ(err, Z_OK);
+    ASSERT_EQ(complen, 0);
+    if (uncomplen > 0) {
+      EXPECT_THAT(
+          uncompressed,
+          testing::StartsWith(std::string(uncompbuf).substr(0, uncomplen)));
+    }
+    ASSERT_FALSE(zlib.UncompressChunkDone());
+  }
+}
+
+TEST(ZLibTest, GzipUncompressedLength) {
+  ZLib zlib;
+
+  // "Hello, World!", compressed.
+  std::string hello_world(
+      "\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03\xf3\x48\xcd\xc9\xc9"
+      "\xd7\x51\x08\xcf\x2f\xca\x49\x51\x04\x00\xd0\xc3\x4a\xec\x0d"
+      "\x00\x00\x00",
+      33);
+  EXPECT_EQ(13, zlib.GzipUncompressedLength(
+                    reinterpret_cast<const Bytef*>(hello_world.c_str()),
+                    hello_world.size()));
+
+  // Empty string, "", compressed.
+  std::string empty(
+      "\x1f\x8b\x08\x00\x00\x00\x00\x00\x00\x03\x03\x00\x00\x00\x00"
+      "\x00\x00\x00\x00\x00",
+      20);
+  EXPECT_EQ(0,
+            zlib.GzipUncompressedLength(
+                reinterpret_cast<const Bytef*>(empty.c_str()), empty.size()));
+
+  std::string bad_data("\x01\x01\x01\x01", 4);
+  for (int len = 0; len <= bad_data.size(); len++) {
+    EXPECT_EQ(0, zlib.GzipUncompressedLength(
+                     reinterpret_cast<const Bytef*>(bad_data.c_str()), len));
+  }
+}
+
 }  // namespace
 }  // namespace net_http
 }  // namespace serving