fix: harden against confirmed review findings

SECURITY.md

@@ -8,6 +8,13 @@ repeated-line patterns) causing a crash, hang, or unbounded memory
 growth, which the fuzz suite (`test/fuzz_runner.mojo`) specifically
 targets.
 
+To bound worst-case memory, the Myers diff caps its exact search at
+`_MAX_EDIT_DISTANCE` steps (`src/diff/diff.mojo`). Inputs whose edit
+distance stays within the cap diff exactly (byte-compatible with
+difflib); more-dissimilar inputs fall back to a coarse
+prefix/suffix-preserving result instead of allocating the O(D^2) trace,
+so two large unrelated files can no longer exhaust memory.
+
 If you find an input that crashes, hangs, or otherwise misbehaves in a
 way that looks security-relevant, please report it via a
 [GitHub issue](https://github.com/conorbronsdon/mojo-diff/issues),

pixi.toml

@@ -12,4 +12,4 @@ fuzz = "mojo run -I src test/fuzz_runner.mojo test/data/code_edit.a.txt"
 demo = "mojo run -I src examples/diff_files.mojo test/data/code_edit.a.txt test/data/code_edit.b.txt"
 
 [dependencies]
-mojo = ">=1.0.0b3,<2"
+mojo = ">=1.0.0b3.dev0,<2"

src/diff/diff.mojo

@@ -10,12 +10,31 @@ with `difflib.unified_diff` — same hunk grouping and `@@` range math).
 `difflib` fed by `readlines()`. A final line without a newline is emitted without
 one — matching difflib, which does not print a "\\ No newline at end of file"
 marker (that is a GNU-diff feature, not a difflib one).
+
+Memory bound: the Myers backtrack stores one `V` snapshot per edit-distance
+step `d`, so the trace is `O(d^2)` in the edit distance. On adversarial,
+fully-dissimilar inputs `d` approaches `len(a) + len(b)`, which would allocate
+tens of GB for two large unrelated files (a DoS). To keep memory bounded,
+`matching_blocks` caps the exact search at `max_edit_distance` steps
+(default `_MAX_EDIT_DISTANCE`). Past the cap it stops the quadratic search and
+returns a *coarse* result: the shared leading/trailing lines are preserved as
+equal blocks and the differing middle is reported as one block-level change.
+This keeps output byte-compatible with `difflib` for every input whose edit
+distance is within the cap (all realistic edits), and degrades gracefully —
+never OOMs — beyond it.
 """
 
 from diff.model import OpCode, Match
 
 comptime _NL = UInt8(0x0A)
 
+# Cap on the Myers edit-distance search. The backtrack trace holds one V
+# snapshot per step, so peak trace memory is ~`(_MAX_EDIT_DISTANCE + 1)^2`
+# ints (~170 MB at 4096). Diffs whose edit distance stays under this run
+# exactly (difflib-identical); more-dissimilar inputs fall back to a coarse
+# prefix/suffix-preserving result instead of allocating unbounded memory.
+comptime _MAX_EDIT_DISTANCE = 4096
+
 
 def _imax(a: Int, b: Int) -> Int:
     return a if a > b else b
@@ -49,12 +68,21 @@ def splitlines_keepends(text: String) -> List[String]:
     return lines^
 
 
-def matching_blocks(a: List[String], b: List[String]) -> List[Match]:
+def matching_blocks(
+    a: List[String],
+    b: List[String],
+    max_edit_distance: Int = _MAX_EDIT_DISTANCE,
+) -> List[Match]:
     """Matching blocks of `a` vs `b` via Myers O(ND), difflib-shaped.
 
     Returns maximal matching runs in increasing order, terminated by the
     `Match(len(a), len(b), 0)` sentinel (as `SequenceMatcher.get_matching_blocks`
     does).
