Add cert generation script and harden sample client

generate-certs.sh creates the certs/ directory with all certificates
needed to run the proxy locally. README quick-start section documents
the workflow.

Sample client hardened against livelock (bounded fast-recv loops),
SSLWantWriteError during handshake/send/recv, and unexpected
ConnectionTerminated frames from the proxy.

Author: sl5-guy

README.md

@@ -10,6 +10,16 @@ The proxy accepts incoming mTLS connections, extracts a custom extension value f
 cargo build --release
 ```
 
+## Quick start
+
+```bash
+./examples/generate-certs.sh          # creates certs/ directory
+cp config.example.toml config.toml    # edit to taste
+cargo run -- --config config.toml
+```
+
+The generated filenames match `config.example.toml` so no editing is needed for local development. Pass a custom extension value as an argument: `./examples/generate-certs.sh agent-beta`.
+
 ## Configuration
 
 Copy `config.example.toml` to `config.toml` and edit it. Key sections:

examples/generate-certs.sh

@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+# Generate a self-signed CA, proxy server cert, and client cert for local
+# development. Writes PEM files to a certs/ directory (created if absent).
+#
+# Usage:
+#   ./examples/generate-certs.sh                  # extension_value defaults to "agent-alpha"
+#   ./examples/generate-certs.sh agent-beta        # custom extension value
+#
+# The client cert contains a custom X.509 extension at OID 1.3.6.1.4.1.57264.1.1
+# with the given value encoded as a DER UTF8String.
+
+set -euo pipefail
+
+EXT_VALUE="${1:-agent-alpha}"
+DIR="certs"
+
+mkdir -p "$DIR"
+
+der_utf8string() {
+    local val="$1"
+    local len=${#val}
+    local result
+    result=$(printf '0c:%02x' "$len")
+    for (( i=0; i<len; i++ )); do
+        result=$(printf '%s:%02x' "$result" "'${val:$i:1}")
+    done
+    printf '%s' "$result"
+}
+
+echo "==> Generating CA"
+openssl req -x509 -newkey ec -pkeyopt ec_paramgen_curve:prime256v1 \
+    -keyout "$DIR/client-ca-key.pem" -out "$DIR/client-ca.pem" \
+    -days 365 -nodes -subj "/CN=agent-gateway CA" 2>/dev/null
+
+echo "==> Generating proxy server cert (localhost / 127.0.0.1)"
+openssl req -newkey ec -pkeyopt ec_paramgen_curve:prime256v1 \
+    -keyout "$DIR/proxy-key.pem" -out "$DIR/proxy.csr" \
+    -nodes -subj "/CN=localhost" 2>/dev/null
+openssl x509 -req -in "$DIR/proxy.csr" \
+    -CA "$DIR/client-ca.pem" -CAkey "$DIR/client-ca-key.pem" -CAcreateserial \
+    -out "$DIR/proxy.pem" -days 365 \
+    -extfile <(printf 'subjectAltName=DNS:localhost,IP:127.0.0.1') 2>/dev/null
+rm -f "$DIR/proxy.csr"
+
+echo "==> Generating client cert (extension_value=$EXT_VALUE)"
+DER_HEX=$(der_utf8string "$EXT_VALUE")
+openssl req -newkey ec -pkeyopt ec_paramgen_curve:prime256v1 \
+    -keyout "$DIR/client-key.pem" -out "$DIR/client.csr" \
+    -nodes -subj "/CN=agent-client" 2>/dev/null
+openssl x509 -req -in "$DIR/client.csr" \
+    -CA "$DIR/client-ca.pem" -CAkey "$DIR/client-ca-key.pem" -CAcreateserial \
+    -out "$DIR/client.pem" -days 365 \
+    -extfile <(printf '1.3.6.1.4.1.57264.1.1=DER:%s' "$DER_HEX") 2>/dev/null
+rm -f "$DIR/client.csr" "$DIR/client-ca.srl"
+
+echo ""
+echo "Generated in $DIR/:"
+ls -1 "$DIR"
+echo ""
+echo "Client extension: OID 1.3.6.1.4.1.57264.1.1 = \"$EXT_VALUE\""

examples/sample_client.py

@@ -84,8 +84,24 @@ def send_connect_request(sock, destination):
 }
 
 
+def _ssl_pending(sock):
+    """Bytes already decrypted and waiting in the OpenSSL read buffer (may be non-empty)."""
+    return sock.pending() if isinstance(sock, ssl.SSLSocket) else 0
+
+
 def wait_for_response(sock, conn, stream_id):
+    """Wait for CONNECT response.
+
+    Uses SSLSocket.pending(): we only count fast recv iterations (draining the OpenSSL read
+    buffer without blocking on the network) toward a livelock limit.
+    """
+    fast = 0
+    max_fast = 100_000
     while True:
+        pending_before = _ssl_pending(sock)
+        if pending_before == 0:
+            fast = 0
+
         data = sock.recv(65535)
         if not data:
             raise ConnectionError("Proxy closed the connection before responding")
@@ -94,6 +110,10 @@ def wait_for_response(sock, conn, stream_id):
         sock.sendall(conn.data_to_send())
 
         for event in events:
+            if isinstance(event, h2.events.ConnectionTerminated):
+                raise ConnectionError(
+                    f"Proxy closed HTTP/2 connection: error_code={event.error_code!r}"
+                )
             if isinstance(event, h2.events.ResponseReceived) and event.stream_id == stream_id:
                 headers = dict(event.headers)
                 status = int(headers[":status"])
@@ -102,25 +122,36 @@ def wait_for_response(sock, conn, stream_id):
             if isinstance(event, h2.events.StreamReset):
                 raise ConnectionError(f"Stream reset by proxy: error code {event.error_code}")
 
+        if pending_before > 0:
+            fast += 1
+            if fast >= max_fast:
+                raise ConnectionError(
+                    "Timed out waiting for CONNECT response (excessive HTTP/2 traffic without response)"
+                )
+
 
 class H2Tunnel:
     """Bidirectional byte pipe over an HTTP/2 DATA stream."""
 
+    # If we keep receiving TLS/plaintext without ever getting DATA for this stream (or EOF),
+    # something is wrong — bail out instead of burning CPU forever.
+    _MAX_FAST_RECV_WITHOUT_APP_DATA = 100_000
+
     def __init__(self, sock, h2_conn, stream_id):
         self._sock = sock
         self._conn = h2_conn
         self._sid = stream_id
         self._buffer = b""
         self._eof = False
 
-    def _read_from_tunnel(self):
-        raw = self._sock.recv(65535)
-        if not raw:
-            self._eof = True
-            return
+    def _process_raw(self, raw):
         events = self._conn.receive_data(raw)
         self._sock.sendall(self._conn.data_to_send())
         for ev in events:
+            if isinstance(ev, h2.events.ConnectionTerminated):
+                raise ConnectionError(
+                    f"Proxy closed HTTP/2 connection: error_code={ev.error_code!r}"
+                )
             if isinstance(ev, h2.events.DataReceived) and ev.stream_id == self._sid:
                 self._conn.acknowledge_received_data(ev.flow_controlled_length, ev.stream_id)
                 self._sock.sendall(self._conn.data_to_send())
@@ -129,8 +160,26 @@ def _read_from_tunnel(self):
                 self._eof = True
 
     def recv(self, bufsize):
+        fast = 0
         while not self._buffer and not self._eof:
-            self._read_from_tunnel()
+            pending_before = _ssl_pending(self._sock)
+            if pending_before == 0:
+                fast = 0
+
+            raw = self._sock.recv(65535)
+            if not raw:
+                self._eof = True
+                break
+
+            self._process_raw(raw)
+
+            if not self._buffer and not self._eof and pending_before > 0:
+                fast += 1
+                if fast >= self._MAX_FAST_RECV_WITHOUT_APP_DATA:
+                    raise ConnectionError(
+                        "HTTP/2 tunnel livelock: received many frames but no payload for CONNECT stream"
+                    )
+
         if not self._buffer:
             return b""
         out = self._buffer[:bufsize]
@@ -167,6 +216,10 @@ def tunnel_tls_handshake(tunnel, dest_host, dest_ca=None):
             if not data:
                 raise ConnectionError(f"tunnel closed during TLS handshake with {dest_host}")
             incoming_bio.write(data)
+        except ssl.SSLWantWriteError:
+            out = outgoing_bio.read()
+            if out:
+                tunnel.send(out)
 
     out = outgoing_bio.read()
     if out:
@@ -175,8 +228,25 @@ def tunnel_tls_handshake(tunnel, dest_host, dest_ca=None):
     return ssl_obj, incoming_bio, outgoing_bio
 
 
-def ssl_send(ssl_obj, outgoing_bio, tunnel, data):
-    ssl_obj.write(data)
+def ssl_send(ssl_obj, incoming_bio, outgoing_bio, tunnel, data):
+    view = memoryview(data)
+    while len(view):
+        try:
+            n = ssl_obj.write(view)
+            view = view[n:]
+        except ssl.SSLWantReadError:
+            out = outgoing_bio.read()
+            if out:
+                tunnel.send(out)
+            chunk = tunnel.recv(65535)
+            if not chunk:
+                raise ConnectionError("tunnel closed while sending to destination TLS")
+            incoming_bio.write(chunk)
+        except ssl.SSLWantWriteError:
+            out = outgoing_bio.read()
+            if out:
+                tunnel.send(out)
+
     out = outgoing_bio.read()
     if out:
         tunnel.send(out)
@@ -200,6 +270,10 @@ def ssl_recv(ssl_obj, incoming_bio, outgoing_bio, tunnel, bufsize=4096):
             if not data:
                 break
             incoming_bio.write(data)
+        except ssl.SSLWantWriteError:
+            out = outgoing_bio.read()
+            if out:
+                tunnel.send(out)
         except ssl.SSLZeroReturnError:
             break
     return b"".join(chunks)
@@ -212,7 +286,7 @@ def send_get_request(ssl_obj, incoming_bio, outgoing_bio, tunnel, host):
         f"Connection: close\r\n"
         f"\r\n"
     ).encode()
-    ssl_send(ssl_obj, outgoing_bio, tunnel, request)
+    ssl_send(ssl_obj, incoming_bio, outgoing_bio, tunnel, request)
 
     response_parts = []
     while True: