diff --git a/src/flyte/remote/_client/auth/_token_client.py b/src/flyte/remote/_client/auth/_token_client.py index 331fc9eb8..775554071 100644 --- a/src/flyte/remote/_client/auth/_token_client.py +++ b/src/flyte/remote/_client/auth/_token_client.py @@ -70,6 +70,34 @@ def from_json_response(cls, j: typing.Dict) -> "DeviceCodeResponse": ) +def _body_snippet(response: httpx.Response, limit: int = 200) -> str: + """Describe a response body compactly enough to put in an error message.""" + content_type = response.headers.get("content-type", "unknown") + text = " ".join(response.text.split()) + if len(text) > limit: + text = text[:limit] + "..." + return f"content-type: {content_type}, body: {text!r}" if text else f"content-type: {content_type}, empty body" + + +def _json_object_or_none(response: httpx.Response) -> typing.Optional[typing.Dict]: + """Parse an IDP response body as a JSON object, or None when it is anything else. + + The OAuth endpoints are specified to answer in JSON, but what actually reaches the SDK is + whatever sits between it and the IDP: a load balancer's HTML 502 page, a proxy's plain-text + "Internal Server Error", an SSO interstitial. That is a deployment problem, and the branches + below already know how to report it -- they only have to survive reading the body first. + + Returns None rather than raising, and only for a JSON *object*: a body that parses to a bare + string would still answer `"error" in j` by substring, which is not the membership test the + caller means. + """ + try: + parsed = response.json() + except ValueError: + return None + return parsed if isinstance(parsed, dict) else None + + def get_basic_authorization_header(client_id: str, client_secret: str) -> str: """ This function transforms the client id and the client secret into a header that conforms with http basic auth. @@ -148,15 +176,25 @@ async def get_token( response = await http_session.post(token_endpoint, data=body, headers=headers) if not response.is_success: - j = response.json() - if "error" in j: + j = _json_object_or_none(response) + if j is not None and "error" in j: err = j["error"] if err == error_auth_pending or err == error_slow_down: raise AuthenticationPending(f"Token not yet available, try again in some time {err}") logger.error("Status Code ({}) received from IDP: {}".format(response.status_code, response.text)) raise AuthenticationError("Status Code ({}) received from IDP: {}".format(response.status_code, response.text)) - j = response.json() + j = _json_object_or_none(response) + if j is None or "access_token" not in j: + # A 2xx that is not a usable token response: an authenticating proxy answering with its + # own login page, or an endpoint that is not the IDP's at all. Saying so beats the + # KeyError/JSONDecodeError that used to escape from here. + raise AuthenticationError( + f"Token endpoint {token_endpoint} returned {response.status_code} but not an access " + f"token ({_body_snippet(response)}). Check that the endpoint in your config points at " + f"the identity provider and that nothing is intercepting the request." + ) + new_refresh_token = None if "refresh_token" in j: new_refresh_token = j["refresh_token"] @@ -199,9 +237,16 @@ async def get_device_code( if not resp.is_success: raise AuthenticationError( f"Unable to retrieve Device Authentication Code for {payload}," - f" Status Code {resp.status_code} Reason {resp.json()}" + f" Status Code {resp.status_code} Reason {_body_snippet(resp)}" + ) + j = _json_object_or_none(resp) + if j is None: + raise AuthenticationError( + f"Device authorization endpoint {device_auth_endpoint} returned {resp.status_code} " + f"with a body that is not a JSON object ({_body_snippet(resp)}). Check that the " + f"endpoint in your config points at the identity provider." ) - return DeviceCodeResponse.from_json_response(resp.json()) + return DeviceCodeResponse.from_json_response(j) async def poll_token_endpoint( diff --git a/tests/flyte/remote/test_token_client.py b/tests/flyte/remote/test_token_client.py new file mode 100644 index 000000000..4eb0ba289 --- /dev/null +++ b/tests/flyte/remote/test_token_client.py @@ -0,0 +1,169 @@ +"""Tests for how the OAuth token/device-code endpoints handle a body that is not JSON. + +Both endpoints are specified to answer in JSON, but the SDK talks to them through whatever a +deployment puts in the way: a load balancer's HTML 502 page, a proxy's plain-text "Internal +Server Error", an SSO interstitial. Reading such a body used to raise `json.JSONDecodeError` +straight out of the error branch that was about to raise a perfectly good `AuthenticationError` +(FLYTE-SDK-60). +""" + +from unittest.mock import AsyncMock, MagicMock + +import httpx +import pytest + +from flyte.remote._client.auth._token_client import ( + GrantType, + _body_snippet, + _json_object_or_none, + get_device_code, + get_token, +) +from flyte.remote._client.auth.errors import AuthenticationError, AuthenticationPending + +# What a proxy in front of a broken IDP actually returns -- the FLYTE-SDK-60 event. +NGINX_502 = "502 Bad Gateway502 Bad Gateway" + + +def _session(response: httpx.Response) -> MagicMock: + session = MagicMock() + session.post = AsyncMock(return_value=response) + return session + + +class TestJsonObjectOrNone: + @pytest.mark.parametrize( + "response, expected", + [ + (httpx.Response(200, json={"access_token": "t"}), {"access_token": "t"}), + (httpx.Response(500, text=NGINX_502), None), + (httpx.Response(500, text="Internal Server Error"), None), + (httpx.Response(200, text=""), None), + # Valid JSON that is not an object: `"error" in j` would answer by substring on a + # bare string, which is not the membership test the caller means. + (httpx.Response(400, json="error"), None), + (httpx.Response(400, json=["error"]), None), + ], + ) + def test_only_json_objects_survive(self, response, expected): + assert _json_object_or_none(response) == expected + + +class TestBodySnippet: + def test_reports_content_type_and_body(self): + snippet = _body_snippet(httpx.Response(500, text="boom", headers={"content-type": "text/plain"})) + assert "text/plain" in snippet + assert "boom" in snippet + + def test_empty_body_is_said_to_be_empty(self): + assert "empty body" in _body_snippet(httpx.Response(500, text="")) + + def test_long_body_is_truncated(self): + snippet = _body_snippet(httpx.Response(500, text="x" * 5000), limit=50) + assert len(snippet) < 200 + assert snippet.endswith("...'") + + +class TestGetTokenNonJsonBody: + @pytest.mark.asyncio + async def test_error_status_with_html_body_reports_the_status(self): + """The failure branch already had the right error to raise; it just had to get there.""" + with pytest.raises(AuthenticationError) as excinfo: + await get_token("https://idp.example.com/token", _session(httpx.Response(500, text=NGINX_502))) + + assert "Status Code (500)" in str(excinfo.value) + + @pytest.mark.asyncio + async def test_error_status_with_empty_body_reports_the_status(self): + with pytest.raises(AuthenticationError) as excinfo: + await get_token("https://idp.example.com/token", _session(httpx.Response(503, text=""))) + + assert "Status Code (503)" in str(excinfo.value) + + @pytest.mark.asyncio + async def test_success_status_with_non_json_body_names_the_endpoint(self): + """A 200 carrying a login page is not an SDK bug, and must not read as one.""" + with pytest.raises(AuthenticationError) as excinfo: + await get_token("https://idp.example.com/token", _session(httpx.Response(200, text=NGINX_502))) + + message = str(excinfo.value) + assert "https://idp.example.com/token" in message + assert "not an access token" in message + + @pytest.mark.asyncio + async def test_success_status_without_access_token_is_reported(self): + with pytest.raises(AuthenticationError) as excinfo: + await get_token("https://idp.example.com/token", _session(httpx.Response(200, json={"scope": "all"}))) + + assert "not an access token" in str(excinfo.value) + + +class TestGetTokenStillWorks: + @pytest.mark.asyncio + async def test_access_token_is_returned(self): + response = httpx.Response(200, json={"access_token": "abc", "refresh_token": "r", "expires_in": 3600}) + + access, refresh, expires = await get_token("https://idp.example.com/token", _session(response)) + + assert (access, refresh, expires) == ("abc", "r", 3600) + + @pytest.mark.asyncio + async def test_missing_refresh_token_is_fine(self): + response = httpx.Response(200, json={"access_token": "abc", "expires_in": 3600}) + + access, refresh, expires = await get_token("https://idp.example.com/token", _session(response)) + + assert (access, refresh, expires) == ("abc", None, 3600) + + @pytest.mark.asyncio + @pytest.mark.parametrize("err", ["authorization_pending", "slow_down"]) + async def test_device_flow_pending_still_raises_authentication_pending(self, err): + """The JSON error branch is the one that keeps the device-code poll loop alive.""" + response = httpx.Response(400, json={"error": err}) + + with pytest.raises(AuthenticationPending): + await get_token( + "https://idp.example.com/token", + _session(response), + grant_type=GrantType.DEVICE_CODE, + device_code="dc", + ) + + +class TestGetDeviceCodeNonJsonBody: + @pytest.mark.asyncio + async def test_error_status_does_not_crash_building_its_own_message(self): + """`Reason {resp.json()}` was interpolated into the error it was raising.""" + with pytest.raises(AuthenticationError) as excinfo: + await get_device_code( + "https://idp.example.com/device", "client", _session(httpx.Response(502, text=NGINX_502)) + ) + + assert "Status Code 502" in str(excinfo.value) + + @pytest.mark.asyncio + async def test_success_status_with_non_json_body_names_the_endpoint(self): + with pytest.raises(AuthenticationError) as excinfo: + await get_device_code( + "https://idp.example.com/device", "client", _session(httpx.Response(200, text=NGINX_502)) + ) + + assert "https://idp.example.com/device" in str(excinfo.value) + + @pytest.mark.asyncio + async def test_valid_device_code_response_is_parsed(self): + response = httpx.Response( + 200, + json={ + "device_code": "dc", + "user_code": "UC", + "verification_uri": "https://idp.example.com/activate", + "expires_in": 600, + "interval": 5, + }, + ) + + result = await get_device_code("https://idp.example.com/device", "client", _session(response)) + + assert result.device_code == "dc" + assert result.interval == 5