fix(nft): address review findings before 0.10.0

.github/workflows/test.yml

@@ -15,7 +15,7 @@ jobs:
               sudo apt-get update && sudo apt install libjpeg-dev zlib1g-dev
         - name: Build and test WebARKitLib
           run: |
-              cd tests && mkdir build && cd build && cmake -DEMSCRIPTEN_COMP=0 .. && make && ./webarkit_test && ./webarkit_nft_test
+              cd tests && mkdir build && cd build && cmake -DEMSCRIPTEN_COMP=0 .. && make && ./webarkit_test && ./webarkit_nft_test && ./webarkit_nft_limit_test
         - name: Build WebARKitLib with Emscripten (Docker)
           run: |
              cd ..

WebARKit/WebARKitTrackers/WebARKitNFT/include/WebARKitTrackers/WebARKitNFT/markerDecompress.h

@@ -3,6 +3,11 @@
 
 #include <stdio.h>
 
+/* Largest decompressed .zft accepted by decompressMarkers(); larger archives are rejected. */
+#ifndef MARKER_DECOMPRESS_MAX_SIZE
+#define MARKER_DECOMPRESS_MAX_SIZE (128u * 1024u * 1024u)
+#endif
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -16,6 +21,12 @@ typedef struct
 
 char* nameConcat(const char *s1, const char *s2);
 FILE *openZFT( const char *filename, const char *ext);
+/*
+ * Unpack <src>.zft into <outTemp>.iset, <outTemp>.fset and <outTemp>.fset3.
+ * Returns 0 on success, -1 on a missing, malformed or oversized archive or a
+ * write error. On failure no output file is left behind, so <outTemp> should
+ * name new files: an existing marker set at that path is not preserved.
+ */
 int decompressMarkers(const char* src, const char* outTemp);
 int extractDataAndSave(const char* str, const char* name);
 

WebARKit/WebARKitTrackers/WebARKitNFT/include/WebARKitTrackers/WebARKitNFT/trackingMod.h

@@ -57,6 +57,10 @@
 #include <AR2/marker.h>
 #include <AR2/tracking.h>
 
+#if AR2_CAPABLE_ADAPTIVE_TEMPLATE
+#  error "WebARKitNFT trackingMod does not support AR2_CAPABLE_ADAPTIVE_TEMPLATE"
+#endif
+
 #define    AR2_TRACKING_6DOF                   1
 #define    AR2_TRACKING_HOMOGRAPHY             2
 

WebARKit/WebARKitTrackers/WebARKitNFT/markerDecompress.c

@@ -13,16 +13,23 @@
 #include <AR/ar.h>
 #include <zlib.h>
 
+#if MARKER_DECOMPRESS_MAX_SIZE < 1
+#  error "MARKER_DECOMPRESS_MAX_SIZE must be at least 1"
+#endif
+
 static const size_t inflate_chunk = 4*1024*1024;
