From 89e842b1475017da104eb3c5d7684138313132b6 Mon Sep 17 00:00:00 2001
From: am0o0 <77095239+am0o0@users.noreply.github.com>
Date: Tue, 3 Sep 2024 09:12:13 +0200
Subject: [PATCH] finilize tests for zlib

---
 .../CWE/CWE-409/DecompressionBombs.ql         |  2 +-
 .../Security/CWE/CWE-409/ZlibUncompress.qll   |  2 +-
 .../Security/CWE/CWE-409/zlibTest.cpp         | 24 ++++++++++++++++++-
 3 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/cpp/ql/src/experimental/query-tests/Security/CWE/CWE-409/DecompressionBombs.ql b/cpp/ql/src/experimental/query-tests/Security/CWE/CWE-409/DecompressionBombs.ql
index 05f39c47e13..659e92d8498 100644
--- a/cpp/ql/src/experimental/query-tests/Security/CWE/CWE-409/DecompressionBombs.ql
+++ b/cpp/ql/src/experimental/query-tests/Security/CWE/CWE-409/DecompressionBombs.ql
@@ -22,7 +22,7 @@ module DecompressionTaintConfig implements DataFlow::ConfigSig {
 
   predicate isSink(DataFlow::Node sink) {
     exists(FunctionCall fc, DecompressionFunction f | fc.getTarget() = f |
-      fc.getArgument(f.getArchiveParameterIndex()) = sink.asExpr() 
+      fc.getArgument(f.getArchiveParameterIndex()) = [sink.asExpr(), sink.asIndirectExpr()]
     )
   }
 
diff --git a/cpp/ql/src/experimental/query-tests/Security/CWE/CWE-409/ZlibUncompress.qll b/cpp/ql/src/experimental/query-tests/Security/CWE/CWE-409/ZlibUncompress.qll
index f69c71ec01c..60943495d45 100644
--- a/cpp/ql/src/experimental/query-tests/Security/CWE/CWE-409/ZlibUncompress.qll
+++ b/cpp/ql/src/experimental/query-tests/Security/CWE/CWE-409/ZlibUncompress.qll
@@ -13,5 +13,5 @@ import DecompressionBomb
 class UncompressFunction extends DecompressionFunction {
   UncompressFunction() { this.hasGlobalName(["uncompress", "uncompress2"]) }
 
-  override int getArchiveParameterIndex() { result = 0 }
+  override int getArchiveParameterIndex() { result = 2 }
 }
diff --git a/cpp/ql/test/experimental/query-tests/Security/CWE/CWE-409/zlibTest.cpp b/cpp/ql/test/experimental/query-tests/Security/CWE/CWE-409/zlibTest.cpp
index 07805f6c83f..9e23944e7ce 100644
--- a/cpp/ql/test/experimental/query-tests/Security/CWE/CWE-409/zlibTest.cpp
+++ b/cpp/ql/test/experimental/query-tests/Security/CWE/CWE-409/zlibTest.cpp
@@ -1,4 +1,3 @@
-
 #define Z_NULL  0
 #  define FAR
 typedef unsigned char Byte;
@@ -145,9 +144,32 @@ int UnsafeGzgets(char *fileName) {
     return 0;
 }
 
+typedef unsigned long uLong;
+typedef long unsigned int size_t;
+typedef uLong uLongf;
+typedef unsigned char Bytef;
+#define Z_OK            0
+
+int uncompress(Bytef *dest, uLongf *destLen,
+               const Bytef *source, uLong sourceLen) { return 0; }
+
+bool InflateString(const unsigned char *input, const unsigned char *output, size_t output_length) {
+    uLong source_length;
+    source_length = (uLong) 500;
+    uLong destination_length;
+    destination_length = (uLong) output_length;
+
+    int result = uncompress((Bytef *) output, &destination_length,
+                            (Bytef *) input, source_length);
+
+    return result == Z_OK;
+}
+
 int main(int argc, char **argv) {
     UnsafeGzfread(argv[2]);
     UnsafeGzgets(argv[2]);
     UnsafeInflate(argv[2]);
     UnsafeGzread(argv[2]);
+    const unsigned char *output;
+    InflateString(reinterpret_cast<const unsigned char *>(argv[1]), output, 1024 * 1024 * 1024);
 }