fix(connect): proxy OIDC token requests and close refresh sessions

mock_tests/test_auth.py

@@ -6,6 +6,7 @@
 from typing import List, Union
 
 import grpc
+import httpx
 import pytest
 from pytest_httpserver import HTTPServer
 from werkzeug import Request, Response
@@ -85,6 +86,203 @@ def test_client_credentials(weaviate_auth_mock: HTTPServer, start_grpc_server: g
     weaviate_auth_mock.check_assertions()
 
 
+OIDC_CREDENTIALS = {
+    "bearer_token": lambda: weaviate.auth.AuthBearerToken(
+        ACCESS_TOKEN, refresh_token=REFRESH_TOKEN
+    ),
+    "client_password": lambda: weaviate.auth.AuthClientPassword("AUsername", "SomePassWord"),
+    "client_credentials": lambda: weaviate.auth.AuthClientCredentials(
+        client_secret=CLIENT_SECRET, scope=SCOPE
+    ),
+}
+
+
+def _expect_oidc_token(weaviate_auth_mock: HTTPServer) -> None:
+    weaviate_auth_mock.expect_request("/auth").respond_with_json(
+        {"access_token": ACCESS_TOKEN, "expires_in": 500, "refresh_token": REFRESH_TOKEN}
+    )
+    weaviate_auth_mock.expect_request(
+        "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN}
+    ).respond_with_json({"classes": []})
+
+
+@pytest.mark.parametrize("credentials", OIDC_CREDENTIALS.keys())
+def test_oidc_requests_use_own_http_client(
+    weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, credentials: str
+) -> None:
+    """Requests must go through the client's own httpx client, never the authlib session.
+
+    Authlib >= 1.8 builds its sessions on httpx2 when it is installed; handing such a
+    session the client's httpx Timeout failed every request with a TypeError (#2220).
+    """
+    _expect_oidc_token(weaviate_auth_mock)
+
+    with weaviate.connect_to_local(
+        host=MOCK_IP,
+        port=MOCK_PORT,
+        grpc_port=MOCK_PORT_GRPC,
+        auth_credentials=OIDC_CREDENTIALS[credentials](),
+    ) as client:
+        assert type(client._connection._client) is httpx.Client
+        client.collections.list_all()
+    assert client._connection._oauth_session is None  # close() closed the session too
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("credentials", OIDC_CREDENTIALS.keys())
+async def test_oidc_requests_use_own_http_client_async(
+    weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, credentials: str
+) -> None:
+    """Async twin of the test above."""
+    _expect_oidc_token(weaviate_auth_mock)
+
+    async with weaviate.use_async_with_local(
+        host=MOCK_IP,
+        port=MOCK_PORT,
+        grpc_port=MOCK_PORT_GRPC,
+        auth_credentials=OIDC_CREDENTIALS[credentials](),
+    ) as client:
+        assert type(client._connection._client) is httpx.AsyncClient
+        await client.collections.list_all()
+    assert client._connection._oauth_session is None  # close() closed the session too
+
+
+def _as_oidc_proxy(proxy: HTTPServer) -> None:
+    """Serve Weaviate and the IdP from a proxy, behind hosts only the proxy can reach."""
+    proxy.expect_request("/v1/meta").respond_with_json({"version": "1.36"})
+    proxy.expect_request("/v1/.well-known/openid-configuration").respond_with_json(
+        {"href": "http://idp.invalid/endpoints", "clientId": CLIENT_ID}
+    )
+    proxy.expect_request("/endpoints").respond_with_json(
+        {"token_endpoint": "http://idp.invalid/auth"}
+    )
+    proxy.expect_request("/auth", method="POST").respond_with_json(
+        {"access_token": ACCESS_TOKEN, "expires_in": 500, "refresh_token": REFRESH_TOKEN}
+    )
+    proxy.expect_request(
+        "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN}
+    ).respond_with_json({"classes": []})
+
+
+@pytest.fixture
+def oidc_proxy():
+    proxy = HTTPServer(host=MOCK_IP, port=0)
+    proxy.start()
+    _as_oidc_proxy(proxy)
+    yield proxy
+    proxy.clear()
+    if proxy.is_running():
+        proxy.stop()
+
+
+def _proxied_custom_kwargs(proxy: HTTPServer) -> dict:
+    return {
+        "http_host": "weaviate.invalid",
+        "http_port": 80,
+        "http_secure": False,
+        "grpc_host": MOCK_IP,
+        "grpc_port": MOCK_PORT_GRPC,
+        "grpc_secure": False,
+        "skip_init_checks": True,
+        "auth_credentials": weaviate.auth.AuthClientPassword("AUsername", "SomePassWord"),
+        "additional_config": weaviate.classes.init.AdditionalConfig(
+            proxies=weaviate.classes.init.Proxies(http=proxy.url_for("/"))
+        ),
+    }
+
+
+def test_oidc_token_requests_use_proxies(oidc_proxy: HTTPServer) -> None:
+    """Token requests go through the configured proxies, like every other request."""
+    with weaviate.connect_to_custom(**_proxied_custom_kwargs(oidc_proxy)) as client:
+        client.collections.list_all()
+    assert any(req.path == "/auth" for req, _ in oidc_proxy.log)
+
+
+@pytest.mark.asyncio
+async def test_oidc_token_requests_use_proxies_async(oidc_proxy: HTTPServer) -> None:
+    """Async twin of the test above."""
+    async with weaviate.use_async_with_custom(**_proxied_custom_kwargs(oidc_proxy)) as client:
+        await client.collections.list_all()
+    assert any(req.path == "/auth" for req, _ in oidc_proxy.log)
+
+
+def _track_sessions(monkeypatch: pytest.MonkeyPatch) -> list:
+    """Record every authlib session the client creates."""
+    sessions: list = []
+    for name in ("OAuth2Client", "AsyncOAuth2Client"):
+        base = getattr(weaviate.connect.authentication, name)
+        tracked = type(name, (base,), {})
+        original_init = base.__init__
+
+        def init(self, *args, __init=original_init, **kwargs) -> None:
+            __init(self, *args, **kwargs)
+            sessions.append(self)
+
+        tracked.__init__ = init
+        monkeypatch.setattr(weaviate.connect.authentication, name, tracked)
+    return sessions
+
+
+def _expect_short_lived_client_credentials(weaviate_auth_mock: HTTPServer) -> List[float]:
+    """Client-credentials tokens without a refresh token, expiring right away; returns the hit times."""
+    hits: List[float] = []
+
+    def handler(request: Request) -> Response:
+        hits.append(time.monotonic())
+        return Response(
+            json.dumps({"access_token": ACCESS_TOKEN, "expires_in": 1}),
+            content_type="application/json",
+        )
+
+    weaviate_auth_mock.expect_request("/auth").respond_with_handler(handler)
+    return hits
+
+
+def test_client_credentials_refresh_closes_sessions(
+    weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Each refresh builds a session for one token; it must be closed, and fetch the token once."""
+    sessions = _track_sessions(monkeypatch)
+    hits = _expect_short_lived_client_credentials(weaviate_auth_mock)
+
+    with weaviate.connect_to_local(
+        host=MOCK_IP,
+        port=MOCK_PORT,
+        grpc_port=MOCK_PORT_GRPC,
+        auth_credentials=weaviate.auth.AuthClientCredentials(
+            client_secret=CLIENT_SECRET, scope=SCOPE
+        ),
+    ):
+        time.sleep(3)  # refresh interval is max(expires_in - 30, 1) -> 1s
+
+    assert len(sessions) >= 3  # the refresher ran
+    assert len(hits) == len(sessions)  # one token request per session
+    assert all(session.is_closed for session in sessions)
+
+
+@pytest.mark.asyncio
+async def test_client_credentials_refresh_closes_sessions_async(
+    weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    """Async twin of the test above."""
+    sessions = _track_sessions(monkeypatch)
+    hits = _expect_short_lived_client_credentials(weaviate_auth_mock)
+
+    async with weaviate.use_async_with_local(
+        host=MOCK_IP,
+        port=MOCK_PORT,
+        grpc_port=MOCK_PORT_GRPC,
+        auth_credentials=weaviate.auth.AuthClientCredentials(
+            client_secret=CLIENT_SECRET, scope=SCOPE
+        ),
+    ):
+        await asyncio.sleep(3)
+
+    assert len(sessions) >= 3
+    assert len(hits) == len(sessions)
+    assert all(session.is_closed for session in sessions)
+
+
 @pytest.mark.asyncio
 async def test_client_credentials_refresh_async(
     weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server

weaviate/connect/authentication.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
-from typing import Awaitable, Callable, Dict, List, Optional, Union
+import importlib
+from types import ModuleType
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Union
 
 import httpx
 from authlib.integrations.httpx_client import (  # type: ignore
@@ -24,10 +26,24 @@
 OIDC_CONFIG = Dict[str, Union[str, List[str]]]
 
 Result = Union[OAuth2Client, Awaitable[AsyncOAuth2Client]]
-MountsMaker = Union[
-    Callable[[], Dict[str, httpx.AsyncHTTPTransport]],
-    Callable[[], Dict[str, httpx.HTTPTransport]],
-]
+# builds the proxy transports with the given library (httpx or httpx2)
+MountsMaker = Callable[[ModuleType], Dict[str, Any]]
+
+
+def _session_http_module() -> ModuleType:
+    """Return the library authlib builds its sessions on.
+
+    Authlib >= 1.8 uses httpx2 when it is installed, older versions httpx. A session's transports must come
+    from the same library.
+    """
+    for base in OAuth2Client.__mro__:
+        root = base.__module__.split(".")[0]
+        if root in ("httpx", "httpx2"):
+            return importlib.import_module(root)
+    return httpx
+
+
+_SESSION_HTTP = _session_http_module()
 
 
 class _Auth:
@@ -121,19 +137,11 @@ def resp(res: httpx.Response) -> str:
         if self.__colour == "async":
 
             async def _execute() -> str:
-                mounts: Dict[str, httpx.AsyncBaseTransport] = {}
-                for key, mount in self.__make_mounts().items():
-                    assert isinstance(mount, httpx.AsyncHTTPTransport)
-                    mounts[key] = mount
-                async with httpx.AsyncClient(mounts=mounts) as client:
+                async with httpx.AsyncClient(mounts=self.__make_mounts(httpx)) as client:
                     return resp(await client.get(self._open_id_config_url))
 
             return _execute()
-        mounts: Dict[str, httpx.BaseTransport] = {}
-        for key, mount in self.__make_mounts().items():
-            assert isinstance(mount, httpx.BaseTransport)
-            mounts[key] = mount
-        with httpx.Client(mounts=mounts) as client:
+        with httpx.Client(mounts=self.__make_mounts(httpx)) as client:
             return resp(client.get(self._open_id_config_url))
 
     def get_auth_session(self) -> Result:
@@ -164,6 +172,7 @@ async def _execute() -> AsyncOAuth2Client:
                     token_endpoint=await executor.aresult(self._get_token_endpoint()),
                     client_id=self._client_id,
                     default_timeout=AUTH_DEFAULT_TIMEOUT,
+                    mounts=self.__make_mounts(_SESSION_HTTP),
                 )
 
             return _execute()
@@ -172,6 +181,7 @@ async def _execute() -> AsyncOAuth2Client:
             token_endpoint=executor.result(self._get_token_endpoint()),
             client_id=self._client_id,
             default_timeout=AUTH_DEFAULT_TIMEOUT,
+            mounts=self.__make_mounts(_SESSION_HTTP),
         )
 
     def _get_session_user_pw(self, config: AuthClientPassword) -> Result:
@@ -186,6 +196,7 @@ async def _execute() -> AsyncOAuth2Client:
                     grant_type="password",
                     scope=scope,
                     default_timeout=AUTH_DEFAULT_TIMEOUT,
+                    mounts=self.__make_mounts(_SESSION_HTTP),
                 )
                 token: dict = await session.fetch_token(
                     username=config.username, password=config.password
@@ -202,6 +213,7 @@ async def _execute() -> AsyncOAuth2Client:
             grant_type="password",
             scope=scope,
             default_timeout=AUTH_DEFAULT_TIMEOUT,
+            mounts=self.__make_mounts(_SESSION_HTTP),
         )
         token: dict = session.fetch_token(username=config.username, password=config.password)
         if "refresh_token" not in token:
@@ -239,6 +251,7 @@ async def _execute() -> AsyncOAuth2Client:
                     grant_type="client_credentials",
                     token={"access_token": None, "expires_in": -100},
                     default_timeout=AUTH_DEFAULT_TIMEOUT,
+                    mounts=self.__make_mounts(_SESSION_HTTP),
                 )
                 # explicitly fetch tokens. Otherwise, authlib will do it in the background and we might have race-conditions
                 await session.fetch_token()
@@ -254,6 +267,7 @@ async def _execute() -> AsyncOAuth2Client:
             grant_type="client_credentials",
             token={"access_token": None, "expires_in": -100},
             default_timeout=AUTH_DEFAULT_TIMEOUT,
+            mounts=self.__make_mounts(_SESSION_HTTP),
         )
         # explicitly fetch tokens. Otherwise, authlib will do it in the background and we might have race-conditions
         session.fetch_token()