From 6a93c45cf488ec98eef2f742fe6e3b4036dcdb0e Mon Sep 17 00:00:00 2001
From: stringhandler <stringhandler@protonmail.com>
Date: Thu, 8 Oct 2026 15:39:07 +0200
Subject: [PATCH] add check for daglen exaggeration

When decoding a simplicity program, the dag length is read first. Currently, we
allocate memory according to that length without checking whether the program
is actually that long. This could be used to allocate more memory than needed.
Here we add a check that the length of the DAG on stream is within the size of
that length before allocating.
---
 .gitignore      |  1 +
 C/deserialize.c |  5 +++++
 C/test.c        | 31 +++++++++++++++++++++++++++++++
 3 files changed, 37 insertions(+)

diff --git a/.gitignore b/.gitignore
index fc4d1734f..e1690adef 100644
--- a/.gitignore
+++ b/.gitignore
@@ -40,3 +40,4 @@ CMakeFiles
 /*.c
 /*.h
 /*.inc
+.vscode
\ No newline at end of file
diff --git a/C/deserialize.c b/C/deserialize.c
index 182f23ba1..f6f236cdd 100644
--- a/C/deserialize.c
+++ b/C/deserialize.c
@@ -191,6 +191,11 @@ int_fast32_t simplicity_decodeMallocDag(dag_node** dag, simplicity_callback_deco
   if (dagLen <= 0) return dagLen;
   static_assert(DAG_LEN_MAX <= (uint32_t)INT32_MAX, "DAG_LEN_MAX exceeds supported parsing range.");
   if (DAG_LEN_MAX < (uint32_t)dagLen) return SIMPLICITY_ERR_DATA_OUT_OF_RANGE;
+  /* Every node's encoding is at least 2 bits long ('decodeNode' reads at least 2 bits before any call to 'decodeJet'),
+   * so a 'stream' too short to hold 'dagLen' nodes is bound to fail with 'SIMPLICITY_ERR_BITSTREAM_EOF'.
+   * Fail before allocating, so that the size of the allocation is proportional to the length of the 'stream'.
+   */
+  if (stream->len < ((size_t)dagLen * 2 + stream->offset + CHAR_BIT - 1) / CHAR_BIT) return SIMPLICITY_ERR_BITSTREAM_EOF;
   static_assert(DAG_LEN_MAX <= SIZE_MAX / sizeof(dag_node), "dag array too large.");
   static_assert(1 <= DAG_LEN_MAX, "DAG_LEN_MAX is zero.");
   static_assert(DAG_LEN_MAX - 1 <= UINT32_MAX, "dag array index does not fit in uint32_t.");
diff --git a/C/test.c b/C/test.c
index b8b6eb418..becde90d6 100644
--- a/C/test.c
+++ b/C/test.c
@@ -71,6 +71,36 @@ static void test_decodeUptoMaxInt(void) {
   }
 }
 
+static void test_decodeMallocDag_length(void) {
+  printf("Test decodeMallocDag length prefix\n");
+  static const struct {
+    const char* name;
+    unsigned char buf[5];
+    size_t len;
+    int_fast32_t expected;
+  } cases[] =
+  /* A length prefix of 8,000,000 in a 5 byte stream must be rejected before allocating the dag. */
+  { { "8000000 nodes in 5 bytes", { 0xf0, 0x6e, 0x84, 0x80, 0x00 }, 5, SIMPLICITY_ERR_BITSTREAM_EOF }
+  /* A length prefix of 4 needs at least 6 + 2*4 bits, i.e. 2 bytes. */
+  , { "4 nodes in 1 byte", { 0xc0 }, 1, SIMPLICITY_ERR_BITSTREAM_EOF }
+  /* With 2 bytes the stream is long enough to start decoding, which then finds the first node referencing a non-existent child. */
+  , { "4 nodes in 2 bytes", { 0xc0, 0x00 }, 2, SIMPLICITY_ERR_DATA_OUT_OF_RANGE }
+  };
+
+  for (size_t i = 0; i < sizeof(cases)/sizeof(cases[0]); ++i) {
+    dag_node* dag;
+    bitstream stream = initializeBitstream(cases[i].buf, cases[i].len);
+    int_fast32_t result = simplicity_decodeMallocDag(&dag, simplicity_elements_decodeJet, NULL, &stream);
+    if (cases[i].expected == result && NULL == dag) {
+      successes++;
+    } else {
+      failures++;
+      printf("Unexpected result for %s.  Expected %d and received %d\n", cases[i].name, (int)cases[i].expected, (int)result);
+    }
+    simplicity_free(dag);
+  }
+}
+
 static void test_hashBlock(void) {
   printf("Test hashBlock\n");
   dag_node* dag;
@@ -745,6 +775,7 @@ int main(int argc, char **argv) {
     timing_flag = 0;
   }
   test_decodeUptoMaxInt();
+  test_decodeMallocDag_length();
   test_hashBlock();
   test_occursCheck();