From baec365fd4112fa7886e816a0cf9c743271d81f1 Mon Sep 17 00:00:00 2001
From: Soma Aishwarya <soma@digiscrypt.com>
Date: Sat, 18 Jul 2026 14:38:49 +0530
Subject: [PATCH] validate trace packet header fields against payload size

---
 util/HalideTraceUtils.cpp | 29 +++++++++++++++++++++++++++++
 util/HalideTraceViz.cpp   | 23 +++++++++++++++++++++++
 2 files changed, 52 insertions(+)

diff --git a/util/HalideTraceUtils.cpp b/util/HalideTraceUtils.cpp
index dfeb34799bc8..fe982d737546 100644
--- a/util/HalideTraceUtils.cpp
+++ b/util/HalideTraceUtils.cpp
@@ -16,6 +16,10 @@ bool Packet::read_from_filedesc(FILE *fdesc) {
     if (!Packet::read(this, header_size, fdesc)) {
         return false;
     }
+    if (size < header_size) {
+        fprintf(stderr, "Malformed trace packet: size (%d) smaller than header\n", (int)size);
+        return false;
+    }
     size_t payload_size = size - header_size;
     if (payload_size > sizeof(payload)) {
         fprintf(stderr, "Payload larger than %d bytes in trace stream (%d)\n", (int)sizeof(payload), (int)payload_size);
@@ -26,6 +30,31 @@ bool Packet::read_from_filedesc(FILE *fdesc) {
         fprintf(stderr, "Unexpected EOF mid-packet");
         return false;
     }
+    // dimensions and type come straight from the untrusted header, and
+    // coordinates()/value()/func()/trace_tag() use them to index into payload.
+    // Reject any packet whose declared layout doesn't fit the bytes we read so
+    // those accessors stay in bounds.
+    if (dimensions < 0) {
+        fprintf(stderr, "Malformed trace packet: negative dimensions (%d)\n", (int)dimensions);
+        return false;
+    }
+    const size_t fixed_bytes = (size_t)dimensions * sizeof(int32_t) + value_bytes();
+    if (fixed_bytes > payload_size) {
+        fprintf(stderr, "Malformed trace packet: coordinates/value exceed payload\n");
+        return false;
+    }
+    // func() and trace_tag() are NUL-terminated strings following the value;
+    // both terminators must lie within the payload or the walks run past it.
+    int terminators = 0;
+    for (size_t i = fixed_bytes; i < payload_size; i++) {
+        if (payload[i] == 0 && ++terminators == 2) {
+            break;
+        }
+    }
+    if (terminators < 2) {
+        fprintf(stderr, "Malformed trace packet: func/trace_tag not terminated within payload\n");
+        return false;
+    }
     return true;
 }
 
diff --git a/util/HalideTraceViz.cpp b/util/HalideTraceViz.cpp
index e0f24687c174..f53f2be6db28 100644
--- a/util/HalideTraceViz.cpp
+++ b/util/HalideTraceViz.cpp
@@ -173,11 +173,34 @@ struct PacketAndPayload : public halide_trace_packet_t {
             return false;  // EOF
         }
 
+        if (this->size < header_size) {
+            fail() << "Malformed trace packet: size " << this->size << " smaller than header";
+        }
         const size_t payload_size = this->size - header_size;
         if (payload_size > sizeof(this->payload) || !read_or_die(this->payload, payload_size)) {
             // Shouldn't ever get EOF here
             fail() << "Unable to read packet payload of size " << payload_size;
         }
+        // dimensions and type come straight from the untrusted header, and
+        // coordinates()/value()/func()/trace_tag() use them to index into
+        // payload. Reject any packet whose declared layout doesn't fit the
+        // bytes we read so those accessors stay in bounds.
+        if (this->dimensions < 0) {
+            fail() << "Malformed trace packet: negative dimensions " << this->dimensions;
+        }
+        const size_t fixed_bytes = (size_t)this->dimensions * sizeof(int32_t) + this->value_bytes();
+        if (fixed_bytes > payload_size) {
+            fail() << "Malformed trace packet: coordinates/value exceed payload";
+        }
+        int terminators = 0;
+        for (size_t i = fixed_bytes; i < payload_size; i++) {
+            if (this->payload[i] == 0 && ++terminators == 2) {
+                break;
+            }
+        }
+        if (terminators < 2) {
+            fail() << "Malformed trace packet: func/trace_tag not terminated within payload";
+        }
         return true;
     }
 };