diff --git a/mock_tests/test_auth.py b/mock_tests/test_auth.py index cc4480c93..cb42338dc 100644 --- a/mock_tests/test_auth.py +++ b/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 diff --git a/weaviate/connect/authentication.py b/weaviate/connect/authentication.py index b7ca53996..11493fd57 100644 --- a/weaviate/connect/authentication.py +++ b/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() diff --git a/weaviate/connect/v4.py b/weaviate/connect/v4.py index 5dcadfaeb..e6549c060 100644 --- a/weaviate/connect/v4.py +++ b/weaviate/connect/v4.py @@ -7,6 +7,7 @@ from dataclasses import dataclass, field from ssl import SSLZeroReturnError from threading import Event, Thread +from types import ModuleType from typing import ( Any, AsyncGenerator, @@ -17,6 +18,7 @@ Literal, Optional, Tuple, + Type, TypeVar, Union, cast, @@ -24,6 +26,7 @@ ) import grpc +import httpx from authlib.integrations.httpx_client import ( # type: ignore AsyncOAuth2Client, OAuth2Client, @@ -42,8 +45,6 @@ HTTPError, HTTPStatusError, HTTPTransport, - Limits, - Proxy, ReadError, ReadTimeout, RemoteProtocolError, @@ -105,7 +106,18 @@ Session = Union[Client, OAuth2Client] AsyncSession = Union[AsyncClient, AsyncOAuth2Client] -HttpClient = Union[AsyncClient, AsyncOAuth2Client, Client, OAuth2Client] +HttpClient = Union[AsyncClient, Client] +# Holds and refreshes the OIDC token only; requests go through HttpClient. Authlib >= 1.8 builds these +# sessions on httpx2 when it is installed, so they must never be handed httpx objects (timeouts, transports). +OAuthSession = Union[AsyncOAuth2Client, OAuth2Client] + +try: + # an httpx2-based session raises httpx2 errors when fetching a token + from httpx2 import HTTPError as _Httpx2HTTPError # type: ignore + + _OAUTH_HTTP_ERRORS: Tuple[Type[Exception], ...] = (HTTPError, _Httpx2HTTPError) +except ImportError: + _OAUTH_HTTP_ERRORS = (HTTPError,) PERMISSION_DENIED = "PERMISSION_DENIED" @@ -143,6 +155,7 @@ def __init__( self.__additional_headers = {} self._auth = auth_client_secret self._client: Optional[HttpClient] = None + self._oauth_session: Optional[OAuthSession] = None self._connection_params = connection_params self._grpc_stub: Optional[weaviate_pb2_grpc.WeaviateStub] = None self._grpc_channel: Union[AsyncChannel, SyncChannel, None] = None @@ -239,34 +252,24 @@ def _make_mounts(self, colour: Literal["sync"]) -> Dict[str, HTTPTransport]: ... def _make_mounts( self, colour: executor.Colour ) -> Union[Dict[str, AsyncHTTPTransport], Dict[str, HTTPTransport]]: - if colour == "async": - return { - (f"{key}://" if key == "http" or key == "https" else key): AsyncHTTPTransport( - limits=Limits( - max_connections=self.__connection_config.session_pool_maxsize, - max_keepalive_connections=self.__connection_config.session_pool_connections, - ), - proxy=Proxy(url=proxy), - retries=self.__connection_config.session_pool_max_retries, - trust_env=self.__trust_env, - ) - for key, proxy in self._proxies.items() - if key != "grpc" - } - if colour == "sync": - return { - f"{key}://" if key == "http" or key == "https" else key: HTTPTransport( - limits=Limits( - max_connections=self.__connection_config.session_pool_maxsize, - max_keepalive_connections=self.__connection_config.session_pool_connections, - ), - proxy=Proxy(url=proxy), - retries=self.__connection_config.session_pool_max_retries, - trust_env=self.__trust_env, - ) - for key, proxy in self._proxies.items() - if key != "grpc" - } + return self._make_mounts_with(httpx, colour) + + def _make_mounts_with(self, http: ModuleType, colour: executor.Colour) -> Dict[str, Any]: + """Build the proxy transports with `http`: httpx, or httpx2 for an httpx2-based authlib session.""" + transport = http.AsyncHTTPTransport if colour == "async" else http.HTTPTransport + return { + (f"{key}://" if key == "http" or key == "https" else key): transport( + limits=http.Limits( + max_connections=self.__connection_config.session_pool_maxsize, + max_keepalive_connections=self.__connection_config.session_pool_connections, + ), + proxy=http.Proxy(url=proxy), + retries=self.__connection_config.session_pool_max_retries, + trust_env=self.__trust_env, + ) + for key, proxy in self._proxies.items() + if key != "grpc" + } def is_connected(self) -> bool: return self._connected @@ -274,8 +277,8 @@ def is_connected(self) -> bool: def get_current_bearer_token(self) -> str: if "authorization" in self._headers: return self._headers["authorization"] - elif isinstance(self._client, (OAuth2Client, AsyncOAuth2Client)): - return f"Bearer {self._client.token['access_token']}" + elif self._oauth_session is not None: + return f"Bearer {self._oauth_session.token['access_token']}" return "" def _prepare_grpc_headers(self) -> None: @@ -479,16 +482,17 @@ async def execute() -> None: _Auth.use( oidc_config=resp, credentials=auth_client_secret, - make_mounts=lambda: self._make_mounts("async"), + make_mounts=lambda http: self._make_mounts_with(http, "async"), colour=colour, ) ) try: - self._client = await _auth.aresult(_auth.get_auth_session()) - except HTTPError as e: + self._oauth_session = await _auth.aresult(_auth.get_auth_session()) + except _OAUTH_HTTP_ERRORS as e: raise AuthenticationFailedError( f"Failed to authenticate with OIDC: {repr(e)}" ) + self.__make_clients(colour) if isinstance(auth_client_secret, AuthClientCredentials): # credentials should only be saved for client credentials, otherwise use refresh token @@ -502,16 +506,17 @@ async def execute() -> None: _Auth.use( oidc_config=resp, credentials=auth_client_secret, - make_mounts=lambda: self._make_mounts("sync"), + make_mounts=lambda http: self._make_mounts_with(http, "sync"), colour=colour, ) ) try: - self._client = _auth.result(_auth.get_auth_session()) - except HTTPError as e: + self._oauth_session = _auth.result(_auth.get_auth_session()) + except _OAUTH_HTTP_ERRORS as e: raise AuthenticationFailedError( f"Failed to authenticate with OIDC: {repr(e)}" ) + self.__make_clients(colour) if isinstance(auth_client_secret, AuthClientCredentials): # credentials should only be saved for client credentials, otherwise use refresh token @@ -546,12 +551,12 @@ async def execute() -> None: def _create_background_token_refresh(self, _auth: Optional[_Auth] = None) -> None: """Create a background thread that periodically refreshes access and refresh tokens. - While the underlying library refreshes tokens, it does not have an internal cronjob that checks every - X-seconds if a token has expired. If there is no activity for longer than the refresh tokens lifetime, it will - expire. Therefore, refresh manually shortly before expiration time is up. + Requests go through our own HTTP client with the session's current token, so the session never refreshes + on its own. Refresh shortly before the token expires, and keep refreshing so the refresh token does not + expire either. """ - assert isinstance(self._client, (OAuth2Client, AsyncOAuth2Client)) - if "refresh_token" not in self._client.token and _auth is None: + assert isinstance(self._oauth_session, (OAuth2Client, AsyncOAuth2Client)) + if "refresh_token" not in self._oauth_session.token and _auth is None: return # stop the refresher from an earlier connect(), if any @@ -559,13 +564,13 @@ def _create_background_token_refresh(self, _auth: Optional[_Auth] = None) -> Non # refresh 30s before the token expires (assume 1 minute if the token does not say); # the loops below always wait at least 1s - refresh_in: int = self._client.token.get("expires_in", 60) - 30 + refresh_in: int = self._oauth_session.token.get("expires_in", 60) - 30 # the refresher keeps its own event: after close() + connect() the old refresher # must stop on this one instead of running on with the new one shutdown = Event() self._shutdown_background_event = shutdown - if isinstance(self._client, AsyncOAuth2Client): + if isinstance(self._oauth_session, AsyncOAuth2Client): # async client: refresh in an asyncio task on the current loop, not in a thread # (threads cannot start under WASM/Pyodide) self.__token_refresh_task = asyncio.get_running_loop().create_task( @@ -578,8 +583,8 @@ def periodic_refresh_token(refresh_time: int, _auth: Optional[_Auth]) -> None: while not shutdown.wait(timeout=max(refresh_time, 1)): try: # use one client for the whole round: close()/connect() may replace - # self._client while a refresh is running - client = self._client + # self._oauth_session while a refresh is running + client = self._oauth_session if not isinstance(client, OAuth2Client): continue if "refresh_token" in client.token: @@ -588,8 +593,12 @@ def periodic_refresh_token(refresh_time: int, _auth: Optional[_Auth]) -> None: # client credentials usually does not contain a refresh token => get a # new token using the saved credentials assert _auth is not None + # the new session already fetched a token; only that token is kept new_session = _auth.result(_auth.get_auth_session()) - client.token = new_session.fetch_token() + try: + client.token = new_session.token + finally: + new_session.close() refresh_time = client.token.get("expires_in", 60) - 30 except Exception as exc: # retry in one second; any error must keep the refresher alive, not only @@ -628,7 +637,7 @@ async def __periodic_token_refresh_async( while not shutdown.is_set(): await asyncio.sleep(max(refresh_time, 1)) try: - client = self._client + client = self._oauth_session if not isinstance(client, AsyncOAuth2Client): continue if "refresh_token" in client.token: @@ -637,8 +646,12 @@ async def __periodic_token_refresh_async( # client credentials usually does not contain a refresh token => get a # new token using the saved credentials assert _auth is not None + # the new session already fetched a token; only that token is kept new_session = await _Auth.aresult(_auth.get_auth_session()) - client.token = await new_session.fetch_token() + try: + client.token = new_session.token + finally: + await new_session.aclose() refresh_time = client.token.get("expires_in", 60) - 30 except asyncio.CancelledError: raise @@ -769,6 +782,10 @@ async def execute() -> None: assert isinstance(self._client, AsyncClient) await self._client.aclose() self._client = None + if self._oauth_session is not None: + assert isinstance(self._oauth_session, AsyncOAuth2Client) + await self._oauth_session.aclose() + self._oauth_session = None if self._grpc_stub is not None: assert self._grpc_channel is not None assert isinstance(self._grpc_channel, AsyncChannel) @@ -782,6 +799,10 @@ async def execute() -> None: assert isinstance(self._client, Client) self._client.close() self._client = None + if self._oauth_session is not None: + assert isinstance(self._oauth_session, OAuth2Client) + self._oauth_session.close() + self._oauth_session = None if self._grpc_stub is not None: assert self._grpc_channel is not None assert isinstance(self._grpc_channel, SyncChannel)