learning/deep-probes/inspect_memory.py

"""Inspect the section forms used by this course's Wasm artifact.

This is a small read-only inspector, not a general Wasm validator.
"""
import argparse
import hashlib
import io
import json
from pathlib import Path


def integer(stream, signed=False):
    value = shift = 0
    while True:
        raw = stream.read(1)
        if not raw:
            raise ValueError("unexpected EOF in LEB integer")
        byte = raw[0]
        value |= (byte & 127) << shift
        shift += 7
        if shift > 70:
            raise ValueError("LEB integer too long")
        if byte < 128:
            return value - (1 << shift) if signed and byte & 64 else value


def limits(stream):
    flags = integer(stream)
    if flags not in (0, 1):
        raise ValueError("inspector supports only non-shared 32-bit limits")
    result = {"minimum": integer(stream)}
    if flags & 1:
        result["maximum"] = integer(stream)
    return result


def expression(stream):
    opcode = stream.read(1)
    if opcode not in (b"\x41", b"\x42"):
        raise ValueError("inspector supports integer constant initializers")
    result = integer(stream, signed=True)
    if stream.read(1) != b"\x0b":
        raise ValueError("unexpected constant expression")
    return result


def inspect(filename):
    payload = Path(filename).read_bytes()
    stream = io.BytesIO(payload)
    if stream.read(8) != b"\0asm\x01\0\0\0":
        raise ValueError("unexpected Wasm header")
    result: dict[str, object] = {"sha256": hashlib.sha256(payload).hexdigest()}
    while tag := stream.read(1):
        size = integer(stream)
        section = stream.read(size)
        if len(section) != size:
            raise ValueError("truncated section")
        part = io.BytesIO(section)
        if tag == b"\x02":
            imports = []
            for _ in range(integer(part)):
                module = part.read(integer(part)).decode()
                name = part.read(integer(part)).decode()
                if part.read(1) != b"\0":
                    raise ValueError("inspector supports function imports only")
                imports.append({"module": module, "name": name, "type_index": integer(part)})
            result["imports"] = imports
        elif tag == b"\x04":
            result["tables"] = [{"element_type": part.read(1).hex(), **limits(part)} for _ in range(integer(part))]
        elif tag == b"\x05":
            result["memories"] = [limits(part) for _ in range(integer(part))]
        elif tag == b"\x06":
            result["globals"] = [{"type": part.read(1).hex(), "mutable": bool(part.read(1)[0]), "initializer": expression(part)} for _ in range(integer(part))]
        elif tag == b"\x0b":
            count = integer(part)
            active = []
            passive_sizes = []
            for _ in range(count):
                mode = integer(part)
                if mode == 0:
                    offset = expression(part)
                elif mode == 1:
                    offset = None
                elif mode == 2:
                    if integer(part) != 0:
                        raise ValueError("inspector supports memory zero only")
                    offset = expression(part)
                else:
                    raise ValueError("unknown data segment mode")
                length = integer(part)
                if len(part.read(length)) != length:
                    raise ValueError("truncated data segment")
                if offset is None:
                    passive_sizes.append(length)
                else:
                    active.append((offset, length))
            result["data"] = {"segments": count, "active_segments": len(active), "passive_sizes": passive_sizes, "active_payload_bytes": sum(n for _, n in active), "first_active_offset": min((p for p, _ in active), default=None), "last_active_end": max((p+n for p, n in active), default=None)}
        else:
            continue
        if part.tell() != size:
            raise ValueError("unconsumed bytes in inspected section")
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("wasm", nargs="?", default="dist/spine.wasm")
    print(json.dumps(inspect(parser.parse_args().wasm), indent=2))