[rigging] Cleanup: drop dead JwtTokenManager wrappers, simplify resolver API

- Rename `vm_address(name, provider="gcp", ...)` → `gcp_vm_address(name, ...)`,
  drop the speculative provider arg.
- Lazy-import `google.auth`/`httpx` inside their only call sites so importing
  `iris.client` no longer pulls them in eagerly.
- Add `is_registered(scheme)` so callers (and tests) don't have to peek at
  `_HANDLERS`.
- Replace the uvicorn-based `test_resolver_plugin.py` integration test with
  an in-process unit test that monkeypatches `ControllerServiceClientSync`.
- Drop `JwtTokenManager.signing_key` and `JwtTokenManager.verifier`
  properties; surface the signing key on `ControllerAuth` instead so the
  log-server subprocess wiring goes through one explicit field.
- Trim the `VerifiedIdentity` re-export comment in `iris/rpc/auth.py`.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: rjpower

lib/iris/src/iris/client/resolver_plugin.py

@@ -3,7 +3,7 @@
 
 """Registers ``iris://<cluster>?endpoint=<name>`` with ``rigging.resolver``."""
 
-from rigging.resolver import ServiceURL, register_scheme, vm_address
+from rigging.resolver import ServiceURL, gcp_vm_address, register_scheme
 
 from iris.rpc.controller_connect import ControllerServiceClientSync
 from iris.rpc.controller_pb2 import Controller as _Controller
@@ -16,7 +16,7 @@ def _resolve_iris(url: ServiceURL) -> tuple[str, int]:
     name = url.query.get("endpoint")
     if not name:
         raise ValueError(f"iris:// URL requires ?endpoint=<name>: {url!r}")
-    controller_host, controller_port = vm_address(f"iris-controller-{cluster}", provider="gcp")
+    controller_host, controller_port = gcp_vm_address(f"iris-controller-{cluster}")
     with ControllerServiceClientSync(address=f"http://{controller_host}:{controller_port}") as client:
         response = client.list_endpoints(_Controller.ListEndpointsRequest(prefix=name, exact=True))
     if not response.endpoints:

lib/iris/src/iris/cluster/controller/auth.py

@@ -181,22 +181,12 @@ class JwtTokenManager:
     """
 
     def __init__(self, signing_key: str, db: ControllerDB | None = None):
+        self._signing_key = signing_key
         self._verifier = JwtVerifier(signing_key)
         self._db = db
         # Tracks the last wall-clock time we wrote last_used_at per jti.
         self._last_touched: dict[str, float] = {}
 
-    @property
-    def signing_key(self) -> str:
-        """HMAC secret used to sign and verify JWTs. Do not log or serialize."""
-        return self._verifier.signing_key
-
-    @property
-    def verifier(self) -> JwtVerifier:
-        """Underlying stateless verifier, suitable for handing to a log server
-        or other process that should validate but not issue tokens."""
-        return self._verifier
-
     def create_token(
         self,
         user_id: str,
@@ -212,7 +202,7 @@ def create_token(
             "iat": int(now),
             "exp": int(now + ttl_seconds),
         }
-        return jwt.encode(payload, self._verifier.signing_key, algorithm=JWT_ALGORITHM)
+        return jwt.encode(payload, self._signing_key, algorithm=JWT_ALGORITHM)
 
     def verify(self, token: str) -> VerifiedIdentity:
         """Verify JWT signature and claims, check revocation.
@@ -276,6 +266,9 @@ class ControllerAuth:
     login_verifier: TokenVerifier | None = None
     gcp_project_id: str | None = None
     jwt_manager: JwtTokenManager | None = None
+    # HMAC signing key — handed to the log-server subprocess via env var. Do
+    # not log or serialize.
+    signing_key: str | None = None
     optional: bool = False
 
 
