# /// script
# requires-python = ">=3.10"
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""The length-prefix attack, and the frame that never resyncs.

A TCP server reads a 4-byte length and then that many bytes. A client sends
FF FF FF FF (about 4 GB) and nothing else: the naive server reserves a
buffer of that size for one connection and its thread hangs waiting. With a
MAX_FRAME check it rejects the frame at once. Then one wrong length in the
middle of a valid stream turns every later frame into garbage.
"""
import argparse, json, socket, struct, threading, time, tracemalloc

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--max-frame", type=int, default=1 << 20, help="MAX_FRAME for the checked server, in bytes.")
p.add_argument("--wait", type=float, default=2.0, help="Seconds to watch the server before calling it hung.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()


def serve(check, results, frames_out=None):
    srv = socket.socket(); srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("127.0.0.1", 0)); srv.listen(1)
    def handle():
        conn, _ = srv.accept()
        f = conn.makefile("rb")
        try:
            while True:
                head = f.read(4)
                if len(head) < 4: break
                n = struct.unpack(">I", head)[0]
                if check and n > args.max_frame:
                    results["rejected"] = n; conn.close(); return
                results["claimed"] = n
                tracemalloc.start()
                buf = bytearray(n)                       # the naive server trusts the header
                results["reserved"] = tracemalloc.get_traced_memory()[1]; tracemalloc.stop()
                got = f.readinto(memoryview(buf))        # ... and waits for n bytes that never come
                if frames_out is not None: frames_out.append(bytes(buf[:got]))
                del buf
        except Exception as e:
            results["error"] = type(e).__name__
    threading.Thread(target=handle, daemon=True).start()
    return srv.getsockname()[1]


rows = []
for check in (False, True):
    res = {}
    port = serve(check, res)
    c = socket.create_connection(("127.0.0.1", port)); c.sendall(b"\xff\xff\xff\xff")
    time.sleep(args.wait)
    label = f"MAX_FRAME = {args.max_frame:,} B" if check else "trusts the header"
    if check:
        rows.append({"server": label, "outcome": f"rejected a {res.get('rejected', 0):,}-byte frame at once", "reserved_mb": 0.0})
    else:
        rows.append({"server": label, "outcome": f"reserved {res.get('reserved', 0) / 1e6:,.0f} MB for one connection; thread still waiting after {args.wait:g} s",
                     "reserved_mb": res.get("reserved", 0) / 1e6})
    print(f"{label:<26} {rows[-1]['outcome']}")
    c.close()

# A valid stream with one wrong length in the middle
frames = [b"alpha", b"bravo", b"charlie", b"delta", b"echo"]
res, out = {}, []
port = serve(True, res, out)
c = socket.create_connection(("127.0.0.1", port))
for i, fr in enumerate(frames):
    n = len(fr) + (2 if i == 1 else 0)                 # frame 2 claims two bytes more than it has
    c.sendall(struct.pack(">I", n) + fr)
time.sleep(0.5); c.close(); time.sleep(0.3)
print(f"sent {[f.decode() for f in frames]}; frame 2's length is off by 2")
print(f"server read {[f.decode('latin-1') for f in out]}")
garbage = sum(1 for a, b in zip(frames, out) if a != b) + (len(frames) - len(out))
rows.append({"server": "one wrong length mid-stream", "outcome": f"{garbage} of {len(frames)} frames wrong or lost after it", "reserved_mb": 0.0})

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"servers": rows, "frames": [{"sent": a.decode(), "read": out[i].decode('latin-1') if i < len(out) else "(nothing)"}
                   for i, a in enumerate(frames)]}}, f)