Merge pull request 'feat: optional mTLS for felt-core only (WP-62 M2b B5.2)' (#1) from benylin/wp62-m2b-b52-mtls into main

Reviewed-on: Benylin/drive-rag#1

.gitea/workflows/ci.yml

@@ -15,4 +15,6 @@ jobs:
           python-version: "3.12"
       - run: uv sync --frozen --group server
       # no unit tests yet: the app must at least import (model + DB load in lifespan)
-      - run: uv run python -c "import server.app"
+      - run: uv run python -c "import server.app, server.serve"
+      # WP-62 M2b B5.2: the mTLS server against a real test PKI (needs the openssl CLI)
+      - run: uv run --with pytest python -m pytest -v tests/

Dockerfile

@@ -18,5 +18,5 @@ RUN uv run python -c "from model2vec import StaticModel; StaticModel.from_pretra
 COPY server/ ./server/
 
 EXPOSE 8000
-CMD ["uv", "run", "--no-dev", "--group", "server", "uvicorn", "server.app:app", \
-     "--host", "0.0.0.0", "--port", "8000"]
+# plain uvicorn, or mTLS when DRIVE_RAG_TLS_BUNDLE is set (server/serve.py, WP-62 M2b)
+CMD ["uv", "run", "--no-dev", "--group", "server", "python", "-m", "server.serve"]

server/mtls.py

@@ -0,0 +1,108 @@
+"""mTLS for drive-rag's API (WP-62 M2b B5.2, decision D10.2).
+
+Same module as felt-voice's (felt_voice/mtls.py, B5.1): felt-core is the only
+caller. She presents her 2-hour Felt client cert; drive-rag presents its own
+(``drive-rag.felt.internal``). Both certs come from clockmost's server-side
+helper as ``bundle.pem`` (key + cert + intermediate) next to ``felt-ca.pem``
+(the pki_felt intermediate: the ONLY trust anchor - never the clockmost root).
+
+- The TLS handshake requires a client cert that chains to ``felt-ca.pem``.
+- Every connection's peer must carry exactly one SPIFFE URI and it must be in
+  ``allowed`` (felt-core's), checked before a single request byte is read.
+- The server cert is re-read when the bundle changes (renewed hourly).
+
+Off unless DRIVE_RAG_TLS_BUNDLE is set (plain HTTP as before).
+"""
+
+import logging
+import os
+import ssl
+import threading
+
+import uvicorn
+from uvicorn.protocols.http.auto import AutoHTTPProtocol
+
+log = logging.getLogger(__name__)
+
+
+def peer_uris(transport) -> list[str]:
+    cert = transport.get_extra_info("peercert") or {}
+    return [v for k, v in cert.get("subjectAltName", ()) if k == "URI"]
+
+
+def make_protocol(allowed: frozenset[str]):
+    """The HTTP protocol, refusing any peer whose identity isn't allowed."""
+
+    class MTLSProtocol(AutoHTTPProtocol):
+        def connection_made(self, transport):  # after the TLS handshake
+            uris = peer_uris(transport)
+            if len(uris) != 1 or uris[0] not in allowed:
+                log.warning(
+                    "refused TLS client %s: identity %s is not allowed",
+                    transport.get_extra_info("peername"),
+                    uris or "none",
+                )
+                self._refused = True
+                transport.abort()
+                return
+            super().connection_made(transport)
+
+        def connection_lost(self, exc):
+            if getattr(self, "_refused", False):  # nothing was set up to tear down
+                return
+            super().connection_lost(exc)
+
+    return MTLSProtocol
+
+
+def harden(ctx: ssl.SSLContext) -> None:
+    ctx.minimum_version = ssl.TLSVersion.TLSv1_3
+    ctx.verify_mode = ssl.CERT_REQUIRED
+    ctx.verify_flags |= ssl.VERIFY_X509_PARTIAL_CHAIN  # felt-ca.pem is an intermediate
+
+
+def watch_bundle(
+    ctx: ssl.SSLContext, bundle: str, interval: float, stop: threading.Event
+) -> None:
+    """Reload the server cert into the live context whenever the renewer
+    replaces the bundle."""
+    seen = os.stat(bundle).st_mtime_ns
+    while not stop.wait(interval):
+        try:
+            now = os.stat(bundle).st_mtime_ns
+            if now != seen:
+                ctx.load_cert_chain(bundle)
+                seen = now
+                log.info("reloaded the server certificate from %s", bundle)
+        except (OSError, ssl.SSLError) as e:
+            log.error(
+                "server certificate reload failed (keeping the current one): %s", e
+            )
+
+
+def serve(
+    app, host: str, port: int, bundle: str, ca: str, allowed, reload_s: float = 60
+) -> None:
+    config = uvicorn.Config(
+        app,
+        host=host,
+        port=port,
+        ssl_certfile=bundle,
+        ssl_keyfile=bundle,
+        ssl_ca_certs=ca,
+        ssl_cert_reqs=ssl.CERT_REQUIRED,
+        http=make_protocol(frozenset(allowed)),
+    )
+    config.load()
+    harden(config.ssl)
+    stop = threading.Event()
+    threading.Thread(
+        target=watch_bundle, args=(config.ssl, bundle, reload_s, stop), daemon=True
+    ).start()
+    log.info(
+        "mTLS on %s:%s, allowed clients: %s", host, port, ", ".join(sorted(allowed))
+    )
+    try:
+        uvicorn.Server(config).run()
+    finally:
+        stop.set()

server/serve.py

@@ -0,0 +1,37 @@
+"""drive-rag's entry point: plain uvicorn, or mTLS when DRIVE_RAG_TLS_BUNDLE is
+set (WP-62 M2b B5.2). Settings from the environment:
+
+  DRIVE_RAG_TLS_BUNDLE   server bundle.pem (key + cert + intermediate)
+  DRIVE_RAG_TLS_CA       felt-ca.pem (the pki_felt intermediate)
+  DRIVE_RAG_TLS_ALLOWED  allowed client SPIFFE URIs, comma-separated
+                         (default: felt-core)
+"""
+
+import logging
+import os
+
+import uvicorn
+
+from server.app import app
+
+FELT_CORE = "spiffe://clockmost.internal/workload/felt-core"
+
+
+def main() -> None:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
+    host, port = "0.0.0.0", 8000
+    bundle = os.environ.get("DRIVE_RAG_TLS_BUNDLE")
+    if not bundle:
+        uvicorn.run(app, host=host, port=port)
+        return
+    ca = os.environ.get("DRIVE_RAG_TLS_CA") or ""
+    if not ca:
+        raise SystemExit("DRIVE_RAG_TLS_BUNDLE is set but DRIVE_RAG_TLS_CA isn't")
+    allowed = [u.strip() for u in os.environ.get("DRIVE_RAG_TLS_ALLOWED", FELT_CORE).split(",") if u.strip()]
+    from server import mtls
+
+    mtls.serve(app, host, port, bundle, ca, allowed)
+
+
+if __name__ == "__main__":
+    main()

tests/test_mtls.py

@@ -0,0 +1,244 @@
+"""server.mtls against a real test PKI shaped like clockmost's (root ->
+pki_felt-style intermediate with the D6 name constraints -> leaves)."""
+
+import shutil
+import socket
+import ssl
+import subprocess
+import threading
+import time
+
+import pytest
+
+from server import mtls
+
+pytestmark = pytest.mark.skipif(
+    shutil.which("openssl") is None, reason="needs the openssl CLI"
+)
+FELT_CORE = "spiffe://clockmost.internal/workload/felt-core"
+
+
+def _openssl(*a, cwd):
+    subprocess.run(["openssl", *a], cwd=cwd, check=True, capture_output=True)
+
+
+def _leaf(d, name, issuer, ext):
+    _openssl(
+        "req",
+        "-new",
+        "-newkey",
+        "ec",
+        "-pkeyopt",
+        "ec_paramgen_curve:P-256",
+        "-nodes",
+        "-keyout",
+        f"{name}.key",
+        "-subj",
+        f"/CN={name}",
+        "-out",
+        f"{name}.csr",
+        cwd=d,
+    )
+    (d / f"{name}.ext").write_text(ext)
+    _openssl(
+        "x509",
+        "-req",
+        "-in",
+        f"{name}.csr",
+        "-CA",
+        f"{issuer}.pem",
+        "-CAkey",
+        f"{issuer}.key",
+        "-CAcreateserial",
+        "-days",
+        "1",
+        "-extfile",
+        f"{name}.ext",
+        "-out",
+        f"{name}.pem",
+        cwd=d,
+    )
+
+
+@pytest.fixture(scope="module")
+def pki(tmp_path_factory):
+    d = tmp_path_factory.mktemp("pki")
+    for ca in ("root", "wrong"):
+        _openssl(
+            "req",
+            "-x509",
+            "-newkey",
+            "ec",
+            "-pkeyopt",
+            "ec_paramgen_curve:P-256",
+            "-nodes",
+            "-keyout",
+            f"{ca}.key",
+            "-subj",
+            f"/CN={ca}",
+            "-days",
+            "2",
+            "-out",
+            f"{ca}.pem",
+            cwd=d,
+        )
+    _leaf(
+        d,
+        "feltint",
+        "root",
+        "basicConstraints=critical,CA:true,pathlen:0\nkeyUsage=critical,keyCertSign,cRLSign\n"
+        "nameConstraints=critical,permitted;URI:clockmost.internal,permitted;DNS:felt.internal,"
+        "excluded;IP:0.0.0.0/0.0.0.0,excluded;IP:::/::\n",
+    )
+    cli = "basicConstraints=critical,CA:false\nextendedKeyUsage=clientAuth\nsubjectAltName="
+    srv = "basicConstraints=critical,CA:false\nextendedKeyUsage=serverAuth\nsubjectAltName="
+    _leaf(
+        d,
+        "voice",
+        "feltint",
+        srv
+        + "DNS:drive-rag.felt.internal,URI:spiffe://clockmost.internal/workload/drive-rag\n",
+    )
+    _leaf(
+        d,
+        "voice2",
+        "feltint",
+        srv
+        + "DNS:drive-rag.felt.internal,URI:spiffe://clockmost.internal/workload/drive-rag\n",
+    )
+    _leaf(d, "core", "feltint", cli + f"DNS:felt-core.felt.internal,URI:{FELT_CORE}\n")
+    _leaf(
+        d,
+        "tray",
+        "feltint",
+        cli + "DNS:tray.felt.internal,URI:spiffe://clockmost.internal/workload/tray\n",
+    )
+    _leaf(d, "rogue", "wrong", cli + f"URI:{FELT_CORE}\n")
+    for n, iss in (
+        ("voice", "feltint"),
+        ("voice2", "feltint"),
+        ("core", "feltint"),
+        ("tray", "feltint"),
+        ("rogue", "wrong"),
+    ):
+        (d / f"{n}.bundle").write_text(
+            (d / f"{n}.key").read_text()
+            + (d / f"{n}.pem").read_text()
+            + (d / f"{iss}.pem").read_text()
+        )
+    return d
+
+
+async def _app(scope, receive, send):
+    if scope["type"] != "http":
+        return
+    await send(
+        {
+            "type": "http.response.start",
+            "status": 200,
+            "headers": [(b"content-length", b"2")],
+        }
+    )
+    await send({"type": "http.response.body", "body": b"ok"})
+
+
+@pytest.fixture(scope="module")
+def server(pki):
+    s = socket.socket()
+    s.bind(("127.0.0.1", 0))
+    port = s.getsockname()[1]
+    s.close()
+    live = pki / "live.bundle"
+    shutil.copy(pki / "voice.bundle", live)
+    threading.Thread(
+        target=mtls.serve,
+        args=(
+            _app,
+            "127.0.0.1",
+            port,
+            str(live),
+            str(pki / "feltint.pem"),
+            [FELT_CORE],
+            0.2,
+        ),
+        daemon=True,
+    ).start()
+    for _ in range(50):
+        try:
+            socket.create_connection(("127.0.0.1", port), 0.2).close()
+            break
+        except OSError:
+            time.sleep(0.1)
+    return port, live
+
+
+def _get(pki, port, bundle=None, max_tls=None):
+    """GET / as felt-core would: trust felt-ca only, verify the server by name."""
+    ctx = ssl.create_default_context(cafile=str(pki / "feltint.pem"))
+    ctx.verify_flags |= ssl.VERIFY_X509_PARTIAL_CHAIN
+    if max_tls:
+        ctx.maximum_version = max_tls
+    if bundle:
+        ctx.load_cert_chain(str(pki / bundle))
+    with ctx.wrap_socket(
+        socket.create_connection(("127.0.0.1", port), 5),
+        server_hostname="drive-rag.felt.internal",
+    ) as s:
+        cert = s.getpeercert()
+        s.sendall(
+            b"GET / HTTP/1.1\r\nHost: drive-rag.felt.internal\r\nConnection: close\r\n\r\n"
+        )
+        return s.recv(200), cert
+
+
+def test_felt_core_gets_through(pki, server):
+    data, _ = _get(pki, server[0], "core.bundle")
+    assert data.startswith(b"HTTP/1.1 200")
+
+
+def test_no_client_cert_refused(pki, server):
+    try:
+        data, _ = _get(pki, server[0])
+    except (ssl.SSLError, ConnectionError):
+        return
+    assert data == b""  # TLS 1.3: the server's refusal can arrive as a closed read
+
+
+def test_other_felt_identity_refused(pki, server):
+    try:
+        data, _ = _get(pki, server[0], "tray.bundle")
+    except (ssl.SSLError, ConnectionError):
+        return
+    assert (
+        data == b""
+    )  # handshake passes (valid Felt cert), connection dropped before any response
+
+
+def test_wrong_ca_refused(pki, server):
+    try:
+        data, _ = _get(pki, server[0], "rogue.bundle")
+    except (ssl.SSLError, ConnectionError):
+        return
+    assert data == b""
+
+
+def test_tls12_refused(pki, server):
+    with pytest.raises(ssl.SSLError):
+        _get(pki, server[0], "core.bundle", max_tls=ssl.TLSVersion.TLSv1_2)
+
+
+def test_renewed_cert_is_picked_up(pki, server):
+    port, live = server
+    _, before = _get(pki, port, "core.bundle")
+    shutil.copy(pki / "voice2.bundle", live)
+    time.sleep(1.0)
+    _, after = _get(pki, port, "core.bundle")
+    assert before["serialNumber"] != after["serialNumber"]
+
+
+def test_refusal_logs_no_traceback(pki, server, caplog):
+    """A refused peer is one warning line, not an asyncio callback traceback."""
+    test_other_felt_identity_refused(pki, server)
+    time.sleep(0.3)
+    assert "is not allowed" in caplog.text
+    assert "Exception in callback" not in caplog.text