+
+    The exact Myers search is capped at `max_edit_distance` steps to bound
+    memory (see module docstring). If the edit distance would exceed the cap,
+    a *coarse* result is returned: shared leading/trailing lines stay as equal
+    blocks and the differing middle is left as a single block-level change.
     """
     var n = len(a)
     var m = len(b)
@@ -63,6 +91,19 @@ def matching_blocks(a: List[String], b: List[String]) -> List[Match]:
         blocks.append(Match(n, m, 0))  # sentinel only; no matches possible
         return blocks^
 
+    # Common leading/trailing lines. Only used to build the coarse fallback
+    # below; the exact path recovers these naturally via Myers snakes.
+    var prefix = 0
+    while prefix < n and prefix < m and a[prefix] == b[prefix]:
+        prefix += 1
+    var suffix = 0
+    while (
+        suffix < (n - prefix)
+        and suffix < (m - prefix)
+        and a[n - 1 - suffix] == b[m - 1 - suffix]
+    ):
+        suffix += 1
+
     # Greedy forward pass, recording a trimmed V snapshot per edit distance d.
     var maxd = n + m
     var offset = maxd
@@ -71,9 +112,15 @@ def matching_blocks(a: List[String], b: List[String]) -> List[Match]:
         v.append(0)
     var trace = List[List[Int]]()
     var found = False
+    var capped = False
     var d_final = 0
     var d = 0
     while d <= maxd:
+        if d > max_edit_distance:
+            # Exact search would need `(d+1)^2` trace ints; bail out to the
+            # bounded coarse fallback instead of growing memory unbounded.
+            capped = True
+            break
         var snap = List[Int]()
         for k in range(-d, d + 1):
             snap.append(v[k + offset])
@@ -99,6 +146,18 @@ def matching_blocks(a: List[String], b: List[String]) -> List[Match]:
             break
         d += 1
 
+    if capped:
+        # Coarse, bounded fallback: keep the shared prefix/suffix as equal
+        # blocks and leave the differing middle as one block-level change
+        # (`get_opcodes` renders it as a single replace/insert/delete). This
+        # never allocates the quadratic trace, so memory stays bounded.
+        if prefix > 0:
+            blocks.append(Match(0, 0, prefix))
+        if suffix > 0:
+            blocks.append(Match(n - suffix, m - suffix, suffix))
+        blocks.append(Match(n, m, 0))  # sentinel
+        return blocks^
+
     # Backtrack through the trace, collecting matched (x, y) points in
     # decreasing order.
     var px = List[Int]()
@@ -149,10 +208,17 @@ def matching_blocks(a: List[String], b: List[String]) -> List[Match]:
     return blocks^
 
 
-def get_opcodes(a: List[String], b: List[String]) -> List[OpCode]:
+def get_opcodes(
+    a: List[String],
+    b: List[String],
+    max_edit_distance: Int = _MAX_EDIT_DISTANCE,
+) -> List[OpCode]:
     """Edit opcodes turning `a` into `b`, mirroring `SequenceMatcher.get_opcodes`.
+
+    `max_edit_distance` caps the exact search (see `matching_blocks`); beyond it
+    the differing region collapses to a single coarse opcode.
     """
-    var blocks = matching_blocks(a, b)
+    var blocks = matching_blocks(a, b, max_edit_distance)
     var ops = List[OpCode]()
     var i = 0
     var j = 0

test/test_diff.mojo

@@ -155,6 +155,41 @@ def test_opcode_roundtrip() raises:
         assert_equal(rebuilt[i], b[i], "roundtrip line " + String(i))
 
 
+# --- edit-distance cap / DoS guard -----------------------------------------
+
+
+def test_cap_coarse_fallback() raises:
+    # Shared prefix (2 lines) + a differing middle that has an internal match
+    # ("MID") + shared suffix (2 lines). Under a tiny cap the exact search bails
+    # out and the middle collapses to a single coarse `replace`, while the
+    # shared prefix/suffix are still reported as `equal`.
+    var a = _lines("p0", "p1", "a0", "MID", "a1", "s0", "s1")
+    var b = _lines("p0", "p1", "b0", "MID", "b1", "s0", "s1")
+    var coarse = get_opcodes(a, b, 2)
+    var exp = List[OpCode]()
+    exp.append(OpCode(String("equal"), 0, 2, 0, 2))
+    exp.append(OpCode(String("replace"), 2, 5, 2, 5))
+    exp.append(OpCode(String("equal"), 5, 7, 5, 7))
+    _assert_ops(coarse, exp, "cap coarse fallback")
+    # With headroom the same diff stays fine-grained (the internal "MID" match
+    # is preserved), i.e. strictly more opcodes than the coarse form — proving
+    # the cap does not degrade any diff whose edit distance is within it.
+    var exact = get_opcodes(a, b, 100)
+    assert_true(len(exact) > len(coarse), "sub-cap diff stays fine-grained")
+
+
+def test_cap_bounds_dissimilar() raises:
+    # Fully dissimilar inputs (no shared prefix/suffix): the capped search must
+    # return a single coarse `replace` and terminate, instead of allocating the
+    # O(D^2) Myers trace that would OOM on large unrelated files.
+    var a = _lines("x0", "x1", "x2", "x3", "x4", "x5")
+    var b = _lines("y0", "y1", "y2", "y3", "y4", "y5")
+    var ops = get_opcodes(a, b, 2)
+    var exp = List[OpCode]()
+    exp.append(OpCode(String("replace"), 0, 6, 0, 6))
+    _assert_ops(ops, exp, "cap dissimilar single replace")
+
+
 # --- ratio -----------------------------------------------------------------