fix(oracle): validate input words before tool lookup

tools/hw-oracle/hw_oracle.py

@@ -1,6 +1,7 @@
 import argparse
 import hashlib
 import os
+import operator
 import shlex
 import shutil
 import struct
@@ -89,7 +90,16 @@ def assemble(body, work, wave64, *, ieee, denorm32, denorm16, dx10_clamp, round3
 
 
 def run(body, rows, extra=b"", wave64=False, coarse=False, *, ieee, denorm32, denorm16, dx10_clamp, round32, round16, fp16_overflow, lds=DEFAULT_LDS):
-    rows = [tuple(r) for r in rows]
+    validated = []
+    for index, row in enumerate(rows, 1):
+        try:
+            values = tuple(operator.index(value) for value in row)
+        except TypeError as error:
+            raise ValueError(f"row {index} needs 4 unsigned 32-bit integers") from error
+        if len(values) != 4 or any(value < 0 or value > 0xffffffff for value in values):
+            raise ValueError(f"row {index} needs 4 unsigned 32-bit integers")
+        validated.append(values)
+    rows = validated
     wave = 64 if wave64 else 32
     padded = rows + [(0, 0, 0, 0)] * (-len(rows) % wave)
     env = dict(os.environ)
@@ -111,6 +121,13 @@ def run(body, rows, extra=b"", wave64=False, coarse=False, *, ieee, denorm32, de
     return out[:len(rows)]
 
 
+def parse_word(value):
+    try:
+        return int(value, 0)
+    except ValueError:
+        return int(value, 10)
+
+
 def main():
     parser = argparse.ArgumentParser(description="Run an RDNA kernel body on the local AMD GPU, one input row per lane.")
     parser.add_argument("body", type=Path, help="assembly inserted into template.s")
@@ -128,7 +145,7 @@ def main():
     parser.add_argument("--extra", type=Path, help="bytes appended after the rows")
     parser.add_argument("--outs", type=int, choices=range(1, 17), default=16, metavar="N", help="print v10..v(10+N-1)")
     args = parser.parse_args()
-    rows = [tuple(int(v, 0) for v in line.split()) for line in args.rows.read_text().splitlines() if line.strip()]
+    rows = [tuple(parse_word(v) for v in line.split()) for line in args.rows.read_text().splitlines() if line.strip()]
     if any(len(r) != 4 for r in rows):
         sys.exit("every row needs 4 values")
     extra = args.extra.read_bytes() if args.extra else b""

tools/hw-oracle/test_hw_oracle.py

@@ -76,6 +76,45 @@ class OracleCacheTests(unittest.TestCase):
                 self.assertEqual([file for file in cache.rglob("*") if file.is_file()], [binary])
 
 
+class InputRowsTests(unittest.TestCase):
+    def test_invalid_rows_fail_before_assembler_or_gpu_lookup(self):
+        for row in ((1, 2, 3), (1, 2, 3, 4, 5), (-1, 0, 0, 0),
+                    (0x100000000, 0, 0, 0), (1.5, 0, 0, 0), ("1", 0, 0, 0)):
+            with self.subTest(row=row), patch.object(hw_oracle, "assemble", side_effect=AssertionError("assembler called")) as assemble:
+                with self.assertRaisesRegex(ValueError, "row 1"):
+                    hw_oracle.run("s_nop 0", [row], **MODES)
+                assemble.assert_not_called()
+
+    def test_unsigned_bits_and_indexable_values_are_preserved(self):
+        class Indexable:
+            def __index__(self):
+                return 0xffffffff
+
+        def dispatch(command, **kwargs):
+            data = Path(command[3]).read_bytes()
+            self.assertEqual(data[:16], hw_oracle.struct.pack("<4I", 0, 0xffffffff, 0x80000000, 1))
+            self.assertEqual(data[16:32 * 16], bytes(31 * 16))
+            self.assertEqual(data[32 * 16:], b"extra")
+            Path(command[4]).write_bytes(bytes(32 * 64))
+
+        with patch.object(hw_oracle, "assemble", return_value="kernel.co"), \
+                patch.object(hw_oracle, "oracle", return_value="oracle"), \
+                patch.object(hw_oracle.subprocess, "run", side_effect=dispatch):
+            results = hw_oracle.run("s_nop 0", [(0, Indexable(), 0x80000000, True)], b"extra", **MODES)
+        self.assertEqual(results, [tuple([0] * 16)])
+
+    def test_cli_accepts_zero_padded_decimal_rows(self):
+        with tempfile.TemporaryDirectory() as tmp:
+            body, rows = Path(tmp) / "body.s", Path(tmp) / "rows.txt"
+            body.write_text("s_nop 0")
+            rows.write_text("0001 08 0xFFFFFFFF 0b10\n")
+            argv = ["hw_oracle.py", str(body), str(rows)]
+            argv += [item for name, value in MODES.items() for item in ["--" + name.replace("_", "-"), str(value)]]
+            with patch("sys.argv", argv), patch.object(hw_oracle, "run", return_value=[]) as run:
+                hw_oracle.main()
+            self.assertEqual(run.call_args.args[1], [(1, 8, 0xffffffff, 2)])
+
+
 class FloatModeTests(unittest.TestCase):
     def test_python_requires_each_mode(self):
         for name in MODES: