Skip to content

Commit cad7e47

Browse files
committed
Address CodeRabbit review feedback
- raster.py: cap `**` in band math. The AST allowlist permitted ast.Pow with no bound, so `9**9**6` or `9**999999` built a multi-megabyte exact integer and pinned a core for tens of seconds before the post-eval shape check could reject it. An exponent must now be a numeric literal within MAX_POW_EXPONENT (64, far past squares/gamma/roots) unless it derives from a band or curated call, where NumPy's element-wise float power is bounded by the raster shape. - tests/test_raster.py: cover chained (`9**9**6`), oversized (`9**999999`), and computed (`9 ** (500 * 500)`) exponents, plus a positive case asserting ordinary powers including a negative literal exponent still evaluate.
1 parent 21810da commit cad7e47

2 files changed

Lines changed: 83 additions & 0 deletions

File tree

backend/geolibre_server/geolibre_server/app/raster.py

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1066,6 +1066,13 @@ def load_bands(path, letter, namespace, base):
10661066
ast.GtE,
10671067
ast.keyword,
10681068
)
1069+
# ``**`` on two plain integers stays exact, so ``9**9**6`` builds a
1070+
# half-million-digit result and pins a core for tens of seconds before the shape
1071+
# check can reject it. Require the exponent to be a small numeric literal unless
1072+
# it derives from a band or a curated call, whose element-wise float power is
1073+
# bounded by the raster's shape. 64 is far past anything band math needs
1074+
# (squares, gamma, roots).
1075+
MAX_POW_EXPONENT = 64
10691076
for node in ast.walk(tree):
10701077
if not isinstance(node, _allowed_nodes):
10711078
raise SystemExit(
@@ -1077,6 +1084,24 @@ def load_bands(path, letter, namespace, base):
10771084
"Expression may only use numeric constants "
10781085
"(band math only allows curated functions and operators)"
10791086
)
1087+
if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Pow):
1088+
exponent = node.right
1089+
while isinstance(exponent, ast.UnaryOp):
1090+
exponent = exponent.operand
1091+
band_backed = any(
1092+
isinstance(inner, (ast.Name, ast.Call)) for inner in ast.walk(exponent)
1093+
)
1094+
if not band_backed:
1095+
if not isinstance(exponent, ast.Constant):
1096+
raise SystemExit(
1097+
"Exponent must be a plain number or derive from a band "
1098+
f"(band math caps '**' at {MAX_POW_EXPONENT})"
1099+
)
1100+
if abs(exponent.value) > MAX_POW_EXPONENT:
1101+
raise SystemExit(
1102+
f"Exponent may not exceed {MAX_POW_EXPONENT} "
1103+
"(band math caps '**' to keep evaluation bounded)"
1104+
)
10801105
if isinstance(node, ast.Call):
10811106
# ``np.where(...)`` / ``A.tofile(...)`` parse as Attribute-backed calls;
10821107
# Attribute is already rejected above, but keep an explicit Name check.

backend/geolibre_server/tests/test_raster.py

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1093,3 +1093,61 @@ def test_raster_calculator_rejects_unbounded_allocations(tmp_path: Path) -> None
10931093
combined = completed.stdout + completed.stderr
10941094
assert "may not use" in combined or "numeric constants" in combined, expression
10951095
assert not out.exists()
1096+
1097+
1098+
@requires_rasterio
1099+
def test_raster_calculator_rejects_runaway_exponents(tmp_path: Path) -> None:
1100+
"""Chained or oversized ``**`` must be refused before eval computes a bignum."""
1101+
src = _write_dem(tmp_path / "dem.tif")
1102+
out = tmp_path / "calc.tif"
1103+
for expression in (
1104+
"A + 9**9**6",
1105+
"A + 9**999999",
1106+
"A + 9 ** (500 * 500)",
1107+
):
1108+
completed = subprocess.run(
1109+
[
1110+
sys.executable,
1111+
"-c",
1112+
_RASTER_TOOL_SCRIPTS["raster-calc"],
1113+
json.dumps(
1114+
{
1115+
"input_path": str(src),
1116+
"output_path": str(out),
1117+
"expression": expression,
1118+
}
1119+
),
1120+
],
1121+
check=False,
1122+
capture_output=True,
1123+
text=True,
1124+
timeout=60,
1125+
)
1126+
assert completed.returncode != 0, expression
1127+
combined = completed.stdout + completed.stderr
1128+
assert "Exponent" in combined, expression
1129+
assert not out.exists()
1130+
1131+
1132+
@requires_rasterio
1133+
def test_raster_calculator_allows_ordinary_exponents(tmp_path: Path) -> None:
1134+
"""The ``**`` cap must not disturb the small powers band math actually uses."""
1135+
import numpy as np
1136+
import rasterio
1137+
1138+
src = _write_dem(tmp_path / "dem.tif") # band 1 = x + y
1139+
out = tmp_path / "calc.tif"
1140+
_run_script(
1141+
_RASTER_TOOL_SCRIPTS["raster-calc"],
1142+
{
1143+
"input_path": str(src),
1144+
"output_path": str(out),
1145+
# A negative literal exponent must survive the unary-sign strip.
1146+
"expression": "(A ** 2) ** 0.5 * 2 ** -2",
1147+
},
1148+
)
1149+
with rasterio.open(src) as ds:
1150+
a = ds.read(1).astype("float64")
1151+
with rasterio.open(out) as ds:
1152+
result = ds.read(1).astype("float64")
1153+
assert np.allclose(result, (a**2) ** 0.5 * 0.25, atol=1e-4)

0 commit comments

Comments
 (0)