Size the surface_distance Dijkstra heap to its real push bound (#3722)

.claude/sweep-security-state.csv

@@ -41,7 +41,7 @@ resample,2026-04-28,1295,HIGH,1,,"HIGH (fixed #1295): resample() did not bound o
 sieve,2026-04-28,1296,HIGH,1,,"HIGH (fixed #1296): sieve() on numpy and cupy backends had no memory guard. _label_connected allocates parent (int32, 4B/px), rank (int32, 4B/px, reused as root_to_id), region_map_flat (int32, 4B/px), plus a float64 result copy (8B/px) ~ 20 B/pixel of working memory before any check. The dask paths (_sieve_dask line 343 and _sieve_dask_cupy line 366) already raised MemoryError via _available_memory_bytes() at 28 B/pixel budget, but the public sieve() API at line 489 dispatched np.ndarray inputs straight into _sieve_numpy with no guard, and _sieve_cupy at line 308 transferred to host via data.get() then called _sieve_numpy, inheriting the gap. A 50000x50000 numpy raster requested ~50 GB silently. Fixed by extracting _check_memory(rows, cols) and _check_gpu_memory(rows, cols) helpers (mirrors cost_distance #1262 / mahalanobis #1288 / multispectral #1291 / kde #1287 pattern) at 28 B/pixel host budget plus 16 B/pixel GPU round-trip budget at 50% of available memory threshold. _check_memory wired into _sieve_numpy at the top before the float64 copy. _check_gpu_memory wired into _sieve_cupy before data.get(); it also calls _check_memory so the host budget still applies. Consolidated _available_memory_bytes definition (was duplicated). All 47 tests pass including 2 new memory-guard tests for the numpy backend (_sieve_numpy direct call + public sieve() API). No other findings: Cat 2 int32 indexing in _label_connected docstring acknowledges <2.1B pixel limit; the new memory guard rejects rasters that large before the int32 issue can trigger so this is a documentation/clarity follow-up rather than an exploitable bug. Cat 3 NaN handled via valid mask; Cat 4 no CUDA kernels; Cat 5 only /proc/meminfo read; Cat 6 _validate_raster called at line 478."
 sky_view_factor,2026-04-28,1299,HIGH,1,,"Unbounded numpy/cupy allocation; fixed via _check_memory and _check_gpu_memory guards (16 B/pixel, 50% threshold). Dask paths skip the guard."
 slope,2026-04-28,,,,,"Clean. slope() validates input via _validate_raster (line 383) and _validate_boundary (line 389). Cat 1: planar _cpu/_run_cupy allocate output matching input shape; geodesic paths build (3,H,W) float64 stacked array but are gated by _check_geodesic_memory(rows, cols) at line 410 (already fixed under geodesic audit, PR #1285). Cat 2: no int32 flat-index math; all loops 2D with range(). Cat 3: NaN propagates through arctan in planar kernels; geodesic delegates to _local_frame_project_and_fit which has explicit NaN guards and degenerate det check. Cat 4: _run_gpu (line 146) uses combined bounds+stencil guard 'i-di>=0 and i+di<H and j-dj>=0 and j+dj<W'; geodesic GPU kernels imported from geodesic.py and audited there; _geodesic_cuda_dims uses 16x16 blocks to avoid register spill. Cat 5: no file I/O. Cat 6: all backends cast explicitly to float32 (planar) or float64 (geodesic); lat/lon cast to float64 in _extract_latlon_coords."
-surface_distance,2026-04-28,1303,HIGH,1,,Fixed in PR #1305: added _check_memory and _check_gpu_memory guards to _surface_distance_numpy (line ~233) and _surface_distance_cupy (line ~448) before O(H*W) heap+output allocations.  Dask paths inherit via per-chunk numpy call.  Other categories clean.
+surface_distance,2026-08-16,3722,CRITICAL,1;2,,"Re-audit of #1303. #1305 guards still present and correct (_check_memory at _surface_distance_numpy L464, _check_gpu_memory at _surface_distance_cupy L547); dask inherits per chunk. NEW CRITICAL #3722 / PR #3727: _dijkstra L214 and _dijkstra_geodesic L277 sized the lazy-deletion heap height*width, but pushes are bounded by height*width*(n_neighbors+1); dense source masks overflow and _heap_push writes OOB. Repro: 80x80 checkerboard sources over random elevation aborts the interpreter (double free), IndexError under NUMBA_BOUNDSCHECK=1; hits numpy planar, numpy geodesic and dask+numpy tile paths. Fixed by matching cost_distance sizing; _BYTES_PER_PIXEL 80->272 so the #1305 guard still models the allocation. Cat 4 clean: _sd_relax_kernel bounds guard present, compute-sanitizer memcheck 0 errors over 9 adversarial shapes x 3 public funcs on this CUDA host. Cat 3 clean (all NaN guards are inside-the-box isfinite tests, not rejection form). Cat 5 n/a (only literal /proc/meminfo read). Cat 6 clean (_validate_raster on raster and elevation). LOW, not fixed: NaN max_distance silently behaves like inf; negative max_distance yields a negative map_overlap depth on the dask bounded branch. Related un-reproduced site: pathfinding.py L304 has the same height*width heap sizing (single-source A*, could not overflow it) - noted in #3722, not filed separately."
 terrain,2026-05-03,1443,MEDIUM,1;3,,"Re-audit 2026-05-03. MEDIUM Cat 1 + Cat 3 fixed in PR #1444: _terrain_numpy and _terrain_cupy now call _check_memory / _check_gpu_memory (24 B/pixel scratch budget, 50% threshold); generate_terrain rejects non-finite or non-positive lacunarity / persistence. Dask path worley_norm_range pre-pass dask.persist remains documented but not exploitable (caller-controlled). No remaining findings."
 viewshed,2026-04-22,1229,HIGH,1,,"HIGH (fixed #1229): _viewshed_cpu allocated ~500 bytes/pixel of working memory (event_list 3*H*W*7*8 bytes + status_values/status_struct/idle + visibility_grid + lexsort temporary) with no guard. A 20000x20000 raster tried to allocate ~200 GB. Fixed by adding peak-memory guard mirroring the _viewshed_dask pattern (_available_memory_bytes() check, raises MemoryError with max_distance= hint). No other HIGH findings: dask path already guarded, _validate_raster is called, distance-sweep uses dtype=float64, _calc_dist_n_grad guards zero distance."
 visibility,2026-04-28,,,,,"Clean. line_of_sight (line 190) and cumulative_viewshed (line 259) call _validate_raster; visibility_frequency delegates. Cat 1: cumulative_viewshed allocates int32 accumulator (4 B/px) but delegates per-observer to viewshed() which has 500 B/px memory guard at viewshed.py:1523-1531; viewshed will fail first on oversize rasters. _bresenham_line (line 35) and _los_kernel (lines 112-143) bounded by transect length (<=W+H+1). Cat 2: int64 throughout, no int32 overflow path. Cat 3: divisions in _los_kernel guarded (D==0 in _fresnel_radius_1 line 87, distance[i]==0 continue line 133, total_dist>0 check line 123); NaN elevation at observer cell would taint los_height but is a correctness not DoS concern. Cat 4: no CUDA kernels. Cat 5: no file I/O. Cat 6: elevations cast to float64 in _extract_transect line 79."

xrspatial/surface_distance.py

@@ -83,13 +83,14 @@ class cupy:  # type: ignore[no-redef]
 #   src_row (int64)             8
 #   src_col (int64)             8
 #   visited (int8)              1
-#   h_keys (float64)            8
-#   h_rows (int64)              8
-#   h_cols (int64)              8
+#   heap (float64 + 2x int64)   24 per slot
 #   output (float32)            4
 #   direction-mode temps        ~16
-# Total ~80 bytes/pixel.  A 50000x50000 raster needs ~200 GB.
-_BYTES_PER_PIXEL = 80
+# The Dijkstra heap holds up to (n_neighbors + 1) slots per pixel, so the
+# worst case is 8-connectivity: 9 * 24 = 216 bytes/pixel of heap on top of
+# the ~53 bytes of everything else.  Round up to 272.  A 50000x50000 raster
+# needs ~680 GB.
+_BYTES_PER_PIXEL = 272
 
 # CuPy backend skips the explicit binary heap (parallel relaxation instead)
 # but still allocates dist, alloc, srow, scol, src cast, elev cast, mask,
@@ -211,7 +212,13 @@ def _dijkstra(elev_data, height, width, max_distance,
     """
     n_neighbors = len(dy)
 
-    max_heap = height * width
+    # Heap arrays.  This is a lazy-deletion min-heap: a pixel is enqueued
+    # again every time its tentative distance improves, and it can improve
+    # once per settled neighbour.  Total pushes are therefore bounded by the
+    # directed edge count (n_neighbors per pixel) plus one seed push per
+    # pixel.  The old height*width sizing underflows that bound and let
+    # _heap_push write past the end of the arrays, corrupting memory.
+    max_heap = height * width * (n_neighbors + 1)
     h_keys = np.empty(max_heap, dtype=np.float64)
     h_rows = np.empty(max_heap, dtype=np.int64)
     h_cols = np.empty(max_heap, dtype=np.int64)
@@ -274,7 +281,9 @@ def _dijkstra_geodesic(elev_data, height, width, max_distance,
     """
     n_neighbors = len(dy)
 
-    max_heap = height * width
+    # See _dijkstra for why the heap is sized to the push bound rather than
+    # to the pixel count.
+    max_heap = height * width * (n_neighbors + 1)
     h_keys = np.empty(max_heap, dtype=np.float64)
     h_rows = np.empty(max_heap, dtype=np.int64)
     h_cols = np.empty(max_heap, dtype=np.int64)

xrspatial/tests/test_surface_distance.py

@@ -755,10 +755,10 @@ def test_dask_path_bounded_per_chunk(self):
 
         from unittest.mock import patch
 
-        # 200x200 total (~6.4 MB at 80 B/pixel) chunked at 20x20
-        # (~32 KB per chunk).  Mock available memory to 1 MB: the full
+        # 200x200 total (~10.9 MB at 272 B/pixel) chunked at 20x20
+        # (~109 KB per chunk).  Mock available memory to 1 MB: the full
         # array would exceed 50% of that, but each 20x20 chunk needs
-        # only ~32 KB so per-chunk allocation passes.
+        # only ~109 KB so per-chunk allocation passes.
         source = np.zeros((200, 200), dtype=np.float64)
         source[100, 100] = 1.0
         elev = np.zeros((200, 200), dtype=np.float64)
@@ -792,3 +792,153 @@ def test_error_message_mentions_grid_size(self):
                 surface_distance(raster, elevation)
             with pytest.raises(MemoryError, match="dask"):
                 surface_distance(raster, elevation)
+
+
+# ---------------------------------------------------------------------------
+# Regression — Dijkstra heap sizing (issue #3722)
+# ---------------------------------------------------------------------------
+
+
+def _checkerboard_inputs(h, w, backend='numpy', chunks=(3, 3), latlon=False):
+    """Dense checkerboard sources over rough random elevation.
+
+    Half the pixels are sources, which drives the lazy-deletion heap past
+    one slot per pixel.  With the old ``max_heap = height * width`` sizing
+    this walks off the end of the heap arrays.
+    """
+    rng = np.random.default_rng(0)
+    elev = rng.random((h, w)) * 5000.0
+    rr, cc = np.meshgrid(np.arange(h), np.arange(w), indexing='ij')
+    source = ((rr + cc) % 2 == 0).astype(np.float64)
+
+    raster = _make_raster(source, backend=backend, chunks=chunks)
+    elevation = _make_raster(elev, backend=backend, chunks=chunks)
+    if latlon:
+        coords = {'y': np.linspace(10.0, 10.5, h),
+                  'x': np.linspace(20.0, 20.5, w)}
+        raster = raster.assign_coords(coords)
+        elevation = elevation.assign_coords(coords)
+    return raster, elevation, source
+
+
+@pytest.mark.parametrize("connectivity", [4, 8])
+def test_dense_sources_do_not_overflow_heap(connectivity):
+    """A checkerboard source mask must not corrupt the Dijkstra heap.
+
+    Regression for #3722: the heap was sized ``height * width`` but a
+    lazy-deletion Dijkstra pushes a pixel once per improving relaxation,
+    so a dense source mask overflowed it and wrote out of bounds.
+    """
+    raster, elevation, source = _checkerboard_inputs(80, 80)
+
+    result = _compute(surface_distance(raster, elevation,
+                                       connectivity=connectivity))
+
+    assert np.isfinite(result).all()
+    # Every source pixel is at distance zero, every other pixel is reachable.
+    assert (result[source > 0] == 0).all()
+    assert (result[source == 0] > 0).all()
+
+
+def test_dense_sources_do_not_overflow_heap_geodesic():
+    """Same regression for the geodesic kernel (#3722)."""
+    raster, elevation, source = _checkerboard_inputs(80, 80, latlon=True)
+
+    result = _compute(surface_distance(raster, elevation, connectivity=8,
+                                       method='geodesic'))
+
+    assert np.isfinite(result).all()
+    assert (result[source > 0] == 0).all()
+
+
+def test_dense_sources_do_not_overflow_heap_dask():
+    """Same regression through the iterative dask tile path (#3722)."""
+    if da is None:
+        pytest.skip("dask not installed")
+
+    raster, elevation, source = _checkerboard_inputs(
+        80, 80, backend='dask+numpy', chunks=(40, 40))
+
+    with pytest.warns(UserWarning, match="iterative tile Dijkstra"):
+        lazy = surface_distance(raster, elevation, connectivity=8)
+    result = _compute(lazy)
+
+    assert np.isfinite(result).all()
+    assert (result[source > 0] == 0).all()
+
+
+def test_heap_bound_covers_worst_case_push_count():
+    """The allocated heap must cover every push the kernel can make.
+
+    Each pixel is pushed once when seeded and at most once per settled
+    neighbour, so the bound is ``height * width * (n_neighbors + 1)``.
+    Assert the kernel's peak heap usage stays inside that on the dense
+    input that used to overflow.
+    """
+    from xrspatial.cost_distance import _heap_push, _heap_pop
+    from xrspatial.utils import ngjit
+
+    @ngjit
+    def _peak_heap(elev_data, height, width, dy, dx, dd, dist):
+        n_neighbors = len(dy)
+        cap = height * width * (n_neighbors + 1)
+        h_keys = np.empty(cap, dtype=np.float64)
+        h_rows = np.empty(cap, dtype=np.int64)
+        h_cols = np.empty(cap, dtype=np.int64)
+        h_size = 0
+        peak = 0
+        visited = np.zeros((height, width), dtype=np.int8)
+
+        for r in range(height):
+            for c in range(width):
+                if dist[r, c] < np.inf:
+                    h_size = _heap_push(h_keys, h_rows, h_cols, h_size,
+                                        dist[r, c], r, c)
+                    if h_size > peak:
+                        peak = h_size
+
+        while h_size > 0:
+            cost_u, ur, uc, h_size = _heap_pop(h_keys, h_rows, h_cols,
+                                               h_size)
+            if visited[ur, uc]:
+                continue
+            visited[ur, uc] = 1
+            elev_u = elev_data[ur, uc]
+            for i in range(n_neighbors):
+                vr = ur + dy[i]
+                vc = uc + dx[i]
+                if vr < 0 or vr >= height or vc < 0 or vc >= width:
+                    continue
+                if visited[vr, vc]:
+                    continue
+                elev_v = elev_data[vr, vc]
+                if not np.isfinite(elev_v):
+                    continue
+                dz = elev_v - elev_u
+                new_cost = cost_u + np.sqrt(dd[i] * dd[i] + dz * dz)
+                if new_cost < dist[vr, vc]:
+                    dist[vr, vc] = new_cost
+                    h_size = _heap_push(h_keys, h_rows, h_cols, h_size,
+                                        new_cost, vr, vc)
+                    if h_size > peak:
+                        peak = h_size
+        return peak
+
+    h = w = 60
+    rng = np.random.default_rng(0)
+    elev = rng.random((h, w)) * 5000.0
+    rr, cc = np.meshgrid(np.arange(h), np.arange(w), indexing='ij')
+    dist = np.full((h, w), np.inf)
+    dist[(rr + cc) % 2 == 0] = 0.0
+
+    dy = np.array([-1, -1, -1, 0, 0, 1, 1, 1], dtype=np.int64)
+    dx = np.array([-1, 0, 1, -1, 1, -1, 0, 1], dtype=np.int64)
+    d = np.sqrt(2.0)
+    dd = np.array([d, 1.0, d, 1.0, 1.0, d, 1.0, d], dtype=np.float64)
+
+    peak = _peak_heap(elev, h, w, dy, dx, dd, dist)
+
+    # The old sizing (height * width) is not enough for this input ...
+    assert peak > h * w
+    # ... but the new bound is.
+    assert peak <= h * w * (len(dy) + 1)