From 08a5409dc6b369645a20321da4ba9bcc81b2790d Mon Sep 17 00:00:00 2001
From: eastagiletracker <310448263+eastagiletracker@users.noreply.github.com>
Date: Wed, 19 Aug 2026 04:36:44 +0000
Subject: [PATCH] Refuse PyUmiPacket merge when the LEN field would overflow

The merged LEN was written back into an eight-bit field without a range
check, so merging packets whose combined length exceeds 256 words wrapped
around: merging LEN=254 into LEN=5 produced LEN=4 and reported success,
leaving a packet that claimed five words while holding 261 bytes.

Refuse the merge instead, before the packet is resized, and add a test
covering the boundary case, the overflow case and an ordinary merge.
---
 examples/test_umi_packet_merge.py | 63 +++++++++++++++++++++++++++++++
 python/switchboard_pybind.cc      |  6 +++
 switchboard/cpp/umilib.h          |  5 +++
 3 files changed, 74 insertions(+)
 create mode 100644 examples/test_umi_packet_merge.py

diff --git a/examples/test_umi_packet_merge.py b/examples/test_umi_packet_merge.py
new file mode 100644
index 00000000..de4be216
--- /dev/null
+++ b/examples/test_umi_packet_merge.py
@@ -0,0 +1,63 @@
+#!/usr/bin/env python
+
+# Copyright (c) 2024 Zero ASIC Corporation
+# This code is licensed under Apache License 2.0 (see LICENSE for details)
+
+import numpy as np
+from switchboard import PyUmiPacket, UmiCmd, umi_pack, umi_len, umi_eom
+
+# SIZE=0 means one byte per word, so LEN+1 is also the number of data bytes
+SIZE = 0
+
+
+def make_packet(len_field, dstaddr, srcaddr, eom=0):
+    cmd = umi_pack(UmiCmd.UMI_REQ_WRITE, 0, SIZE, len_field, eom, 0)
+    data = np.arange(len_field + 1, dtype=np.uint8)
+    return PyUmiPacket(cmd, dstaddr, srcaddr, data)
+
+
+def test_merge_within_len_limit():
+    # 255 words + 1 word fills the LEN field exactly, so this merge is allowed
+    p = make_packet(254, 0, 0)
+    q = make_packet(0, 255, 255, eom=1)
+
+    assert p.merge(q)
+
+    assert umi_len(p.cmd) == 255
+    assert umi_eom(p.cmd) == 1
+    assert p.data.size == 256
+    assert p.data[255] == q.data[0]
+
+
+def test_merge_beyond_len_limit_rejected():
+    # 255 words + 6 words needs a LEN of 260, which does not fit in the
+    # eight-bit LEN field, so the merge must be refused rather than wrapping
+    p = make_packet(254, 0, 0)
+    q = make_packet(5, 255, 255, eom=1)
+
+    cmd_before = p.cmd
+
+    assert not p.merge(q)
+
+    # the rejected merge must leave the packet untouched
+    assert p.cmd == cmd_before
+    assert umi_len(p.cmd) == 254
+    assert umi_eom(p.cmd) == 0
+    assert p.data.size == 255
+
+
+def test_merge_ordinary_packets():
+    p = make_packet(3, 0, 0)
+    q = make_packet(1, 4, 4, eom=1)
+
+    assert p.merge(q)
+
+    assert umi_len(p.cmd) == 5
+    assert umi_eom(p.cmd) == 1
+    assert np.array_equal(p.data, np.array([0, 1, 2, 3, 0, 1], dtype=np.uint8))
+
+
+if __name__ == '__main__':
+    test_merge_within_len_limit()
+    test_merge_beyond_len_limit_rejected()
+    test_merge_ordinary_packets()
diff --git a/python/switchboard_pybind.cc b/python/switchboard_pybind.cc
index 8a7ac3c6..8a6771cf 100644
--- a/python/switchboard_pybind.cc
+++ b/python/switchboard_pybind.cc
@@ -199,6 +199,12 @@ struct PyUmiPacket {
         uint32_t len = umi_len(cmd);
         uint32_t nbytes = (len + 1) << size;
 
+        if ((len + umi_len(other.cmd) + 1) > UMI_MAX_LEN) {
+            // LEN is an eight-bit field, so the merged packet could not
+            // represent its own length; merging would silently wrap around
+            return false;
+        }
+
         if (other.dstaddr != (dstaddr + nbytes)) {
             // new dstaddr must be next sequentially
             return false;
diff --git a/switchboard/cpp/umilib.h b/switchboard/cpp/umilib.h
index c0bfbcfc..e3d682dd 100644
--- a/switchboard/cpp/umilib.h
+++ b/switchboard/cpp/umilib.h
@@ -130,6 +130,11 @@ static inline uint32_t umi_len(uint32_t cmd) {
     }
 }
 DECL_UMI_SETTER(len, 8, 8)
+
+// largest value that the eight-bit LEN field can hold, corresponding to
+// (UMI_MAX_LEN + 1) words of data in a single transaction
+#define UMI_MAX_LEN 255
+
 DECL_UMI_FIELD(atype, 8, 8)
 DECL_UMI_FIELD(qos, 16, 4)
 DECL_UMI_FIELD(prot, 20, 2)