Files
n_signer/firmware/teensy41/test_pq_profile.py
T

381 lines
13 KiB
Python

#!/usr/bin/env python3
"""Profile each PQ verb ALONE on a fresh boot for the Teensy 4.1 n_signer.
For each test case:
1. Reflash the firmware (arduino-cli upload) so the device is on a fresh boot.
2. Open /dev/ttyACM0, wait 6s for boot, drain boot output
(which includes any "LAST OP BEFORE CRASH" / CRASH REPORT from the
previous fault, plus the boot banner).
3. Send exactly ONE JSON-RPC request over the 4-byte big-endian length
prefix wire protocol.
4. Try to read the response with a generous timeout. If we get a valid
response, record PASS. If we time out / get garbage, record CRASH
and re-open the port to drain the *next* boot's diagnostics
("LAST OP BEFORE CRASH: id=X seq=Y stack_hw_free=Z heap_free_at_crash=W"
plus the CrashReport text).
Usage:
python3 firmware/teensy41/test_pq_profile.py [--port /dev/ttyACM0]
python3 firmware/teensy41/test_pq_profile.py --only ml-kem
"""
import serial, struct, json, time, sys, argparse, re, subprocess, os
DEFAULT_PORT = "/dev/ttyACM0"
BAUD = 115200
SIGNER_DIR = "firmware/teensy41/signer"
FQBN = "teensy:avr:teensy41"
# OP_ID map (from signer.ino) for human-readable crash ids.
OP_ID_NAMES = {
1: "get_info", 2: "get_public_key", 3: "sign", 4: "verify",
5: "derive_shared_secret", 6: "derive", 7: "nostr_get_public_key",
8: "nostr_sign_event", 9: "nostr_mine_event", 10: "nip04",
11: "nip44", 12: "encapsulate", 13: "decapsulate", 14: "otp",
}
# Test cases. Each is ONE verb on a fresh boot.
TESTS = [
{
"name": "ml-dsa-65 get_public_key",
"method": "get_public_key",
"params": [{"algorithm": "ml-dsa-65", "index": 0}],
"expect_op_id": 2,
},
{
"name": "slh-dsa-128s get_public_key",
"method": "get_public_key",
"params": [{"algorithm": "slh-dsa-128s", "index": 0}],
"expect_op_id": 2,
},
{
"name": "ml-kem-768 get_public_key",
"method": "get_public_key",
"params": [{"algorithm": "ml-kem-768", "index": 0}],
"expect_op_id": 2,
},
{
"name": "ml-dsa-65 sign",
"method": "sign",
"params": ["746573742065643235353139206d657373616765",
{"algorithm": "ml-dsa-65", "index": 0}],
"expect_op_id": 3,
},
{
"name": "slh-dsa-128s sign",
"method": "sign",
"params": ["746573742065643235353139206d657373616765",
{"algorithm": "slh-dsa-128s", "index": 0}],
"expect_op_id": 3,
},
]
def wait_for_port(port, timeout=30.0, must_exist=True):
"""Wait for the serial device node to (re)appear after a reboot/reflash."""
t0 = time.time()
while time.time() - t0 < timeout:
if os.path.exists(port):
return True
time.sleep(0.25)
if must_exist:
print(f" [port] {port} did not appear within {timeout:.0f}s")
return False
def reflash(port, retries=3):
"""Upload firmware and return True on success. Retries on failure
(the device may be mid-reboot or the port held open)."""
for attempt in range(1, retries + 1):
print(f" [reflash] (attempt {attempt}/{retries}) arduino-cli upload -p {port} --fqbn {FQBN} {SIGNER_DIR}")
try:
r = subprocess.run(
["arduino-cli", "upload", "-p", port, "--fqbn", FQBN, SIGNER_DIR],
cwd=os.getcwd(), capture_output=True, text=True, timeout=120,
)
except Exception as e:
print(f" [reflash] exception: {e}")
time.sleep(2.0)
continue
if r.returncode == 0:
print(f" [reflash] OK")
return True
# Common failure: device mid-reboot / port gone. Wait for it to
# come back, then retry.
print(f" [reflash] rc={r.returncode}")
tail = (r.stdout + r.stderr)[-400:]
if tail.strip():
print(f" [reflash] tail: {tail.strip()}")
# If the port is gone, wait for re-enumeration first.
wait_for_port(port, timeout=15.0, must_exist=False)
time.sleep(1.0)
print(f" [reflash] FAILED after {retries} attempts")
return False
def open_port(port):
return serial.Serial(port, BAUD, timeout=2.0)
def drain(ser, label, seconds=2.0):
end = time.time() + seconds
buf = b""
while time.time() < end:
n = ser.in_waiting
if n:
buf += ser.read(n)
else:
time.sleep(0.05)
if buf:
txt = buf.decode("utf-8", errors="replace")
print(f"--- {label} ---")
print(txt)
print(f"--- end {label} ---")
return txt
return ""
def send_request(ser, req, timeout=90.0):
payload = json.dumps(req).encode("utf-8")
ser.write(struct.pack(">I", len(payload)) + payload)
ser.flush()
h = b""
deadline = time.time() + timeout
while len(h) < 4 and time.time() < deadline:
c = ser.read(4 - len(h))
if c:
h += c
else:
time.sleep(0.02)
if len(h) < 4:
raise TimeoutError("response header timeout")
n = struct.unpack(">I", h)[0]
if n == 0 or n > 65536:
raise ValueError(f"bad response len {n} (raw hdr={h!r})")
p = b""
while len(p) < n and time.time() < deadline:
c = ser.read(n - len(p))
if c:
p += c
else:
time.sleep(0.02)
if len(p) < n:
raise TimeoutError(f"response body timeout ({len(p)}/{n})")
return json.loads(p.decode("utf-8"))
LAST_OP_RE = re.compile(
r"LAST OP BEFORE CRASH: id=(\d+)\s+seq=(\d+)\s+stack_hw_free=(\d+)\s+heap_free_at_crash=(\d+)"
)
def parse_last_op(text):
m = LAST_OP_RE.search(text)
if not m:
return None
return {
"op_id": int(m.group(1)),
"op_name": OP_ID_NAMES.get(int(m.group(1)), "?"),
"seq": int(m.group(2)),
"stack_hw_free": int(m.group(3)),
"heap_free_at_crash": int(m.group(4)),
}
def extract_crash_report(text):
if "CRASH REPORT" not in text:
return None
idx = text.find("CRASH REPORT")
# Capture up to the next "n_signer booting..." or 1500 chars.
end = text.find("n_signer booting", idx)
if end == -1:
end = idx + 1500
return text[idx:end].strip()
def run_one_case(port, case):
print(f"\n{'='*70}")
print(f"TEST: {case['name']}")
print(f" method={case['method']} params={json.dumps(case['params'])}")
print(f"{'='*70}")
result = {
"case": case["name"],
"method": case["method"],
"params": case["params"],
"outcome": None, # "PASS" | "CRASH" | "ERROR"
"response": None,
"boot_output": None,
"crash_report": None,
"last_op_before_crash": None,
"next_boot_last_op": None,
"next_boot_crash_report": None,
"error": None,
}
# 1. Reflash for a guaranteed fresh boot.
if not reflash(port):
result["outcome"] = "ERROR"
result["error"] = "reflash failed"
return result
# 2. Wait for the port to re-enumerate after the reboot, then for
# boot (auto-generate mnemonic + derive keys).
if not wait_for_port(port, timeout=20.0):
result["outcome"] = "ERROR"
result["error"] = f"{port} did not re-appear after reflash"
print(f" !! {port} did not re-appear after reflash")
return result
time.sleep(6)
ser = None
for attempt in range(1, 6):
try:
ser = open_port(port)
break
except Exception as e:
print(f" !! open failed (attempt {attempt}/5): {e}")
time.sleep(1.0)
if ser is None:
result["outcome"] = "ERROR"
result["error"] = "open failed after 5 attempts"
print(f" !! open failed after 5 attempts")
return result
try:
boot_txt = drain(ser, "BOOT OUTPUT", seconds=2.0)
result["boot_output"] = boot_txt
if "CRASH REPORT" in boot_txt:
result["crash_report"] = extract_crash_report(boot_txt)
lop = parse_last_op(boot_txt)
if lop:
result["last_op_before_crash"] = lop
print(f" [boot] LAST OP: {lop}")
# 3. Send exactly one request.
req = {"jsonrpc": "2.0", "id": 1, "method": case["method"],
"params": case["params"]}
print(f" -> send {case['method']}")
t0 = time.time()
try:
resp = send_request(ser, req, timeout=90.0)
dt = time.time() - t0
result["response"] = resp
if "result" in resp:
result["outcome"] = "PASS"
rstr = json.dumps(resp)
print(f" <- OK in {dt:.1f}s ({len(rstr)} bytes): {rstr[:120]}...")
else:
result["outcome"] = "ERROR"
print(f" <- ERR in {dt:.1f}s: {json.dumps(resp)[:200]}")
except (TimeoutError, ValueError, Exception) as e:
dt = time.time() - t0
result["outcome"] = "CRASH"
result["error"] = f"{type(e).__name__}: {e} (after {dt:.1f}s)"
print(f" !! CRASH during {case['method']} after {dt:.1f}s: {e}")
try:
drain(ser, "PARTIAL OUTPUT AFTER CRASH", seconds=2.0)
except Exception:
pass
try:
ser.close()
except Exception:
pass
# Wait for the port to come back (Teensy CrashReport usually
# auto-reboots; the device may briefly disappear).
wait_for_port(port, timeout=20.0, must_exist=False)
# Re-open IMMEDIATELY — the device's setup() only waits 3s
# (while (!Serial && millis() < 3000)) for the host to open
# the port before it stops buffering early boot prints. If we
# sleep too long we miss the "LAST OP BEFORE CRASH" /
# "CRASH REPORT" lines that are emitted at the very top of
# setup().
ser2 = None
for attempt in range(1, 10):
try:
ser2 = open_port(port)
break
except Exception:
time.sleep(0.3)
if ser2 is None:
result["error"] += " | next-boot port never came back"
print(f" !! next-boot port never came back")
return result
try:
# Drain immediately and keep draining for a while to
# capture the full boot banner.
next_txt = drain(ser2, "NEXT-BOOT OUTPUT", seconds=8.0)
result["next_boot_last_op"] = parse_last_op(next_txt)
result["next_boot_crash_report"] = extract_crash_report(next_txt)
if result["next_boot_last_op"]:
print(f" [next-boot] LAST OP: {result['next_boot_last_op']}")
if result["next_boot_crash_report"]:
print(f" [next-boot] CRASH REPORT:\n{result['next_boot_crash_report'][:600]}")
else:
print(f" [next-boot] NO CRASH REPORT / LAST OP line found in boot output")
except Exception as e2:
result["error"] += f" | next-boot drain failed: {e2}"
print(f" !! next-boot drain failed: {e2}")
finally:
try:
ser2.close()
except Exception:
pass
return result
# After a successful verb, drain any trailing log output.
drain(ser, "TRAILING OUTPUT", seconds=1.0)
finally:
try:
ser.close()
except Exception:
pass
return result
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--port", default=DEFAULT_PORT)
ap.add_argument("--only", default=None,
help="substring filter on case name; run only matching cases")
args = ap.parse_args()
cases = TESTS
if args.only:
cases = [c for c in TESTS if args.only.lower() in c["name"].lower()]
if not cases:
print(f"No cases match --only {args.only!r}")
return 1
print(f"Running {len(cases)} PQ profiling case(s) on {args.port}.")
print("Each case reflashes the firmware first, so each verb runs on a fresh boot.")
results = []
for case in cases:
r = run_one_case(args.port, case)
results.append(r)
print(f"\n{'='*70}")
print("SUMMARY")
print(f"{'='*70}")
for r in results:
line = f" {r['case']:30s} -> {r['outcome']}"
if r["next_boot_last_op"]:
l = r["next_boot_last_op"]
line += (f" | last_op={l['op_name']} seq={l['seq']} "
f"stack_hw_free={l['stack_hw_free']} "
f"heap_free_at_crash={l['heap_free_at_crash']}")
elif r["last_op_before_crash"]:
l = r["last_op_before_crash"]
line += (f" | boot_last_op={l['op_name']} seq={l['seq']} "
f"stack_hw_free={l['stack_hw_free']} "
f"heap_free_at_crash={l['heap_free_at_crash']}")
if r["error"]:
line += f" | err={r['error'][:80]}"
print(line)
with open("firmware/teensy41/pq_profile_results.json", "w") as f:
json.dump(results, f, indent=2)
print(f"\nFull results written to firmware/teensy41/pq_profile_results.json")
return 0
if __name__ == "__main__":
sys.exit(main())