Files
n_signer/firmware/teensy41/test_dsa_sign_diag.py
T

134 lines
4.3 KiB
Python

#!/usr/bin/env python3
"""ml-dsa-65 sign test with soft-reboot diagnostic recovery.
Runs sign; on hang/timeout, triggers a Teensy soft-reboot via DTR/RTS toggle
(which preserves DMAMEM) and reads the boot output for the reject-count
diagnostic. Falls back to a hard reboot hint if DTR toggle doesn't work.
"""
import serial, struct, json, time, sys, argparse
DEFAULT_PORT = "/dev/ttyACM0"
BAUD = 115200
def send_request(ser, req, timeout=180.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.01)
if len(h) < 4:
raise TimeoutError("hdr timeout")
n = struct.unpack(">I", h)[0]
if n == 0 or n > 65536:
raise ValueError("bad len %d" % n)
p = b""
while len(p) < n and time.time() < deadline:
c = ser.read(n - len(p))
if c:
p += c
else:
time.sleep(0.01)
if len(p) < n:
raise TimeoutError("body timeout")
return json.loads(p.decode())
def drain(ser, seconds=2.0):
buf = b""
end = time.time() + seconds
while time.time() < end:
if ser.in_waiting:
buf += ser.read(ser.in_waiting)
else:
time.sleep(0.05)
return buf.decode("utf-8", errors="replace")
def soft_reboot_and_read(port, wait_boot=8.0):
"""Toggle DTR/RTS to trigger Teensy bootloader/reset; preserve DMAMEM.
Teensy 4.1: setting DTR low then high does not reboot; but a BREAK condition
or the 1200-bps touch does. We try the 1200-bps open (Teensy bootloader
trigger) — note this may or may not preserve DMAMEM. As a safer soft-reboot
we instead pulse RTS/DTR which on some Teensy USB-CDC builds triggers a
watchdog reset.
"""
# Attempt 1: 1200 bps touch (Teensy bootloader reset). This is a hard reset
# but is the most reliable way to reboot a hung Teensy.
try:
s = serial.Serial(port, 1200, timeout=1.0)
s.dtr = False
s.rts = False
time.sleep(0.1)
s.dtr = True
time.sleep(0.05)
s.dtr = False
s.close()
except Exception as e:
print(f" (1200bps touch failed: {e})", flush=True)
time.sleep(1.0)
# Reopen at normal baud and drain boot
try:
ser = serial.Serial(port, BAUD, timeout=2.0)
except Exception as e:
print(f" (reopen failed: {e}; retrying in 5s)", flush=True)
time.sleep(5)
ser = serial.Serial(port, BAUD, timeout=2.0)
boot = drain(ser, seconds=wait_boot)
ser.close()
return boot
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--port", default=DEFAULT_PORT)
ap.add_argument("--timeout", type=float, default=180.0)
args = ap.parse_args()
ser = serial.Serial(args.port, BAUD, timeout=2.0)
boot = drain(ser, seconds=8.0)
print("=== BOOT OUTPUT (pre-sign) ===")
print(boot)
print("=== END BOOT ===", flush=True)
req = {"jsonrpc": "2.0", "id": 1, "method": "sign",
"params": ["746573742065643235353139206d657373616765",
{"algorithm": "ml-dsa-65", "index": 0}]}
print(f"-> sign ml-dsa-65 (timeout={args.timeout}s)", flush=True)
result = None
try:
resp = send_request(ser, req, timeout=args.timeout)
ok = "result" in resp
print(f"<- {'OK' if ok else 'ERR'} {str(resp)[:300]}", flush=True)
result = "pass" if ok else "err"
except Exception as e:
print(f"!! HANG/TIMEOUT: {e}", flush=True)
result = "hang"
ser.close()
if result == "hang":
print("\n=== Attempting soft reboot to read DMAMEM diagnostic ===", flush=True)
boot = soft_reboot_and_read(args.port, wait_boot=10.0)
print("=== BOOT OUTPUT (post-hang reboot) ===")
print(boot)
print("=== END BOOT ===", flush=True)
# Extract reject count
for line in boot.splitlines():
if "REJECT COUNT" in line:
print(f"\n>>> DIAGNOSTIC: {line.strip()}", flush=True)
if "CRASH" in line:
print(f"\n>>> CRASH: {line.strip()}", flush=True)
return 2
return 0 if result == "pass" else 1
if __name__ == "__main__":
sys.exit(main())