uce/scripts/wasm/check_unit_wasm.py
2026-06-13 02:07:38 +00:00

228 lines
8.2 KiB
Python
Executable File

#!/usr/bin/env python3
"""Validate a W2 UCE PIC unit wasm artifact.
Checks intentionally stay small and explicit: section walk for dylink.0,
uce.abi, imports/exports, plus llvm-nm for allocator definitions.
"""
from __future__ import annotations
import argparse
import shutil
import subprocess
import sys
from pathlib import Path
def read_u32leb(data: bytes, pos: int) -> tuple[int, int]:
result = 0
shift = 0
while True:
if pos >= len(data):
raise ValueError("truncated leb128")
b = data[pos]
pos += 1
result |= (b & 0x7F) << shift
if (b & 0x80) == 0:
return result, pos
shift += 7
if shift > 35:
raise ValueError("oversized leb128")
def read_name(data: bytes, pos: int) -> tuple[str, int]:
n, pos = read_u32leb(data, pos)
end = pos + n
if end > len(data):
raise ValueError("truncated name")
return data[pos:end].decode("utf-8", "replace"), end
def walk_sections(data: bytes):
if not data.startswith(b"\0asm\x01\0\0\0"):
raise ValueError("not a wasm v1 module")
pos = 8
while pos < len(data):
section_id = data[pos]
pos += 1
size, pos = read_u32leb(data, pos)
end = pos + size
if end > len(data):
raise ValueError("section extends past EOF")
payload = data[pos:end]
yield section_id, payload
pos = end
def parse_imports(payload: bytes):
pos = 0
count, pos = read_u32leb(payload, pos)
imports = []
for _ in range(count):
module, pos = read_name(payload, pos)
name, pos = read_name(payload, pos)
if pos >= len(payload):
raise ValueError("truncated import kind")
kind = payload[pos]
pos += 1
# Skip type descriptors. We only need module/name/kind for W2 policy.
if kind == 0: # func type index
_, pos = read_u32leb(payload, pos)
elif kind == 1: # table
if pos >= len(payload): raise ValueError("truncated table import")
pos += 1
flags, pos = read_u32leb(payload, pos)
_, pos = read_u32leb(payload, pos)
if flags & 1: _, pos = read_u32leb(payload, pos)
elif kind == 2: # memory
flags, pos = read_u32leb(payload, pos)
_, pos = read_u32leb(payload, pos)
if flags & 1: _, pos = read_u32leb(payload, pos)
elif kind == 3: # global
pos += 2
else:
raise ValueError(f"unknown import kind {kind}")
imports.append((module, name, kind))
return imports
def parse_exports(payload: bytes):
pos = 0
count, pos = read_u32leb(payload, pos)
exports = []
for _ in range(count):
name, pos = read_name(payload, pos)
if pos >= len(payload):
raise ValueError("truncated export kind")
kind = payload[pos]
pos += 1
_, pos = read_u32leb(payload, pos)
exports.append((name, kind))
return exports
def collect(path: Path):
data = path.read_bytes()
customs: dict[str, list[bytes]] = {}
imports = []
exports = []
for section_id, payload in walk_sections(data):
if section_id == 0:
name, pos = read_name(payload, 0)
customs.setdefault(name, []).append(payload[pos:])
elif section_id == 2:
imports = parse_imports(payload)
elif section_id == 7:
exports = parse_exports(payload)
return customs, imports, exports
def dylink_has_valid_mem_info(payload: bytes) -> bool:
pos = 0
while pos < len(payload):
subsection_id = payload[pos]
pos += 1
size, pos = read_u32leb(payload, pos)
end = pos + size
if end > len(payload):
raise ValueError("dylink.0 subsection extends past section")
if subsection_id == 1:
mem_size, p = read_u32leb(payload, pos)
mem_align, p = read_u32leb(payload, p)
table_size, p = read_u32leb(payload, p)
table_align, p = read_u32leb(payload, p)
if p > end:
raise ValueError("truncated dylink.0 mem_info")
return mem_align < 32 and table_align < 32 and mem_size < (1 << 31) and table_size < (1 << 31)
pos = end
return False
def defined_symbols(path: Path, llvm_nm: str) -> list[str]:
proc = subprocess.run([llvm_nm, "--defined-only", str(path)], text=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
if proc.returncode != 0:
raise RuntimeError(proc.stderr.strip() or "llvm-nm failed")
symbols = []
for line in proc.stdout.splitlines():
parts = line.split()
if parts:
symbols.append(parts[-1])
return symbols
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("wasm", type=Path)
ap.add_argument("--abi-version", default="6")
ap.add_argument("--llvm-nm", default=None)
ap.add_argument("--verbose", action="store_true")
args = ap.parse_args()
try:
customs, imports, exports = collect(args.wasm)
errors = []
dylink_payloads = customs.get("dylink.0", [])
if not dylink_payloads:
errors.append("missing dylink.0 custom section")
elif not any(dylink_has_valid_mem_info(payload) for payload in dylink_payloads):
errors.append("dylink.0 missing valid mem_info subsection")
abi_payloads = customs.get("uce.abi", [])
if not abi_payloads:
errors.append("missing uce.abi custom section")
else:
abi_text = abi_payloads[-1].decode("utf-8", "replace")
required = ["format=uce-wasm-unit-abi-v1", f"unit_abi_version={args.abi_version}", "toolchain="]
for needle in required:
if needle not in abi_text:
errors.append(f"uce.abi missing {needle!r}")
export_names = {name for name, _ in exports}
forbidden_exports = {"uce_alloc", "uce_free"}
for name in sorted(export_names & forbidden_exports):
errors.append(f"forbidden allocator export {name}")
import_map = {(module, name): kind for module, name, kind in imports}
required_imports = {
("env", "memory"): 2,
("env", "__memory_base"): 3,
}
# units without indirect calls / stack spills / table needs
# legitimately omit these; if present, the kind must be right
optional_imports = {
("env", "__indirect_function_table"): 1,
("env", "__stack_pointer"): 3,
("env", "__table_base"): 3,
}
for key, kind in required_imports.items():
if import_map.get(key) != kind:
errors.append(f"missing required import {key[0]}.{key[1]}")
for key, kind in optional_imports.items():
if key in import_map and import_map[key] != kind:
errors.append(f"wrong kind for import {key[0]}.{key[1]}")
for module, name, kind in imports:
if module.startswith("wasi_") or module == "wasi_snapshot_preview1":
errors.append(f"forbidden WASI import {module}.{name}")
if module not in {"env", "GOT.mem", "GOT.func"} and not module.startswith("GOT."):
errors.append(f"unexpected import module {module}.{name}")
if module.startswith("GOT.") and kind != 3:
errors.append(f"GOT import is not a global: {module}.{name}")
llvm_nm = args.llvm_nm or shutil.which("llvm-nm") or "/opt/wasi-sdk/bin/llvm-nm"
if Path(llvm_nm).exists():
bad_prefixes = ("_Znwm", "_Znam", "_ZdlPv", "_ZdaPv", "_ZdlPvm", "_ZdaPvm")
for sym in defined_symbols(args.wasm, llvm_nm):
if sym in {"uce_alloc", "uce_free"} or sym.startswith(bad_prefixes):
errors.append(f"forbidden allocator definition {sym}")
else:
errors.append("llvm-nm not found; cannot verify allocator definitions")
if errors:
for e in errors:
print(f"ERROR: {e}", file=sys.stderr)
return 1
if args.verbose:
print(f"UCE W2 unit check PASS: {args.wasm}")
return 0
except Exception as exc:
print(f"ERROR: {exc}", file=sys.stderr)
return 1
if __name__ == "__main__":
raise SystemExit(main())