task-121: N+1 扫描审计（循环内单查扫描器 + 10 条测试 + 报告快照，不改生产代码）
This commit is contained in:
@@ -0,0 +1,313 @@
|
||||
"""task-121 N+1 扫描审计工具。
|
||||
|
||||
对 backend-java modules 下的 Java Service 源码做机械化扫描：
|
||||
识别循环（for/while/forEach/map）内对 MyBatis-Plus 单行查询
|
||||
（selectById/selectOne/selectList/selectCount/selectObjs/selectMaps/getById/getOne）
|
||||
的调用，产出 N+1 候选清单（位置/对象/影响），供人工审计与 task-122 修复。
|
||||
本工具只读，不改任何代码。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
HIGH = "HIGH"
|
||||
MEDIUM = "MEDIUM"
|
||||
LOW = "LOW"
|
||||
IMPACT_ORDER = {HIGH: 0, MEDIUM: 1, LOW: 2}
|
||||
|
||||
MODULES_ROOT = (
|
||||
Path(__file__).resolve().parent.parent / "src/main/java/com/nanri/aiimage/modules"
|
||||
)
|
||||
|
||||
# 单行查询调用（selectBatchIds/selectByIds 天然批量，不匹配此正则）
|
||||
QUERY_RE = re.compile(
|
||||
r"\.(selectById|selectOne|selectList|selectCount|selectObjs|selectMaps|getById|getOne)\("
|
||||
)
|
||||
# 循环头：for/while 关键字 + 圆括号跨度，或 .forEach( / .map( 方法调用
|
||||
LOOP_RE = re.compile(r"\b(for|while)\s*\(|\.(forEach|map)\s*\(")
|
||||
# 返回集合的查询调用（循环遍历这类结果集 → 影响 HIGH）
|
||||
COLLECTION_QUERY_RE = re.compile(
|
||||
r"(selectList|selectObjs|selectMaps|selectBatchIds|selectByIds|\.list\()"
|
||||
)
|
||||
ASSIGN_RE = re.compile(r"(\w+)\s*=\s*([^;{}=]+);")
|
||||
FIELD_RE = re.compile(r"\b[A-Za-z_]\w*\b")
|
||||
|
||||
|
||||
def _strip_comments(text: str) -> str:
|
||||
"""用空格替换注释（保留换行与偏移），简化后续扫描。"""
|
||||
out = list(text)
|
||||
i, n = 0, len(text)
|
||||
while i < n:
|
||||
if text.startswith("/*", i):
|
||||
j = text.find("*/", i + 2)
|
||||
j = n if j < 0 else j + 2
|
||||
for k in range(i, j):
|
||||
if text[k] != "\n":
|
||||
out[k] = " "
|
||||
i = j
|
||||
elif text.startswith("//", i):
|
||||
j = text.find("\n", i)
|
||||
j = n if j < 0 else j
|
||||
for k in range(i, j):
|
||||
out[k] = " "
|
||||
i = j
|
||||
else:
|
||||
i += 1
|
||||
return "".join(out)
|
||||
|
||||
|
||||
def _match_span(text: str, open_pos: int) -> int | None:
|
||||
"""从 text[open_pos] 的 ( 出发，返回匹配的 ) 偏移（含），不匹配返回 None。"""
|
||||
depth = 0
|
||||
for i in range(open_pos, len(text)):
|
||||
c = text[i]
|
||||
if c == "(":
|
||||
depth += 1
|
||||
elif c == ")":
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return i
|
||||
return None
|
||||
|
||||
|
||||
def _line_of(text: str, offset: int) -> int:
|
||||
return text.count("\n", 0, offset) + 1
|
||||
|
||||
|
||||
class _Frame:
|
||||
__slots__ = ("body_depth", "expr", "loop_type", "start")
|
||||
|
||||
def __init__(self, start: int, body_depth: int, loop_type: str, expr: str):
|
||||
self.start = start
|
||||
self.body_depth = body_depth
|
||||
self.loop_type = loop_type
|
||||
self.expr = expr
|
||||
|
||||
|
||||
def _query_results(stripped: str) -> set[str]:
|
||||
"""收集被赋值给集合查询结果（selectList 等）的变量名。"""
|
||||
names = set()
|
||||
for lhs, rhs in ASSIGN_RE.findall(stripped):
|
||||
if COLLECTION_QUERY_RE.search(rhs):
|
||||
names.add(lhs)
|
||||
return names
|
||||
|
||||
|
||||
def _loop_headers(stripped: str) -> list[dict]:
|
||||
"""预取全部循环头：(start, end, type, expr)（start 为 ( 偏移，end 为匹配的 ) 偏移）。"""
|
||||
headers: list[dict] = []
|
||||
pos = 0
|
||||
while True:
|
||||
m = LOOP_RE.search(stripped, pos)
|
||||
if not m:
|
||||
break
|
||||
paren_pos = m.end() - 1
|
||||
if stripped[paren_pos] == "(":
|
||||
close = _match_span(stripped, paren_pos)
|
||||
if close is not None:
|
||||
headers.append(
|
||||
{
|
||||
"start": paren_pos,
|
||||
"end": close,
|
||||
"type": m.group(1) or m.group(2),
|
||||
"expr": stripped[paren_pos + 1 : close].strip(),
|
||||
}
|
||||
)
|
||||
pos = m.end()
|
||||
return headers
|
||||
|
||||
|
||||
def _header_context(stripped: str, h: dict, offset: int) -> tuple[bool, str, str]:
|
||||
"""查询落在循环头括号跨度内时的归属判定。
|
||||
|
||||
返回 (是否循环内, 循环类型, 循环表达式)：
|
||||
- lambda 箭头之后 → 循环内（无花括号 lambda 体）
|
||||
- 增强 for 的 ` : ` 之后（迭代源 selectList 等）→ 集合装载，非循环内
|
||||
- for/while 初始化或条件 → 视为循环内（条件逐次求值）
|
||||
"""
|
||||
expr = h["expr"]
|
||||
if "->" in stripped[h["start"] : h["end"]]:
|
||||
arrow = stripped.find("->", h["start"], h["end"])
|
||||
if arrow != -1 and offset > arrow:
|
||||
return True, h["type"], expr
|
||||
return False, h["type"], expr
|
||||
if " : " in expr:
|
||||
colon = stripped.find(" : ", h["start"], h["end"])
|
||||
if colon != -1 and offset > colon:
|
||||
return False, h["type"], expr
|
||||
return True, h["type"], expr
|
||||
return True, h["type"], expr
|
||||
|
||||
|
||||
def scan_text(text: str, filename: str = "fixture.java") -> list[dict]:
|
||||
"""扫描单份 Java 源码，返回 N+1 候选清单（按影响降序、文件行号升序）。"""
|
||||
stripped = _strip_comments(text)
|
||||
query_results = _query_results(stripped)
|
||||
headers = _loop_headers(stripped)
|
||||
|
||||
depth = 0
|
||||
frames: list[_Frame] = []
|
||||
single: tuple[str, str, int, int] | None = None # (loop_type, expr, start, end)
|
||||
findings: list[dict] = []
|
||||
|
||||
def in_loop(offset: int) -> bool:
|
||||
if single is not None and single[2] <= offset < single[3]:
|
||||
return True
|
||||
if any(f.start <= offset and depth >= f.body_depth for f in frames):
|
||||
return True
|
||||
for h in reversed(headers):
|
||||
if h["start"] <= offset <= h["end"]:
|
||||
return _header_context(stripped, h, offset)[0]
|
||||
return False
|
||||
|
||||
def enclosing(offset: int) -> tuple[str, str]:
|
||||
if single is not None and single[2] <= offset < single[3]:
|
||||
return single[0], single[1]
|
||||
for h in reversed(headers):
|
||||
if h["start"] <= offset <= h["end"]:
|
||||
in_loop, loop_type, expr = _header_context(stripped, h, offset)
|
||||
if in_loop:
|
||||
return loop_type, expr
|
||||
for f in reversed(frames):
|
||||
if f.start <= offset:
|
||||
return f.loop_type, f.expr
|
||||
return "loop", ""
|
||||
|
||||
i = 0
|
||||
n = len(stripped)
|
||||
while i < n:
|
||||
c = stripped[i]
|
||||
if c == "{":
|
||||
depth += 1
|
||||
i += 1
|
||||
continue
|
||||
if c == "}":
|
||||
depth -= 1
|
||||
if frames and depth == frames[-1].body_depth - 1:
|
||||
frames.pop()
|
||||
i += 1
|
||||
continue
|
||||
|
||||
loop_match = LOOP_RE.match(stripped, i)
|
||||
if loop_match:
|
||||
loop_type = loop_match.group(1) or loop_match.group(2)
|
||||
paren_pos = loop_match.end() - 1
|
||||
close = (
|
||||
_match_span(stripped, paren_pos) if stripped[paren_pos] == "(" else None
|
||||
)
|
||||
if close is not None:
|
||||
expr = stripped[paren_pos + 1 : close].strip()
|
||||
body_start = close + 1
|
||||
while body_start < n and stripped[body_start].isspace():
|
||||
body_start += 1
|
||||
if body_start < n and stripped[body_start] == "{":
|
||||
frames.append(_Frame(body_start + 1, depth + 1, loop_type, expr))
|
||||
else:
|
||||
end = stripped.find(";", body_start)
|
||||
single = (loop_type, expr, body_start, n if end < 0 else end)
|
||||
# 继续扫描括号跨度内容（无花括号 lambda 体的查询落在其中）
|
||||
i = paren_pos + 1
|
||||
continue
|
||||
i = loop_match.end()
|
||||
continue
|
||||
|
||||
query_match = QUERY_RE.match(stripped, i)
|
||||
if query_match:
|
||||
if in_loop(i):
|
||||
line_start = stripped.rfind("\n", 0, i) + 1
|
||||
line_end = stripped.find("\n", i)
|
||||
line_end = n if line_end < 0 else line_end
|
||||
statement = stripped[line_start:line_end].strip()
|
||||
loop_type, expr = enclosing(i)
|
||||
impact = _classify(loop_type, [expr], query_results)
|
||||
findings.append(
|
||||
{
|
||||
"file": filename,
|
||||
"line": _line_of(text, i),
|
||||
"statement": statement,
|
||||
"loop_type": loop_type,
|
||||
"impact": impact,
|
||||
}
|
||||
)
|
||||
i = query_match.end()
|
||||
continue
|
||||
|
||||
i += 1
|
||||
|
||||
findings.sort(key=lambda f: (IMPACT_ORDER[f["impact"]], f["line"]))
|
||||
return findings
|
||||
|
||||
|
||||
def _classify(loop_type: str, exprs: list[str], query_results: set[str]) -> str:
|
||||
"""影响分级：遍历查询结果集 → HIGH；forEach/map lambda → MEDIUM；其余 → LOW。"""
|
||||
for expr in exprs:
|
||||
if expr and COLLECTION_QUERY_RE.search(expr):
|
||||
return HIGH
|
||||
if expr and any(t in query_results for t in FIELD_RE.findall(expr)):
|
||||
return HIGH
|
||||
if loop_type in ("forEach", "map"):
|
||||
return MEDIUM
|
||||
return LOW
|
||||
|
||||
|
||||
@dataclass
|
||||
class Report:
|
||||
scanned_files: int = 0
|
||||
modules: list[str] = field(default_factory=list)
|
||||
findings: list[dict] = field(default_factory=list)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
return {
|
||||
"scanned_files": self.scanned_files,
|
||||
"modules": self.modules,
|
||||
"findings": self.findings,
|
||||
}
|
||||
|
||||
|
||||
def scan_modules(root: Path) -> Report:
|
||||
"""扫描 modules 根目录下全部 Java 文件，产出 N+1 候选报告。"""
|
||||
report = Report()
|
||||
module_set: set[str] = set()
|
||||
for java_file in sorted(root.rglob("*.java")):
|
||||
report.scanned_files += 1
|
||||
rel = java_file.relative_to(root)
|
||||
module_set.add(rel.parts[0])
|
||||
for finding in scan_text(java_file.read_text(encoding="utf-8"), str(rel)):
|
||||
finding["module"] = rel.parts[0]
|
||||
report.findings.append(finding)
|
||||
report.findings.sort(
|
||||
key=lambda f: (IMPACT_ORDER[f["impact"]], f["file"], f["line"])
|
||||
)
|
||||
report.modules = sorted(module_set)
|
||||
return report
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description="N+1 扫描审计")
|
||||
parser.add_argument("--json", help="输出 JSON 报告路径（默认 stdout）")
|
||||
args = parser.parse_args()
|
||||
|
||||
report = scan_modules(MODULES_ROOT)
|
||||
payload = report.to_dict()
|
||||
if args.json:
|
||||
Path(args.json).write_text(
|
||||
json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
|
||||
)
|
||||
else:
|
||||
print(json.dumps(payload, ensure_ascii=False, indent=2))
|
||||
print(
|
||||
f"scanned {report.scanned_files} files, "
|
||||
f"{len(report.findings)} N+1 candidates "
|
||||
f"across {len(report.modules)} modules",
|
||||
file=sys.stderr,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user