@@ -301,7 +294,12 @@ def create_controller_auth(
 
             worker_token = _create_worker_jwt(db, jwt_mgr, now)
             logger.info("Authentication disabled — null-auth mode (workers use JWT)")
-            return ControllerAuth(verifier=jwt_mgr, worker_token=worker_token, jwt_manager=jwt_mgr)
+            return ControllerAuth(
+                verifier=jwt_mgr,
+                worker_token=worker_token,
+                jwt_manager=jwt_mgr,
+                signing_key=signing_key,
+            )
         logger.info("Authentication disabled — null-auth mode, no DB")
         return ControllerAuth()
 
@@ -327,8 +325,8 @@ def create_controller_auth(
 
         verifier: TokenVerifier | None = jwt_mgr
     else:
-        ephemeral_key = secrets.token_hex(32)
-        jwt_mgr = JwtTokenManager(ephemeral_key)
+        signing_key = secrets.token_hex(32)
+        jwt_mgr = JwtTokenManager(signing_key)
         worker_token = jwt_mgr.create_token(WORKER_USER, "worker", f"iris_k_worker_{secrets.token_hex(8)}")
         verifier = None
 
@@ -361,6 +359,7 @@ def create_controller_auth(
         login_verifier=login_verifier,
         gcp_project_id=gcp_project_id,
         jwt_manager=jwt_mgr,
+        signing_key=signing_key,
         optional=optional,
     )
 

lib/iris/src/iris/cluster/controller/main.py

@@ -231,7 +231,7 @@ def run_controller_serve(
         port=log_port,
         log_dir=log_dir,
         remote_log_dir=remote_log_dir,
-        signing_key=auth.jwt_manager.signing_key if auth.jwt_manager else None,
+        signing_key=auth.signing_key,
         strict_auth=auth.provider is not None,
     )
     try:

lib/iris/src/iris/rpc/auth.py

@@ -30,12 +30,6 @@
 SESSION_COOKIE = "iris_session"
 
 
-# `VerifiedIdentity` is defined in `rigging.auth` so the stateless
-# `JwtVerifier` (also in rigging) and the iris-side issuer share the same
-# type. Re-exported here for the many iris call sites that import
-# `from iris.rpc.auth import VerifiedIdentity`.
-
-
 def _extract_cookie(cookie_header: str, name: str) -> str | None:
     """Extract a named cookie value from a raw Cookie header."""
     if not cookie_header:

lib/iris/tests/client/test_resolver_plugin.py

@@ -1,146 +1,89 @@
 # Copyright The Marin Authors
 # SPDX-License-Identifier: Apache-2.0
 
-"""End-to-end test for the iris resolver plugin."""
-
-import socket
-import threading
-from collections.abc import Iterator
-from typing import Any
+"""Unit tests for the iris:// resolver plugin."""
 
 import pytest
-import uvicorn
-from starlette.applications import Starlette
-from starlette.middleware.wsgi import WSGIMiddleware
-from starlette.routing import Mount
 
 import iris.client  # noqa: F401  -- side-effect import: registers iris:// scheme
+from iris.client import resolver_plugin
 from iris.rpc import controller_pb2
-from iris.rpc.controller_connect import ControllerServiceSync, ControllerServiceWSGIApplication
-from rigging import resolver as resolver_module
-from rigging.resolver import resolve
-from rigging.timing import Duration, ExponentialBackoff
+from rigging.resolver import is_registered, resolve
+
 
+class _FakeControllerClient:
+    """Stubs ControllerServiceClientSync + its context manager."""
 
-class _StubControllerService(ControllerServiceSync):
-    """Minimal ``ControllerServiceSync`` that only implements ``list_endpoints``.
+    def __init__(self, endpoints: dict[str, str]):
+        self._endpoints = endpoints
+        self.last_request: controller_pb2.Controller.ListEndpointsRequest | None = None
 
-    Inherits the Protocol base class so all other RPCs default to
-    ``UNIMPLEMENTED`` errors, exactly what we want for an isolated test.
-    """
+    def __enter__(self) -> "_FakeControllerClient":
+        return self
 
-    def __init__(self) -> None:
-        self.endpoints: dict[str, str] = {}
+    def __exit__(self, *_exc) -> None:
+        return None
 
     def list_endpoints(
         self,
         request: controller_pb2.Controller.ListEndpointsRequest,
-        ctx: Any,
     ) -> controller_pb2.Controller.ListEndpointsResponse:
-        results: list[controller_pb2.Controller.Endpoint] = []
-        for name, address in self.endpoints.items():
-            if request.exact:
-                if name == request.prefix:
-                    results.append(controller_pb2.Controller.Endpoint(name=name, address=address))
-            else:
-                if name.startswith(request.prefix):
-                    results.append(controller_pb2.Controller.Endpoint(name=name, address=address))
-        return controller_pb2.Controller.ListEndpointsResponse(endpoints=results)
-
-
-def _free_port() -> int:
-    with socket.socket() as s:
-        s.bind(("127.0.0.1", 0))
-        return s.getsockname()[1]
-
-
-def _build_app(service: _StubControllerService) -> Starlette:
-    wsgi = ControllerServiceWSGIApplication(service=service)
-    return Starlette(routes=[Mount(wsgi.path, app=WSGIMiddleware(wsgi))])
-
-
-class _BackgroundServer:
-    def __init__(self, app: Starlette, port: int) -> None:
-        config = uvicorn.Config(
-            app,
-            host="127.0.0.1",
-            port=port,
-            log_level="error",
-            log_config=None,
-            timeout_keep_alive=5,
-        )
-        self.server = uvicorn.Server(config)
-        self.port = port
-        self._thread = threading.Thread(
-            target=self.server.run,
-            name=f"resolver-plugin-test-{port}",
-            daemon=True,
-        )
-
-    def start(self) -> None:
-        self._thread.start()
-        started = ExponentialBackoff(initial=0.01, maximum=0.2).wait_until(
-            lambda: self.server.started,
-            timeout=Duration.from_seconds(5.0),
-        )
-        if not started:
-            raise RuntimeError(f"uvicorn did not start within 5s on port {self.port}")
-
-    def stop(self) -> None:
-        self.server.should_exit = True
-        self._thread.join(timeout=5.0)
+        self.last_request = request
+        matches = [
+            controller_pb2.Controller.Endpoint(name=n, address=a)
+            for n, a in self._endpoints.items()
+            if n == request.prefix
+        ]
+        return controller_pb2.Controller.ListEndpointsResponse(endpoints=matches)
 
 
 @pytest.fixture
-def stub_controller() -> Iterator[tuple[_StubControllerService, int]]:
-    svc = _StubControllerService()
-    port = _free_port()
-    bg = _BackgroundServer(_build_app(svc), port)
-    bg.start()
-    try:
-        yield svc, port
-    finally:
-        bg.stop()
-
-
-def test_resolve_iris_round_trips(monkeypatch, stub_controller):
-    svc, controller_port = stub_controller
-    svc.endpoints["/system/x"] = "host.example.com:1234"
-
-    captured: list[tuple] = []
+def patch_resolver(monkeypatch):
+    """Replace gcp_vm_address + controller client with in-process stubs."""
+    vm_calls: list[str] = []
+
+    def _install(endpoints: dict[str, str]) -> _FakeControllerClient:
+        fake = _FakeControllerClient(endpoints)
+
+        def _fake_vm_address(name: str, *, port: int = 10002) -> tuple[str, int]:
+            vm_calls.append(name)
+            return ("127.0.0.1", 65000)
+
+        monkeypatch.setattr(resolver_plugin, "gcp_vm_address", _fake_vm_address)
+        monkeypatch.setattr(
+            resolver_plugin,
+            "ControllerServiceClientSync",
+            lambda address: fake,
+        )
+        return fake
 
-    def _fake_vm_address(name: str, provider: str) -> tuple[str, int]:
-        captured.append((name, provider))
-        # Direct test traffic at the in-process stub rather than GCP.
-        return ("127.0.0.1", controller_port)
+    _install.vm_calls = vm_calls  # type: ignore[attr-defined]
+    return _install
 
-    # The plugin binds vm_address as a module-global at import time; patch
-    # it where it's looked up.
-    from iris.client import resolver_plugin
 
-    monkeypatch.setattr(resolver_plugin, "vm_address", _fake_vm_address)
+def test_resolve_iris_returns_endpoint_address(patch_resolver):
+    patch_resolver({"/system/x": "host.example.com:1234"})
+    assert resolve("iris://marin?endpoint=/system/x") == ("host.example.com", 1234)
+    assert patch_resolver.vm_calls == ["iris-controller-marin"]
 
-    host, port = resolve("iris://marin?endpoint=/system/x")
-    assert (host, port) == ("host.example.com", 1234)
-    assert captured == [("iris-controller-marin", "gcp")]
 
+def test_resolve_iris_not_found_raises(patch_resolver):
+    patch_resolver({})
+    with pytest.raises(KeyError, match="iris endpoint not found"):
+        resolve("iris://marin?endpoint=/system/missing")
 
-def test_resolve_iris_not_found(monkeypatch, stub_controller):
-    _svc, controller_port = stub_controller
 
-    from iris.client import resolver_plugin
+def test_resolve_iris_requires_endpoint_query(patch_resolver):
+    patch_resolver({})
+    with pytest.raises(ValueError, match="requires \\?endpoint="):
+        resolve("iris://marin")
 
-    monkeypatch.setattr(
-        resolver_plugin,
-        "vm_address",
-        lambda name, provider: ("127.0.0.1", controller_port),
-    )
 
-    with pytest.raises(KeyError, match="iris endpoint not found"):
-        resolve("iris://marin?endpoint=/system/missing")
+def test_resolve_iris_rejects_port(patch_resolver):
+    patch_resolver({})
+    with pytest.raises(ValueError, match="cannot have a port"):
+        resolve("iris://marin:9000?endpoint=/x")
 
 
 def test_iris_scheme_registered_after_iris_client_import():
-    # Sanity check: importing iris.client (done at the top of this module)
-    # installs the iris:// handler in rigging.resolver's registry.
-    assert "iris" in resolver_module._HANDLERS
+    assert is_registered("iris")

lib/rigging/src/rigging/resolver.py

@@ -12,10 +12,6 @@
 from dataclasses import dataclass
 from urllib.parse import parse_qs, urlsplit
 
-import google.auth
-import google.auth.transport.requests
-import httpx
-
 _COMPUTE_BASE = "https://compute.googleapis.com/compute/v1"
 _OAUTH_SCOPE = "https://www.googleapis.com/auth/cloud-platform"
 _TIMEOUT_SECONDS = 10.0
@@ -52,6 +48,10 @@ def register_scheme(scheme: str, handler: SchemeHandler) -> None:
     _HANDLERS[scheme] = handler
 
 
+def is_registered(scheme: str) -> bool:
+    return scheme in _HANDLERS
+
+
 def resolve(ref: str) -> tuple[str, int]:
     if "://" not in ref:
         host, port = ref.rsplit(":", 1)
@@ -63,9 +63,7 @@ def resolve(ref: str) -> tuple[str, int]:
     return handler(url)
 
 
-def vm_address(name: str, provider: str, *, port: int = _DEFAULT_GCP_PORT) -> tuple[str, int]:
-    if provider != "gcp":
-        raise ValueError(f"unsupported provider: {provider}")
+def gcp_vm_address(name: str, *, port: int = _DEFAULT_GCP_PORT) -> tuple[str, int]:
     return _gcp_internal_ip(name), port
 
 
@@ -84,6 +82,9 @@ def _gcp_internal_ip(name: str) -> str:
 
 
 def _gcp_credentials() -> tuple[str, str]:
+    import google.auth
+    import google.auth.transport.requests
+
     creds, project_id = google.auth.default(scopes=[_OAUTH_SCOPE])
     if not project_id:
         raise LookupError("google.auth.default() returned no project_id; set GOOGLE_CLOUD_PROJECT")
@@ -92,6 +93,8 @@ def _gcp_credentials() -> tuple[str, str]:
 
 
 def _fetch_vm_aggregated(project_id: str, token: str, name: str) -> dict | None:
+    import httpx
+
     url = f"{_COMPUTE_BASE}/projects/{project_id}/aggregated/instances"
     headers = {"Authorization": f"Bearer {token}"}
     params: dict[str, str] = {"filter": f"name eq {name}"}
@@ -110,4 +113,4 @@ def _fetch_vm_aggregated(project_id: str, token: str, name: str) -> dict | None:
             params["pageToken"] = page_token
 
 
-register_scheme("gcp", lambda url: vm_address(url.host, "gcp", port=url.port or _DEFAULT_GCP_PORT))
+register_scheme("gcp", lambda url: gcp_vm_address(url.host, port=url.port or _DEFAULT_GCP_PORT))