+static const size_t inflate_max = MARKER_DECOMPRESS_MAX_SIZE;
 
 /*
  * Inflate a whole zlib stream. On success *outLen is the decompressed size and
  * the buffer has an extra NUL after it, so it can be searched as a string.
+ * Streams that expand past inflate_max are rejected, so a small archive with
+ * a huge expansion ratio cannot exhaust memory.
  */
 static char *inflateAll(const unsigned char *in, size_t inLen, size_t *outLen)
 {
     z_stream strm;
-    size_t cap = inflate_chunk;
+    size_t cap = inflate_chunk < inflate_max ? inflate_chunk : inflate_max;
     char *out = malloc(cap + 1);
     int ret;
 
@@ -37,13 +44,29 @@ static char *inflateAll(const unsigned char *in, size_t inLen, size_t *outLen)
 
     do {
         if (strm.total_out == cap) {
-            char *bigger = realloc(out, cap * 2 + 1);
+            size_t newCap;
+            char *bigger;
+            if (cap >= inflate_max) {
+                // At the limit: let zlib finish the stream (empty final block,
+                // trailer) with the one spare byte after the buffer. Any output
+                // written there means the stream is larger than the limit.
+                strm.next_out = (Bytef *)(out + cap);
+                strm.avail_out = 1;
+                ret = inflate(&strm, Z_NO_FLUSH);
+                if (strm.total_out > cap) {
+                    ARLOGe("Error: .zft data expands past %zu bytes\n", inflate_max);
+                    ret = Z_MEM_ERROR;
+                }
+                continue;
+            }
+            newCap = cap > inflate_max / 2 ? inflate_max : cap * 2;
+            bigger = realloc(out, newCap + 1);
             if (bigger == NULL) {
                 ret = Z_MEM_ERROR;
                 break;
             }
             out = bigger;
-            cap *= 2;
+            cap = newCap;
         }
         strm.next_out = (Bytef *)(out + strm.total_out);
         strm.avail_out = (uInt)(cap - strm.total_out);

WebARKit/WebARKitTrackers/WebARKitNFT/trackingMod2d.c

@@ -45,6 +45,7 @@
 #include <AR2/template.h>
 #include <AR2/searchPoint.h>
 #include <AR2/tracking.h>
+#include <WebARKitTrackers/WebARKitNFT/trackingMod.h>
 
 #if AR2_CAPABLE_ADAPTIVE_TEMPLATE
 int ar2Tracking2dSubMod ( AR2HandleT *handle, AR2SurfaceSetT *surfaceSet, AR2TemplateCandidateT *candidate,

WebARKit/WebARKitTrackers/WebARKitNFT/trackingSub.c

@@ -210,7 +210,7 @@ static void *trackingInitMain( THREAD_HANDLE_T *threadHandle )
         return (NULL);
     }
     trackingInitHandle = (TrackingInitHandle *)threadGetArg(threadHandle);
-    if (!threadHandle) {
+    if (!trackingInitHandle) {
         ARLOGe("Error starting tracking thread: empty trackingInitHandle.\n");
         return (NULL);
     }

tests/CMakeLists.txt

@@ -37,6 +37,20 @@ add_executable(webarkit_nft_test webarkit_nft_test.cc)
 target_compile_definitions(webarkit_nft_test PRIVATE WEBARKIT_NFT_THREADS)
 target_link_libraries(webarkit_nft_test WebARKitNFT GTest::gtest_main)
 
+# markerDecompress built with a 1 MB limit, to test the limit handling cheaply.
+find_package(ZLIB REQUIRED)
+add_executable(webarkit_nft_limit_test
+  webarkit_nft_limit_test.cc
+  ../WebARKit/WebARKitTrackers/WebARKitNFT/markerDecompress.c
+  ../lib/SRC/ARUtil/log.c
+)
+target_compile_definitions(webarkit_nft_limit_test PRIVATE MARKER_DECOMPRESS_MAX_SIZE=1048576)
+target_include_directories(webarkit_nft_limit_test PRIVATE
+  ../WebARKit/WebARKitTrackers/WebARKitNFT/include
+  ../include
+)
+target_link_libraries(webarkit_nft_limit_test ZLIB::ZLIB GTest::gtest_main)
+
 add_executable(
   webarkit_test
   webarkit_test.cc
@@ -86,3 +100,4 @@ target_link_libraries(
 include(GoogleTest)
 gtest_discover_tests(webarkit_test)
 gtest_discover_tests(webarkit_nft_test)
+gtest_discover_tests(webarkit_nft_limit_test)

tests/webarkit_nft_limit_test.cc

@@ -0,0 +1,83 @@
+// markerDecompress compiled with MARKER_DECOMPRESS_MAX_SIZE = 1 MB (see CMakeLists.txt).
+#include <gtest/gtest.h>
+
+#include <WebARKitTrackers/WebARKitNFT/markerDecompress.h>
+
+#include <zlib.h>
+
+#include <cstdio>
+#include <fstream>
+#include <string>
+#include <vector>
+
+static_assert(MARKER_DECOMPRESS_MAX_SIZE == 1024 * 1024, "this test expects a 1 MB limit");
+
+namespace {
+
+const std::string kHead = "{\"iset\":\"";
+const std::string kTail = "\",\"fset\":\"F\",\"fset3\":\"G\"}";
+
+// A marker whose decompressed size is exactly `size` bytes. The data is
+// flushed before the final (empty) block, so zlib still has a block and the
+// trailer to read when the output reaches `size`.
+std::string markerOfSize(size_t size) {
+  return kHead + std::string(size - kHead.size() - kTail.size(), 'A') + kTail;
+}
+
+void writeZftFlushedBeforeEnd(const std::string &basename, const std::string &content) {
+  z_stream strm = {};
+  ASSERT_EQ(deflateInit(&strm, Z_BEST_COMPRESSION), Z_OK);
+  std::string compressed;
+  std::vector<unsigned char> out(64 * 1024);
+  auto run = [&](int flush) {
+    do {
+      strm.next_out = out.data();
+      strm.avail_out = static_cast<uInt>(out.size());
+      deflate(&strm, flush);
+      compressed.append(reinterpret_cast<char *>(out.data()), out.size() - strm.avail_out);
+    } while (strm.avail_out == 0);
+  };
+  strm.next_in = reinterpret_cast<Bytef *>(const_cast<char *>(content.data()));
+  strm.avail_in = static_cast<uInt>(content.size());
+  run(Z_FULL_FLUSH);
+  run(Z_FINISH);
+  deflateEnd(&strm);
+  std::ofstream file(basename + ".zft", std::ios::binary);
+  file.write(compressed.data(), compressed.size());
+}
+
+bool fileExists(const std::string &path) {
+  std::ifstream in(path, std::ios::binary);
+  return in.good();
+}
+
+void removeOutputs(const std::string &out) {
+  std::remove((out + ".iset").c_str());
+  std::remove((out + ".fset").c_str());
+  std::remove((out + ".fset3").c_str());
+}
+
+} // namespace
+
+TEST(MarkerDecompressLimitTest, AcceptsArchiveExactlyAtTheLimit) {
+  writeZftFlushedBeforeEnd("nft_limit_exact", markerOfSize(MARKER_DECOMPRESS_MAX_SIZE));
+  EXPECT_EQ(decompressMarkers("nft_limit_exact", "nft_limit_exact_out"), 0);
+  EXPECT_TRUE(fileExists("nft_limit_exact_out.fset3"));
+  std::remove("nft_limit_exact.zft");
+  removeOutputs("nft_limit_exact_out");
+}
+
+TEST(MarkerDecompressLimitTest, RejectsArchiveOneByteOverTheLimit) {
+  writeZftFlushedBeforeEnd("nft_limit_over", markerOfSize(MARKER_DECOMPRESS_MAX_SIZE + 1));
+  EXPECT_EQ(decompressMarkers("nft_limit_over", "nft_limit_over_out"), -1);
+  EXPECT_FALSE(fileExists("nft_limit_over_out.iset"));
+  std::remove("nft_limit_over.zft");
+}
+
+TEST(MarkerDecompressLimitTest, EnforcesLimitsBelowTheInitialBuffer) {
+  // 2 MB is below the 4 MB initial buffer, but above this build's 1 MB limit.
+  writeZftFlushedBeforeEnd("nft_limit_small", markerOfSize(2 * 1024 * 1024));
+  EXPECT_EQ(decompressMarkers("nft_limit_small", "nft_limit_small_out"), -1);
+  EXPECT_FALSE(fileExists("nft_limit_small_out.iset"));
+  std::remove("nft_limit_small.zft");
+}

tests/webarkit_nft_test.cc

@@ -37,6 +37,11 @@ std::string readFile(const std::string &path) {
   return ss.str();
 }
 
+bool fileExists(const std::string &path) {
+  std::ifstream in(path, std::ios::binary);
+  return in.good();
+}
+
 } // namespace
 
 TEST(MarkerDecompressTest, ExtractsIsetFsetAndFset3) {
@@ -57,17 +62,44 @@ TEST(MarkerDecompressTest, MissingFileReturnsError) {
   EXPECT_EQ(decompressMarkers("does_not_exist", "nft_test_out"), -1);
 }
 
+TEST(MarkerDecompressTest, ArchiveExpandingPastTheLimitReturnsError) {
+  // Deflate MARKER_DECOMPRESS_MAX_SIZE + 1 MB of a valid-looking marker in
+  // chunks, so the test never holds the expanded data in memory.
+  z_stream strm = {};
+  ASSERT_EQ(deflateInit(&strm, Z_BEST_COMPRESSION), Z_OK);
+  std::string compressed;
+  std::vector<unsigned char> out(64 * 1024);
+  auto feed = [&](const std::string &data, int flush) {
+    strm.next_in = reinterpret_cast<Bytef *>(const_cast<char *>(data.data()));
+    strm.avail_in = static_cast<uInt>(data.size());
+    do {
+      strm.next_out = out.data();
+      strm.avail_out = static_cast<uInt>(out.size());
+      deflate(&strm, flush);
+      compressed.append(reinterpret_cast<char *>(out.data()), out.size() - strm.avail_out);
+    } while (strm.avail_out == 0);
+  };
+  const std::string chunk(1024 * 1024, 'A');
+  feed("{\"iset\":\"", Z_NO_FLUSH);
+  for (size_t i = 0; i < MARKER_DECOMPRESS_MAX_SIZE / chunk.size() + 1; i++) feed(chunk, Z_NO_FLUSH);
+  feed("\",\"fset\":\"F\",\"fset3\":\"G\"}", Z_FINISH);
+  deflateEnd(&strm);
+  {
+    std::ofstream file("nft_test_huge.zft", std::ios::binary);
+    file.write(compressed.data(), compressed.size());
+  }
+
+  EXPECT_EQ(decompressMarkers("nft_test_huge", "nft_test_huge_out"), -1);
+  EXPECT_FALSE(fileExists("nft_test_huge_out.iset"));
+  std::remove("nft_test_huge.zft");
+}
+
 TEST(MarkerDecompressTest, MalformedContentReturnsError) {
   writeZft("nft_test_bad", "{\"iset\":\"ISETDATA\"}");
   EXPECT_EQ(decompressMarkers("nft_test_bad", "nft_test_out"), -1);
   std::remove("nft_test_bad.zft");
 }
 
-static bool fileExists(const std::string &path) {
-  std::ifstream in(path, std::ios::binary);
-  return in.good();
-}
-
 TEST(MarkerDecompressTest, ExtractsMarkersLargerThanFourMegabytes) {
   const std::string iset(5 * 1024 * 1024, 0x41);
   writeZft("nft_test_big", "{\"iset\":\"" + iset + "\",\"fset\":\"F\",\"fset3\":\"G\"}");