From eca6b009f3e5382ff9b9cce055ca044f059a655e Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Mon, 21 Sep 2026 12:58:15 +0300 Subject: [PATCH 1/7] feat: add OIDC credentials providers --- CHANGELOG.md | 1 + docs/driver.rst | 54 ++++ examples/oidc-credentials/main.py | 70 +++++ tests/aio/test_credentials.py | 57 ++++ tests/auth/test_credentials.py | 115 ++++++++ tests/oidc/.gitignore | 1 + tests/oidc/README.md | 133 +++++++++ tests/oidc/_smoke_common.py | 171 ++++++++++++ tests/oidc/client_credentials_smoke.py | 29 ++ tests/oidc/compose.yaml | 84 ++++++ tests/oidc/device_authorization_smoke.py | 75 ++++++ tests/oidc/prepare.sh | 146 ++++++++++ tests/oidc/static_token_smoke.py | 25 ++ tests/oidc/ydb.yaml | 330 +++++++++++++++++++++++ ydb/__init__.py | 2 +- ydb/aio/oidc.py | 192 +++++++++++++ ydb/oidc/__init__.py | 13 + ydb/oidc/_common.py | 193 +++++++++++++ ydb/oidc/credentials.py | 203 ++++++++++++++ 19 files changed, 1893 insertions(+), 1 deletion(-) create mode 100644 examples/oidc-credentials/main.py create mode 100644 tests/oidc/.gitignore create mode 100644 tests/oidc/README.md create mode 100644 tests/oidc/_smoke_common.py create mode 100644 tests/oidc/client_credentials_smoke.py create mode 100644 tests/oidc/compose.yaml create mode 100644 tests/oidc/device_authorization_smoke.py create mode 100755 tests/oidc/prepare.sh create mode 100644 tests/oidc/static_token_smoke.py create mode 100644 tests/oidc/ydb.yaml create mode 100644 ydb/aio/oidc.py create mode 100644 ydb/oidc/__init__.py create mode 100644 ydb/oidc/_common.py create mode 100644 ydb/oidc/credentials.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c53465784..4741d0ee6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,4 @@ +* Add static-token, Client Credentials, and Device Authorization OAuth 2.0 providers under `ydb.oidc`, with non-blocking Client Credentials and Device Authorization variants under `ydb.aio.oidc` * Deprecated the table client scan query methods — `TableClient.scan_query`, `TableClient.async_scan_query` and the async `ydb.aio.TableClient.scan_query`: they now emit a `DeprecationWarning` and keep working as before, use QueryService (`ydb.QuerySessionPool` / `ydb.aio.QuerySessionPool`) instead * Mark the package as typed so type checkers use the SDK's inline annotations diff --git a/docs/driver.rst b/docs/driver.rst index d58762400..4ade89a0d 100644 --- a/docs/driver.rst +++ b/docs/driver.rst @@ -182,6 +182,60 @@ Pass a static IAM token or API key directly credentials = ydb.AccessTokenCredentials("your-token") +OIDC and OAuth 2.0 Credentials +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +``ydb.oidc`` provides three OAuth 2.0 credential modes for an external identity +provider advertised through OIDC discovery +(`full example `__). + +Use an access token obtained outside the SDK: + +.. code-block:: python + + import ydb.oidc + + credentials = ydb.oidc.OAuth2TokenCredentials("your-access-token") + +Use the Client Credentials Grant for a service acting on its own behalf: + +.. code-block:: python + + credentials = ydb.oidc.OAuth2ClientCredentials( + issuer="https://identity.example.com/realms/example", + client_id="service-client", + client_secret="service-client-secret", + audience="ydb", + ca_file="/path/to/idp-ca.pem", + ) + +Use the Device Authorization Grant for a CLI or another input-constrained client. The +callback is invoked before polling begins and must show the verification URI and user +code to the user: + +.. code-block:: python + + def show_device_authorization(info): + print(info.verification_uri_complete or info.verification_uri) + print(info.user_code) + + credentials = ydb.oidc.OAuth2DeviceCredentials( + issuer="https://identity.example.com/realms/example", + client_id="public-device-client", + device_authorization_callback=show_device_authorization, + scope=["openid", "offline_access"], + ca_file="/path/to/idp-ca.pem", + ) + +Client Credentials obtains a new access token when needed. Device Authorization uses a +returned refresh token for subsequent refreshes and starts a new user interaction if +the refresh token is no longer valid. Access tokens are sent to YDB using the +``Bearer`` authentication scheme. + +Non-blocking counterparts are available as ``ydb.aio.oidc.OAuth2ClientCredentials`` +and ``ydb.aio.oidc.OAuth2DeviceCredentials``. The asynchronous Device Authorization +callback may be either a regular callable or an async callable. + StaticCredentials (username/password) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/examples/oidc-credentials/main.py b/examples/oidc-credentials/main.py new file mode 100644 index 000000000..5f891da21 --- /dev/null +++ b/examples/oidc-credentials/main.py @@ -0,0 +1,70 @@ +import argparse + +import ydb +import ydb.oidc + + +def parse_args(): + parser = argparse.ArgumentParser(description="OIDC and OAuth 2.0 credentials example") + parser.add_argument("--endpoint", required=True, help="YDB endpoint") + parser.add_argument("--database", required=True, help="YDB database") + parser.add_argument("--mode", required=True, choices=("token", "client", "device")) + parser.add_argument("--access-token", help="Existing OAuth 2.0 access token") + parser.add_argument("--issuer", help="OIDC issuer URL") + parser.add_argument("--client-id", help="OAuth 2.0 client ID") + parser.add_argument("--client-secret", help="OAuth 2.0 client secret") + parser.add_argument("--scope", action="append", help="OAuth 2.0 scope; may be repeated") + parser.add_argument("--audience", help="OAuth 2.0 token audience") + parser.add_argument("--ca-file", help="CA certificate for the identity provider") + args = parser.parse_args() + + if args.mode == "token" and not args.access_token: + parser.error("--access-token is required in token mode") + if args.mode in ("client", "device") and (not args.issuer or not args.client_id): + parser.error("--issuer and --client-id are required in client and device modes") + if args.mode == "client" and not args.client_secret: + parser.error("--client-secret is required in client mode") + + return args + + +def device_authorization_callback(info): + print("Open {}".format(info.verification_uri_complete or info.verification_uri)) + print("Enter code {}".format(info.user_code)) + + +def create_credentials(args): + if args.mode == "token": + return ydb.oidc.OAuth2TokenCredentials(args.access_token) + + common = { + "issuer": args.issuer, + "client_id": args.client_id, + "audience": args.audience, + "ca_file": args.ca_file, + } + if args.scope is not None: + common["scope"] = args.scope + if args.mode == "client": + return ydb.oidc.OAuth2ClientCredentials(client_secret=args.client_secret, **common) + return ydb.oidc.OAuth2DeviceCredentials( + device_authorization_callback=device_authorization_callback, + **common, + ) + + +def main(): + args = parse_args() + with ydb.Driver( + endpoint=args.endpoint, + database=args.database, + credentials=create_credentials(args), + ) as driver: + driver.wait(timeout=10, fail_fast=True) + with ydb.QuerySessionPool(driver) as pool: + result_sets = pool.execute_with_retries("SELECT 42 AS answer") + print(result_sets[0].rows[0].answer) + + +if __name__ == "__main__": + main() diff --git a/tests/aio/test_credentials.py b/tests/aio/test_credentials.py index e73f35221..4e9d0f376 100644 --- a/tests/aio/test_credentials.py +++ b/tests/aio/test_credentials.py @@ -12,6 +12,7 @@ import tests.oauth2_token_exchange import tests.oauth2_token_exchange.test_token_exchange import ydb.aio.iam +import ydb.aio.oidc import ydb.aio.oauth2_token_exchange import ydb.oauth2_token_exchange.token_source @@ -263,3 +264,59 @@ def mock_submit(callback): token3 = await credentials.token() assert token3 == "token_v2" assert call_count == 2 + + +@pytest.mark.asyncio +async def test_oauth2_client_credentials(): + issuer = "https://issuer.example" + credentials = ydb.aio.oidc.OAuth2ClientCredentials(issuer, "client-id", "client-secret") + credentials._request_json = AsyncMock( + side_effect=[ + (200, {"issuer": issuer, "token_endpoint": issuer + "/token"}), + (200, {"access_token": "access-token", "token_type": "Bearer", "expires_in": 300}), + ] + ) + + assert await credentials.get_auth_token() == "Bearer access-token" + assert credentials._request_json.await_count == 2 + + +@pytest.mark.asyncio +async def test_oauth2_device_credentials(): + issuer = "https://issuer.example" + callback_values = [] + + async def callback(info): + callback_values.append(info) + + credentials = ydb.aio.oidc.OAuth2DeviceCredentials(issuer, "device-client", callback) + credentials._request_json = AsyncMock( + side_effect=[ + ( + 200, + { + "issuer": issuer, + "token_endpoint": issuer + "/token", + "device_authorization_endpoint": issuer + "/device", + }, + ), + ( + 200, + { + "device_code": "device-code", + "user_code": "user-code", + "verification_uri": issuer + "/verify", + "expires_in": 600, + "interval": 1, + }, + ), + (400, {"error": "authorization_pending"}), + (200, {"access_token": "access-token", "token_type": "Bearer", "expires_in": 300}), + ] + ) + + with patch("ydb.aio.oidc.asyncio.sleep", new=AsyncMock()) as sleep: + assert await credentials.get_auth_token() == "Bearer access-token" + + assert callback_values[0].user_code == "user-code" + assert sleep.await_count == 2 diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index a78040ce8..3207beeb3 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -2,8 +2,10 @@ import concurrent.futures import grpc import time +from unittest.mock import patch import ydb.iam +import ydb.oidc from yandex.cloud.iam.v1 import iam_token_service_pb2_grpc from yandex.cloud.iam.v1 import iam_token_service_pb2 @@ -74,3 +76,116 @@ def test_yandex_service_account_credentials(): assert t == "test_token" assert credentials.get_expire_time() <= 42 server.stop() + + +def test_oauth2_token_credentials(): + credentials = ydb.oidc.OAuth2TokenCredentials("access-token") + + assert credentials.auth_metadata() == [("x-ydb-auth-ticket", "Bearer access-token")] + assert ydb.oidc.OAuth2TokenCredentials("Bearer access-token").get_auth_token() == "Bearer access-token" + + +def test_oauth2_client_credentials(): + issuer = "https://issuer.example" + requests = [] + responses = iter( + [ + (200, {"issuer": issuer, "token_endpoint": issuer + "/token"}), + (200, {"access_token": "access-token", "token_type": "Bearer", "expires_in": 300}), + ] + ) + credentials = ydb.oidc.OAuth2ClientCredentials( + issuer, + "client-id", + "client-secret", + scope=["openid", "profile"], + audience="ydb", + ) + + def request_json(url, data=None, headers=None): + requests.append((url, data, headers)) + return next(responses) + + credentials._request_json = request_json + + assert credentials.get_auth_token() == "Bearer access-token" + assert requests[0] == (issuer + "/.well-known/openid-configuration", None, None) + assert requests[1][0] == issuer + "/token" + assert requests[1][1] == { + "grant_type": "client_credentials", + "scope": "openid profile", + "audience": "ydb", + } + assert requests[1][2]["Authorization"].startswith("Basic ") + + +def test_oauth2_device_credentials_poll_and_refresh(): + issuer = "https://issuer.example" + callback_values = [] + responses = iter( + [ + ( + 200, + { + "issuer": issuer, + "token_endpoint": issuer + "/token", + "device_authorization_endpoint": issuer + "/device", + }, + ), + ( + 200, + { + "device_code": "device-code", + "user_code": "user-code", + "verification_uri": issuer + "/verify", + "verification_uri_complete": issuer + "/verify?user_code=user-code", + "expires_in": 600, + "interval": 1, + }, + ), + (400, {"error": "authorization_pending"}), + (400, {"error": "slow_down"}), + ( + 200, + { + "access_token": "device-access-token", + "refresh_token": "refresh-token", + "token_type": "Bearer", + "expires_in": 300, + }, + ), + ( + 200, + { + "access_token": "refreshed-access-token", + "refresh_token": "new-refresh-token", + "token_type": "Bearer", + "expires_in": 300, + }, + ), + ] + ) + requests = [] + credentials = ydb.oidc.OAuth2DeviceCredentials( + issuer, + "device-client", + callback_values.append, + ) + + def request_json(url, data=None, headers=None): + requests.append((url, data, headers)) + return next(responses) + + credentials._request_json = request_json + + with patch("ydb.oidc.credentials.time.sleep") as sleep: + assert credentials.get_auth_token() == "Bearer device-access-token" + + assert callback_values[0].user_code == "user-code" + assert [value.args[0] for value in sleep.call_args_list] == [1, 1, 6] + assert credentials._make_token_request() == { + "access_token": "Bearer refreshed-access-token", + "expires_in": 300, + } + assert requests[-1][1]["grant_type"] == "refresh_token" + assert requests[-1][1]["refresh_token"] == "refresh-token" diff --git a/tests/oidc/.gitignore b/tests/oidc/.gitignore new file mode 100644 index 000000000..95f7491ee --- /dev/null +++ b/tests/oidc/.gitignore @@ -0,0 +1 @@ +.state/ diff --git a/tests/oidc/README.md b/tests/oidc/README.md new file mode 100644 index 000000000..2166bb4ca --- /dev/null +++ b/tests/oidc/README.md @@ -0,0 +1,133 @@ +# Local YDB and Keycloak OIDC environment + +This directory provides a self-contained development environment for the OIDC authentication contract shared by YDB SDKs. It runs Keycloak and a trunk build of local YDB in one Compose project. + +The generated credentials and self-signed certificate are for local development only. + +## Generate local configuration + +Run the preparation step once: + +```shell +tests/oidc/prepare.sh +``` + +It generates two ignored files: + +- `.state/oidc.env` is the source of truth for container and SDK settings; +- `.state/keycloak/ydb-realm.json` is the Keycloak realm import assembled from the same generated credentials. + +No client secret or test-user password is stored in the repository. Load the generated variables before running a host-side test: + +```shell +set -a +source tests/oidc/.state/oidc.env +set +a +``` + +The relevant flow variables are: + +| Flow | Variables | +| --- | --- | +| Static token | `YDB_ENDPOINT`, `YDB_DATABASE`; set the access token returned by either flow in the SDK's static-token input | +| Client Credentials | `YDB_OIDC_ISSUER`, `YDB_OIDC_AUDIENCE`, `YDB_OIDC_CLIENT_ID`, `YDB_OIDC_CLIENT_SECRET`, `YDB_OIDC_CA_FILE` | +| Device Authorization | `YDB_OIDC_ISSUER`, `YDB_OIDC_AUDIENCE`, `YDB_OIDC_DEVICE_CLIENT_ID`, `YDB_OIDC_DEVICE_USERNAME`, `YDB_OIDC_DEVICE_PASSWORD`, `YDB_OIDC_CA_FILE` | + +The public device client intentionally has no client secret. + +## Start and verify + +From the repository root: + +```shell +docker compose \ + --env-file tests/oidc/.state/oidc.env \ + -f tests/oidc/compose.yaml \ + up -d --wait +``` + +With the project virtual environment activated and the generated variables loaded, run the three smoke tests: + +```shell +PYTHONPATH=. python tests/oidc/static_token_smoke.py +PYTHONPATH=. python tests/oidc/client_credentials_smoke.py +PYTHONPATH=. python tests/oidc/device_authorization_smoke.py +``` + +Every script exercises both the synchronous and asynchronous drivers. The Client Credentials and Device Authorization scripts also force token expiry and verify refresh. The static-token script verifies that YDB rejects an invalid bearer token. + +The Device Authorization smoke test automatically completes Keycloak login and consent using the generated local test user. This automation belongs only to the test; production applications receive the verification URI and user code through `device_authorization_callback` and leave authentication to the user. + +The public API exercised by the scripts is: + +```python +import ydb.oidc + +ydb.oidc.OAuth2TokenCredentials(access_token) +ydb.oidc.OAuth2ClientCredentials( + issuer=issuer, + client_id=client_id, + client_secret=client_secret, + audience="ydb", + ca_file=ca_file, +) +ydb.oidc.OAuth2DeviceCredentials( + issuer=issuer, + client_id=device_client_id, + device_authorization_callback=show_verification_url_and_code, + audience="ydb", + ca_file=ca_file, +) +``` + +The non-blocking equivalents use the same names under `ydb.aio.oidc`. + +The stack uses non-default host ports: + +- YDB gRPC: `grpc://localhost:22136` +- YDB monitoring: +- Keycloak: + +## Device Authorization + +Request a device code with the generated public client ID: + +```shell +curl --cacert "$YDB_OIDC_CA_FILE" \ + --request POST \ + --data "client_id=$YDB_OIDC_DEVICE_CLIENT_ID" \ + "$YDB_OIDC_ISSUER/protocol/openid-connect/auth/device" +``` + +Open the returned `verification_uri_complete` and sign in with `YDB_OIDC_DEVICE_USERNAME` and `YDB_OIDC_DEVICE_PASSWORD`. Poll the token endpoint using the returned `device_code` and grant type `urn:ietf:params:oauth:grant-type:device_code`. + +## Stop and reset + +Stop the stack while preserving the imported Keycloak realm: + +```shell +docker compose \ + --env-file tests/oidc/.state/oidc.env \ + -f tests/oidc/compose.yaml \ + down +``` + +To rotate all local credentials, remove the Keycloak volume and regenerate the state: + +```shell +docker compose \ + --env-file tests/oidc/.state/oidc.env \ + -f tests/oidc/compose.yaml \ + down --volumes +tests/oidc/prepare.sh --force +``` + +YDB uses in-memory pdisks and intentionally has no data volume. Keycloak keeps its imported realm in a named volume, so `down --volumes` is required before importing newly generated credentials. + +## Configuration notes + +YDB currently requires an HTTPS issuer for external IdP discovery. Compose creates a self-signed Keycloak certificate on first start, and `.state/oidc.env` points clients to its CA file. + +YDB and Keycloak share a network namespace so both YDB inside Docker and an SDK on the host see the exact issuer `https://localhost:28443/realms/ydb`. OIDC issuer matching includes the scheme, hostname, port, and path. + +`enforce_user_token_check_requirement` is enabled in `ydb.yaml`, so a supplied invalid token is rejected instead of falling back to anonymous access. Anonymous requests remain available for local YDB initialization and health checks. diff --git a/tests/oidc/_smoke_common.py b/tests/oidc/_smoke_common.py new file mode 100644 index 000000000..1b4a0d488 --- /dev/null +++ b/tests/oidc/_smoke_common.py @@ -0,0 +1,171 @@ +import asyncio +import base64 +import http.cookiejar +import json +import os +import ssl +import urllib.parse +import urllib.request +from html.parser import HTMLParser +from typing import Dict, List, Optional + +import ydb + + +def required_env(name: str) -> str: + value = os.getenv(name) + if not value: + raise RuntimeError("Missing {}: run prepare.sh and load .state/oidc.env".format(name)) + return value + + +def oidc_settings() -> Dict[str, str]: + return { + "issuer": required_env("YDB_OIDC_ISSUER"), + "audience": required_env("YDB_OIDC_AUDIENCE"), + "ca_file": required_env("YDB_OIDC_CA_FILE"), + } + + +def client_credentials() -> ydb.oidc.OAuth2ClientCredentials: + return ydb.oidc.OAuth2ClientCredentials( + client_id=required_env("YDB_OIDC_CLIENT_ID"), + client_secret=required_env("YDB_OIDC_CLIENT_SECRET"), + **oidc_settings(), + ) + + +def raw_access_token() -> str: + token = os.getenv("YDB_ACCESS_TOKEN") + if token: + return token.removeprefix("Bearer ") + return client_credentials().get_auth_token().removeprefix("Bearer ") + + +def assert_token_claims(token: str) -> None: + encoded_payload = token.split(".")[1] + padding = "=" * (-len(encoded_payload) % 4) + claims = json.loads(base64.urlsafe_b64decode(encoded_payload + padding)) + audience = claims.get("aud", []) + if isinstance(audience, str): + audience = [audience] + if claims.get("iss") != required_env("YDB_OIDC_ISSUER"): + raise RuntimeError("Unexpected issuer: {!r}".format(claims.get("iss"))) + if required_env("YDB_OIDC_AUDIENCE") not in audience: + raise RuntimeError("Unexpected audience: {!r}".format(audience)) + + +def execute_sync(credentials, label: str) -> None: + with ydb.Driver( + endpoint=required_env("YDB_ENDPOINT"), + database=required_env("YDB_DATABASE"), + credentials=credentials, + ) as driver: + driver.wait(timeout=10, fail_fast=True) + with ydb.QuerySessionPool(driver) as pool: + result_sets = pool.execute_with_retries("SELECT 42 AS answer") + answer = result_sets[0].rows[0].answer + if answer != 42: + raise RuntimeError("{} returned {!r}".format(label, answer)) + print("{} sync query succeeded: answer=42".format(label)) + + +async def execute_async(credentials, label: str) -> None: + async with ydb.aio.Driver( + endpoint=required_env("YDB_ENDPOINT"), + database=required_env("YDB_DATABASE"), + credentials=credentials, + ) as driver: + await driver.wait(timeout=10, fail_fast=True) + async with ydb.aio.QuerySessionPool(driver) as pool: + result_sets = await pool.execute_with_retries("SELECT 42 AS answer") + answer = result_sets[0].rows[0].answer + if answer != 42: + raise RuntimeError("{} returned {!r}".format(label, answer)) + print("{} async query succeeded: answer=42".format(label)) + + +class _Form: + def __init__(self, form_id: Optional[str], action: str): + self.form_id = form_id + self.action = action + self.inputs: Dict[str, str] = {} + + +class _FormParser(HTMLParser): + def __init__(self): + super(_FormParser, self).__init__() + self.forms: List[_Form] = [] + self._current: Optional[_Form] = None + + def handle_starttag(self, tag, attrs): + attributes = dict(attrs) + if tag == "form": + action = attributes.get("action") + if action: + self._current = _Form(attributes.get("id"), action) + self.forms.append(self._current) + elif tag == "input" and self._current is not None: + name = attributes.get("name") + if name: + self._current.inputs[name] = attributes.get("value", "") + + def handle_endtag(self, tag): + if tag == "form": + self._current = None + + +def _forms(content: bytes) -> List[_Form]: + parser = _FormParser() + parser.feed(content.decode("utf-8")) + return parser.forms + + +def approve_keycloak_device(info: ydb.oidc.DeviceAuthorizationInfo) -> None: + if not info.verification_uri_complete: + raise RuntimeError("Keycloak did not return verification_uri_complete") + + context = ssl.create_default_context(cafile=required_env("YDB_OIDC_CA_FILE")) + opener = urllib.request.build_opener( + urllib.request.HTTPSHandler(context=context), + urllib.request.HTTPCookieProcessor(http.cookiejar.CookieJar()), + ) + + with opener.open(info.verification_uri_complete, timeout=10) as response: + current_url = response.geturl() + login_forms = [form for form in _forms(response.read()) if form.form_id == "kc-form-login"] + if len(login_forms) != 1: + raise RuntimeError("Expected one Keycloak login form") + + login = login_forms[0] + login.inputs.update( + { + "username": required_env("YDB_OIDC_DEVICE_USERNAME"), + "password": required_env("YDB_OIDC_DEVICE_PASSWORD"), + "credentialId": "", + } + ) + request = urllib.request.Request( + urllib.parse.urljoin(current_url, login.action), + data=urllib.parse.urlencode(login.inputs).encode("utf-8"), + ) + with opener.open(request, timeout=10) as response: + current_url = response.geturl() + consent_forms = [form for form in _forms(response.read()) if "login-actions/consent" in form.action] + if len(consent_forms) != 1: + raise RuntimeError("Expected one Keycloak device consent form") + + consent = consent_forms[0] + consent.inputs["accept"] = "Yes" + request = urllib.request.Request( + urllib.parse.urljoin(current_url, consent.action), + data=urllib.parse.urlencode(consent.inputs).encode("utf-8"), + ) + with opener.open(request, timeout=10) as response: + content = response.read() + if b"Device Login Successful" not in content: + raise RuntimeError("Keycloak did not confirm Device Authorization") + + +async def approve_keycloak_device_async(info: ydb.oidc.DeviceAuthorizationInfo) -> None: + await asyncio.to_thread(approve_keycloak_device, info) diff --git a/tests/oidc/client_credentials_smoke.py b/tests/oidc/client_credentials_smoke.py new file mode 100644 index 000000000..625177863 --- /dev/null +++ b/tests/oidc/client_credentials_smoke.py @@ -0,0 +1,29 @@ +import asyncio + +import ydb.aio.oidc + +from _smoke_common import client_credentials, execute_async, execute_sync, oidc_settings, required_env + + +async def run_async() -> None: + credentials = ydb.aio.oidc.OAuth2ClientCredentials( + client_id=required_env("YDB_OIDC_CLIENT_ID"), + client_secret=required_env("YDB_OIDC_CLIENT_SECRET"), + **oidc_settings(), + ) + await execute_async(credentials, "Client Credentials") + credentials._expires_in = 0 + await execute_async(credentials, "Client Credentials refresh") + + +def main() -> None: + credentials = client_credentials() + execute_sync(credentials, "Client Credentials") + credentials._expires_in = 0 + execute_sync(credentials, "Client Credentials refresh") + + asyncio.run(run_async()) + + +if __name__ == "__main__": + main() diff --git a/tests/oidc/compose.yaml b/tests/oidc/compose.yaml new file mode 100644 index 000000000..515b0dbb9 --- /dev/null +++ b/tests/oidc/compose.yaml @@ -0,0 +1,84 @@ +name: ydb-sdk-oidc + +services: + keycloak-cert-init: + image: quay.io/keycloak/keycloak:26.4.0 + user: "0:0" + entrypoint: ["/bin/bash", "-ec"] + environment: + KEYCLOAK_KEYSTORE_PASSWORD: ${KEYCLOAK_KEYSTORE_PASSWORD:?run prepare.sh first} + command: + - | + if [[ ! -f /certs/keycloak.p12 ]]; then + keytool -genkeypair \ + -alias keycloak \ + -keyalg RSA \ + -keysize 2048 \ + -validity 3650 \ + -dname "CN=localhost" \ + -ext "SAN=dns:localhost,ip:127.0.0.1" \ + -storetype PKCS12 \ + -keystore /certs/keycloak.p12 \ + -storepass "$${KEYCLOAK_KEYSTORE_PASSWORD}" \ + -keypass "$${KEYCLOAK_KEYSTORE_PASSWORD}" + keytool -exportcert \ + -rfc \ + -alias keycloak \ + -keystore /certs/keycloak.p12 \ + -storepass "$${KEYCLOAK_KEYSTORE_PASSWORD}" \ + -file /certs/keycloak-ca.pem + chmod 0644 /certs/keycloak.p12 /certs/keycloak-ca.pem + fi + volumes: + - ./.state/certs:/certs + + keycloak: + image: quay.io/keycloak/keycloak:26.4.0 + hostname: localhost + depends_on: + keycloak-cert-init: + condition: service_completed_successfully + command: + - start-dev + - --import-realm + - --hostname=https://localhost:28443 + - --https-port=28443 + - --https-key-store-file=/certs/keycloak.p12 + - --https-key-store-password=${KEYCLOAK_KEYSTORE_PASSWORD:?run prepare.sh first} + environment: + KC_BOOTSTRAP_ADMIN_USERNAME: ${KEYCLOAK_ADMIN_USERNAME:?run prepare.sh first} + KC_BOOTSTRAP_ADMIN_PASSWORD: ${KEYCLOAK_ADMIN_PASSWORD:?run prepare.sh first} + KC_FEATURES: device-flow + healthcheck: + test: ["CMD", "/bin/bash", "-c", "exec 3<>/dev/tcp/127.0.0.1/28443"] + interval: 5s + timeout: 2s + retries: 30 + start_period: 20s + ports: + - "28443:28443" + - "22136:22136" + - "28765:8765" + volumes: + - ./.state/certs:/certs:ro + - ./.state/keycloak/ydb-realm.json:/opt/keycloak/data/import/ydb-realm.json:ro + - keycloak-data:/opt/keycloak/data + + ydb: + image: ydbplatform/local-ydb:trunk + platform: linux/amd64 + depends_on: + keycloak: + condition: service_healthy + network_mode: service:keycloak + command: ["--config-path", "/ydb-config/ydb.yaml"] + environment: + GRPC_PORT: "22136" + GRPC_TLS_PORT: "0" + YDB_GRPC_ENABLE_TLS: "false" + YDB_USE_IN_MEMORY_PDISKS: "true" + volumes: + - ./ydb.yaml:/ydb-config/ydb.yaml:ro + +volumes: + keycloak-data: diff --git a/tests/oidc/device_authorization_smoke.py b/tests/oidc/device_authorization_smoke.py new file mode 100644 index 000000000..a2077836c --- /dev/null +++ b/tests/oidc/device_authorization_smoke.py @@ -0,0 +1,75 @@ +import asyncio + +import ydb +import ydb.aio.oidc + +from _smoke_common import ( + approve_keycloak_device, + approve_keycloak_device_async, + execute_async, + execute_sync, + oidc_settings, + required_env, +) + + +def sync_credentials(): + callbacks = [] + + def callback(info): + callbacks.append(info) + print("Device verification URL: {}".format(info.verification_uri_complete or info.verification_uri)) + print("Device user code: {}".format(info.user_code)) + approve_keycloak_device(info) + + return ( + ydb.oidc.OAuth2DeviceCredentials( + client_id=required_env("YDB_OIDC_DEVICE_CLIENT_ID"), + device_authorization_callback=callback, + **oidc_settings(), + ), + callbacks, + ) + + +def async_credentials(): + callbacks = [] + + async def callback(info): + callbacks.append(info) + print("Async device verification URL: {}".format(info.verification_uri_complete or info.verification_uri)) + print("Async device user code: {}".format(info.user_code)) + await approve_keycloak_device_async(info) + + return ( + ydb.aio.oidc.OAuth2DeviceCredentials( + client_id=required_env("YDB_OIDC_DEVICE_CLIENT_ID"), + device_authorization_callback=callback, + **oidc_settings(), + ), + callbacks, + ) + + +async def run_async() -> None: + credentials, callbacks = async_credentials() + await execute_async(credentials, "Device Authorization") + credentials._expires_in = 0 + await execute_async(credentials, "Device Authorization refresh") + if len(callbacks) != 1: + raise RuntimeError("Device Authorization callback ran again instead of using refresh_token") + + +def main() -> None: + credentials, callbacks = sync_credentials() + execute_sync(credentials, "Device Authorization") + credentials._expires_in = 0 + execute_sync(credentials, "Device Authorization refresh") + if len(callbacks) != 1: + raise RuntimeError("Device Authorization callback ran again instead of using refresh_token") + + asyncio.run(run_async()) + + +if __name__ == "__main__": + main() diff --git a/tests/oidc/prepare.sh b/tests/oidc/prepare.sh new file mode 100755 index 000000000..66ba57d54 --- /dev/null +++ b/tests/oidc/prepare.sh @@ -0,0 +1,146 @@ +#!/usr/bin/env bash + +set -euo pipefail + +script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +state_dir="${script_dir}/.state" +env_file="${state_dir}/oidc.env" +realm_file="${state_dir}/keycloak/ydb-realm.json" +force=false + +if [[ "${1:-}" == "--force" ]]; then + force=true +elif [[ $# -ne 0 ]]; then + echo "Usage: $0 [--force]" >&2 + exit 2 +fi + +if [[ -f "${env_file}" && -f "${realm_file}" && "${force}" == false ]]; then + echo "OIDC test configuration already exists: ${env_file}" + exit 0 +fi + +if [[ "${force}" == false && ( -e "${env_file}" || -e "${realm_file}" ) ]]; then + echo "OIDC test state is incomplete; rerun with --force to replace it" >&2 + exit 1 +fi + +if ! command -v openssl >/dev/null 2>&1; then + echo "openssl is required to generate local test credentials" >&2 + exit 1 +fi + +umask 077 +mkdir -p "${state_dir}/certs" "${state_dir}/keycloak" +rm -f "${state_dir}/certs/keycloak.p12" "${state_dir}/certs/keycloak-ca.pem" + +keycloak_admin_password="$(openssl rand -hex 24)" +keycloak_keystore_password="$(openssl rand -hex 24)" +client_secret="$(openssl rand -hex 32)" +device_user_password="$(openssl rand -hex 24)" + +{ + printf '%s\n' \ + "KEYCLOAK_ADMIN_USERNAME=admin" \ + "KEYCLOAK_ADMIN_PASSWORD=${keycloak_admin_password}" \ + "KEYCLOAK_KEYSTORE_PASSWORD=${keycloak_keystore_password}" \ + "YDB_ENDPOINT=grpc://localhost:22136" \ + "YDB_DATABASE=/local" \ + "YDB_OIDC_ISSUER=https://localhost:28443/realms/ydb" \ + "YDB_OIDC_AUDIENCE=ydb" \ + "YDB_OIDC_CA_FILE=${state_dir}/certs/keycloak-ca.pem" \ + "YDB_OIDC_CLIENT_ID=ydb-client-credentials" \ + "YDB_OIDC_CLIENT_SECRET=${client_secret}" \ + "YDB_OIDC_DEVICE_CLIENT_ID=ydb-device" \ + "YDB_OIDC_DEVICE_USERNAME=developer" \ + "YDB_OIDC_DEVICE_PASSWORD=${device_user_password}" +} >"${env_file}" + +{ + printf '%s\n' \ + '{' \ + ' "realm": "ydb",' \ + ' "enabled": true,' \ + ' "displayName": "YDB local OIDC",' \ + ' "accessTokenLifespan": 300,' \ + ' "oauth2DeviceCodeLifespan": 600,' \ + ' "oauth2DevicePollingInterval": 5,' \ + ' "groups": [{"name": "developers"}],' \ + ' "users": [' \ + ' {' \ + ' "username": "developer",' \ + ' "enabled": true,' \ + ' "email": "developer@example.test",' \ + ' "emailVerified": true,' \ + ' "firstName": "YDB",' \ + ' "lastName": "Developer",' \ + " \"credentials\": [{\"type\": \"password\", \"value\": \"${device_user_password}\", \"temporary\": false}]," \ + ' "groups": ["/developers"]' \ + ' }' \ + ' ],' \ + ' "clients": [' \ + ' {' \ + ' "clientId": "ydb-client-credentials",' \ + ' "name": "YDB client credentials",' \ + ' "enabled": true,' \ + ' "protocol": "openid-connect",' \ + ' "publicClient": false,' \ + " \"secret\": \"${client_secret}\"," \ + ' "serviceAccountsEnabled": true,' \ + ' "standardFlowEnabled": false,' \ + ' "directAccessGrantsEnabled": false,' \ + ' "protocolMappers": [' \ + ' {' \ + ' "name": "ydb-audience",' \ + ' "protocol": "openid-connect",' \ + ' "protocolMapper": "oidc-audience-mapper",' \ + ' "config": {' \ + ' "included.custom.audience": "ydb",' \ + ' "id.token.claim": "false",' \ + ' "access.token.claim": "true"' \ + ' }' \ + ' }' \ + ' ]' \ + ' },' \ + ' {' \ + ' "clientId": "ydb-device",' \ + ' "name": "YDB device authorization",' \ + ' "enabled": true,' \ + ' "protocol": "openid-connect",' \ + ' "publicClient": true,' \ + ' "standardFlowEnabled": false,' \ + ' "directAccessGrantsEnabled": false,' \ + ' "attributes": {' \ + ' "oauth2.device.authorization.grant.enabled": "true"' \ + ' },' \ + ' "protocolMappers": [' \ + ' {' \ + ' "name": "ydb-audience",' \ + ' "protocol": "openid-connect",' \ + ' "protocolMapper": "oidc-audience-mapper",' \ + ' "config": {' \ + ' "included.custom.audience": "ydb",' \ + ' "id.token.claim": "false",' \ + ' "access.token.claim": "true"' \ + ' }' \ + ' },' \ + ' {' \ + ' "name": "groups",' \ + ' "protocol": "openid-connect",' \ + ' "protocolMapper": "oidc-group-membership-mapper",' \ + ' "config": {' \ + ' "claim.name": "groups",' \ + ' "full.path": "false",' \ + ' "id.token.claim": "false",' \ + ' "access.token.claim": "true"' \ + ' }' \ + ' }' \ + ' ]' \ + ' }' \ + ' ]' \ + '}' +} >"${realm_file}" + +chmod 0600 "${env_file}" "${realm_file}" +echo "Generated OIDC test configuration: ${env_file}" +echo "Load it with: set -a; source '${env_file}'; set +a" diff --git a/tests/oidc/static_token_smoke.py b/tests/oidc/static_token_smoke.py new file mode 100644 index 000000000..40160e005 --- /dev/null +++ b/tests/oidc/static_token_smoke.py @@ -0,0 +1,25 @@ +import asyncio + +import ydb +import ydb.aio.oidc + +from _smoke_common import assert_token_claims, execute_async, execute_sync, raw_access_token + + +def main() -> None: + access_token = raw_access_token() + assert_token_claims(access_token) + + execute_sync(ydb.oidc.OAuth2TokenCredentials(access_token), "Static token") + asyncio.run(execute_async(ydb.aio.oidc.OAuth2TokenCredentials(access_token), "Static token")) + + try: + execute_sync(ydb.oidc.OAuth2TokenCredentials("this-is-not-a-jwt"), "Invalid static token") + except ydb.ConnectionFailure: + print("Invalid static token rejected by YDB") + else: + raise RuntimeError("YDB unexpectedly accepted an invalid static token") + + +if __name__ == "__main__": + main() diff --git a/tests/oidc/ydb.yaml b/tests/oidc/ydb.yaml new file mode 100644 index 000000000..badf68d9c --- /dev/null +++ b/tests/oidc/ydb.yaml @@ -0,0 +1,330 @@ +actor_system_config: + batch_executor: 2 + executor: + - name: System + spin_threshold: 0 + threads: 2 + type: BASIC + - name: User + spin_threshold: 0 + threads: 3 + type: BASIC + - name: Batch + spin_threshold: 0 + threads: 2 + type: BASIC + - name: IO + threads: 1 + time_per_mailbox_micro_secs: 100 + type: IO + - name: IC + spin_threshold: 10 + threads: 1 + time_per_mailbox_micro_secs: 100 + type: BASIC + io_executor: 3 + scheduler: + progress_threshold: 10000 + resolution: 1024 + spin_threshold: 0 + service_executor: + - executor_id: 4 + service_name: Interconnect + sys_executor: 0 + user_executor: 1 +auth_config: + external_idp_authentication_domain: sso + external_idp_config: + allowed_clock_skew: 30s + audience: ydb + groups_claim_name: groups + issuer: https://localhost:28443/realms/ydb + jwks_cache_settings: + timeout: 2h + subject_claim_name: sub +blob_storage_config: + service_set: + availability_domains: 1 + groups: + - erasure_species: 0 + group_generation: 1 + group_id: 0 + rings: + - fail_domains: + - vdisk_locations: + - node_id: 1 + pdisk_guid: 1 + pdisk_id: 1 + vdisk_slot_id: 0 + pdisks: + - node_id: 1 + path: SectorMap:1:64 + pdisk_category: 0 + pdisk_guid: 1 + pdisk_id: 1 + vdisks: + - vdisk_id: + domain: 0 + group_generation: 1 + group_id: 0 + ring: 0 + vdisk: 0 + vdisk_location: + node_id: 1 + pdisk_guid: 1 + pdisk_id: 1 + vdisk_slot_id: 0 +channel_profile_config: + profile: + - channel: + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + profile_id: 0 + - channel: + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + - erasure_species: none + pdisk_category: 0 + storage_pool_kind: hdd + profile_id: 1 +domains_config: + domain: + - domain_id: 1 + name: local + storage_pool_types: + - kind: hdd + pool_config: + box_id: 1 + erasure_species: none + kind: hdd + pdisk_filter: + - property: + - type: ROT + vdisk_kind: Default + - kind: hdd1 + pool_config: + box_id: 1 + erasure_species: none + kind: hdd + pdisk_filter: + - property: + - type: ROT + vdisk_kind: Default + - kind: hdd2 + pool_config: + box_id: 1 + erasure_species: none + kind: hdd + pdisk_filter: + - property: + - type: ROT + vdisk_kind: Default + - kind: hdde + pool_config: + box_id: 1 + encryption_mode: 1 + erasure_species: none + kind: hdd + pdisk_filter: + - property: + - type: ROT + vdisk_kind: Default + security_config: + enforce_user_token_check_requirement: true + state_storage: + - ring: + nto_select: 1 + ring: + - node: + - 1 + use_ring_specific_node_selection: true + ssid: 1 +feature_flags: + enable_drain_on_shutdown: false + enable_mvcc_snapshot_reads: true + enable_persistent_query_stats: true + enable_public_api_external_blobs: false + enable_scheme_transactions_at_scheme_shard: true +federated_query_config: + audit: + enabled: false + uaconfig: + uri: '' + checkpoint_coordinator: + checkpointing_period_millis: 1000 + enabled: true + max_inflight: 1 + storage: + endpoint: '' + common: + ids_prefix: pt + use_bearer_for_ydb: true + control_plane_proxy: + enabled: true + request_timeout: 30s + control_plane_storage: + available_binding: + - DATA_STREAMS + - OBJECT_STORAGE + available_connection: + - YDB_DATABASE + - CLICKHOUSE_CLUSTER + - DATA_STREAMS + - OBJECT_STORAGE + - MONITORING + enabled: true + storage: + endpoint: '' + db_pool: + enabled: true + storage: + endpoint: '' + enabled: false + gateways: + dq: + default_settings: [] + enabled: true + pq: + cluster_mapping: [] + solomon: + cluster_mapping: [] + nodes_manager: + enabled: true + pending_fetcher: + enabled: true + pinger: + ping_period: 30s + private_api: + enabled: true + private_proxy: + enabled: true + resource_manager: + enabled: true + token_accessor: + enabled: true +grpc_config: + host: '[::]' + services: + - nbs + - legacy + - tablet_service + - yql + - discovery + - cms + - locking + - kesus + - pq + - pqcd + - pqv1 + - topic + - datastreams + - scripting + - clickhouse_internal + - rate_limiter + - analytics + - export + - import + - yq + - keyvalue + - monitoring + - auth + - query_service + - view +interconnect_config: + start_tcp: true +kafka_proxy_config: + enable_kafka_proxy: true + listening_port: 9092 +kqpconfig: + settings: + - name: _ResultRowsLimit + value: '1000' +log_config: + default_level: 5 + entry: [] + sys_log: false +nameservice_config: + node: + - address: ::1 + host: localhost + node_id: 1 + port: 19001 + walle_location: + body: 1 + data_center: '1' + rack: '1' +net_classifier_config: + cms_config_timeout_seconds: 30 + net_data_file_path: /ydb_data/netData.tsv + updater_config: + net_data_update_interval_seconds: 60 + retry_interval_seconds: 30 +pqcluster_discovery_config: + enabled: false +pqconfig: + check_acl: false + cluster_table_path: '' + clusters_update_timeout_sec: 1 + enable_proto_source_id_info: true + enabled: true + max_storage_node_port: 65535 + meta_cache_timeout_sec: 1 + quoting_config: + enable_quoting: false + require_credentials_in_new_protocol: false + root: '' + topics_are_first_class_citizen: true + version_table_path: '' +sqs_config: + enable_dead_letter_queues: true + enable_sqs: false + force_queue_creation_v2: true + force_queue_deletion_v2: true + scheme_cache_hard_refresh_time_seconds: 0 + scheme_cache_soft_refresh_time_seconds: 0 +static_erasure: none +system_tablets: + default_node: + - 1 + flat_schemeshard: + - info: + tablet_id: 72057594046678944 + flat_tx_coordinator: + - node: + - 1 + tx_allocator: + - node: + - 1 + tx_mediator: + - node: + - 1 +table_service_config: + resource_manager: + channel_buffer_size: 262144 + mkql_heavy_program_memory_limit: 1048576 + mkql_light_program_memory_limit: 65536 + verbose_memory_limit_exception: true + sql_version: 1 diff --git a/ydb/__init__.py b/ydb/__init__.py index e0eff6c79..b094c079d 100644 --- a/ydb/__init__.py +++ b/ydb/__init__.py @@ -32,7 +32,7 @@ pass -_LAZY_MODULES = {"iam"} +_LAZY_MODULES = {"iam", "oidc"} def __getattr__(name): diff --git a/ydb/aio/oidc.py b/ydb/aio/oidc.py new file mode 100644 index 000000000..0e515aed5 --- /dev/null +++ b/ydb/aio/oidc.py @@ -0,0 +1,192 @@ +# -*- coding: utf-8 -*- +import asyncio +import inspect +import typing + +import aiohttp + +from ydb import issues +from ydb.aio.credentials import AbstractExpiringTokenCredentials +from ydb.oidc._common import DeviceAuthorizationInfo, OAuth2CredentialsBase +from ydb.oidc.credentials import OAuth2TokenCredentials + + +class _OAuth2Credentials(AbstractExpiringTokenCredentials, OAuth2CredentialsBase): + def __init__( + self, + issuer: str, + client_id: str, + scope: typing.Union[str, typing.Sequence[str], None], + audience: typing.Optional[str], + ca_file: typing.Optional[str], + request_timeout: float, + ): + AbstractExpiringTokenCredentials.__init__(self) + OAuth2CredentialsBase.__init__(self, issuer, client_id, scope, audience, ca_file, request_timeout) + + async def _request_json( + self, + url: str, + data: typing.Optional[typing.Mapping[str, str]] = None, + headers: typing.Optional[typing.Mapping[str, str]] = None, + ) -> typing.Tuple[int, typing.Dict[str, typing.Any]]: + timeout = aiohttp.ClientTimeout(total=self._request_timeout) + ssl_context = self._ssl_context if url.startswith("https://") else None + try: + async with aiohttp.ClientSession(timeout=timeout) as session: + async with session.request( + "POST" if data is not None else "GET", + url, + data=data, + headers=headers, + ssl=ssl_context, + ) as response: + return response.status, self._decode_json(await response.read(), url) + except (aiohttp.ClientError, asyncio.TimeoutError, OSError) as error: + raise issues.Unavailable("OAuth 2.0 endpoint is unavailable at {}: {}".format(url, error)) + + async def _discovery(self) -> typing.Dict[str, typing.Any]: + if self._discovery_document is None: + url = self._issuer + "/.well-known/openid-configuration" + status, response = await self._request_json(url) + self._discovery_document = self._process_discovery_response(status, response) + return self._discovery_document + + +class OAuth2ClientCredentials(_OAuth2Credentials): + """Asynchronous OAuth 2.0 Client Credentials Grant through OIDC discovery.""" + + def __init__( + self, + issuer: str, + client_id: str, + client_secret: str, + scope: typing.Union[str, typing.Sequence[str], None] = None, + audience: typing.Optional[str] = None, + ca_file: typing.Optional[str] = None, + request_timeout: float = 10, + ): + if not client_secret: + raise ValueError("OAuth 2.0 client secret must not be empty") + super(OAuth2ClientCredentials, self).__init__(issuer, client_id, scope, audience, ca_file, request_timeout) + self._client_secret = client_secret + + async def _make_token_request(self): + token_endpoint = (await self._discovery())["token_endpoint"] + headers = {"Authorization": self._client_authorization_header(self._client_id, self._client_secret)} + status, response = await self._request_json(token_endpoint, self._client_credentials_data(), headers) + self._raise_for_status(status, response) + return self._process_token_response(response) + + +class OAuth2DeviceCredentials(_OAuth2Credentials): + """Asynchronous OAuth 2.0 Device Authorization Grant through OIDC discovery.""" + + def __init__( + self, + issuer: str, + client_id: str, + device_authorization_callback: typing.Callable[ + [DeviceAuthorizationInfo], typing.Union[typing.Awaitable[None], None] + ], + scope: typing.Union[str, typing.Sequence[str], None] = "openid", + audience: typing.Optional[str] = None, + client_secret: typing.Optional[str] = None, + ca_file: typing.Optional[str] = None, + request_timeout: float = 10, + device_flow_timeout: typing.Optional[float] = None, + ): + if not callable(device_authorization_callback): + raise ValueError("Device Authorization callback must be callable") + if device_flow_timeout is not None and device_flow_timeout <= 0: + raise ValueError("Device Authorization timeout must be positive") + super(OAuth2DeviceCredentials, self).__init__(issuer, client_id, scope, audience, ca_file, request_timeout) + self._device_authorization_callback = device_authorization_callback + self._client_secret = client_secret + self._device_flow_timeout = device_flow_timeout + self._refresh_token_value: typing.Optional[str] = None + + def _client_headers(self) -> typing.Dict[str, str]: + if self._client_secret is None: + return {} + return {"Authorization": self._client_authorization_header(self._client_id, self._client_secret)} + + def _save_token_response(self, response: typing.Mapping[str, typing.Any]) -> typing.Dict[str, typing.Any]: + refresh_token = response.get("refresh_token") + if refresh_token is not None: + if not isinstance(refresh_token, str) or not refresh_token: + raise issues.Error("OAuth 2.0 token response contains an invalid refresh_token") + self._refresh_token_value = refresh_token + return self._process_token_response(response) + + async def _try_refresh(self, token_endpoint: str) -> typing.Optional[typing.Dict[str, typing.Any]]: + if self._refresh_token_value is None: + return None + status, response = await self._request_json( + token_endpoint, + self._refresh_token_data(self._refresh_token_value), + self._client_headers(), + ) + if status >= 400 and response.get("error") == "invalid_grant": + self._refresh_token_value = None + return None + self._raise_for_status(status, response) + return self._save_token_response(response) + + async def _make_token_request(self): + discovery = await self._discovery() + token_endpoint = discovery["token_endpoint"] + refreshed = await self._try_refresh(token_endpoint) + if refreshed is not None: + return refreshed + + device_endpoint = discovery.get("device_authorization_endpoint") + if not isinstance(device_endpoint, str) or not device_endpoint: + raise issues.Error("OIDC discovery response does not contain a device_authorization_endpoint") + + status, response = await self._request_json( + device_endpoint, + self._device_authorization_data(), + self._client_headers(), + ) + self._raise_for_status(status, response) + device_code, info = self._process_device_authorization_response(response) + callback_result = self._device_authorization_callback(info) + if inspect.isawaitable(callback_result): + await callback_result + + timeout = info.expires_in + if self._device_flow_timeout is not None: + timeout = min(timeout, self._device_flow_timeout) + deadline = asyncio.get_running_loop().time() + timeout + interval = info.interval + + while asyncio.get_running_loop().time() < deadline: + await asyncio.sleep(interval) + status, response = await self._request_json( + token_endpoint, + self._device_token_data(device_code), + self._client_headers(), + ) + if 200 <= status < 300: + return self._save_token_response(response) + + error = response.get("error") + if error == "authorization_pending": + continue + if error == "slow_down": + interval += 5 + continue + if error == "expired_token": + raise issues.Unauthenticated("OAuth 2.0 device code expired") + self._raise_for_status(status, response) + + raise issues.Unauthenticated("OAuth 2.0 Device Authorization timed out") + + +__all__ = [ + "DeviceAuthorizationInfo", + "OAuth2ClientCredentials", + "OAuth2DeviceCredentials", + "OAuth2TokenCredentials", +] diff --git a/ydb/oidc/__init__.py b/ydb/oidc/__init__.py new file mode 100644 index 000000000..6ecbaea2f --- /dev/null +++ b/ydb/oidc/__init__.py @@ -0,0 +1,13 @@ +# -*- coding: utf-8 -*- +from ._common import DeviceAuthorizationInfo # noqa +from .credentials import OAuth2ClientCredentials # noqa +from .credentials import OAuth2DeviceCredentials # noqa +from .credentials import OAuth2TokenCredentials # noqa + + +__all__ = [ + "DeviceAuthorizationInfo", + "OAuth2ClientCredentials", + "OAuth2DeviceCredentials", + "OAuth2TokenCredentials", +] diff --git a/ydb/oidc/_common.py b/ydb/oidc/_common.py new file mode 100644 index 000000000..a1cc2f2c0 --- /dev/null +++ b/ydb/oidc/_common.py @@ -0,0 +1,193 @@ +# -*- coding: utf-8 -*- +import base64 +import json +import os +import ssl +import typing +from dataclasses import dataclass +from urllib.parse import quote_plus + +from ydb import issues + + +DEVICE_GRANT_TYPE = "urn:ietf:params:oauth:grant-type:device_code" + + +@dataclass(frozen=True) +class DeviceAuthorizationInfo: + verification_uri: str + user_code: str + verification_uri_complete: typing.Optional[str] + expires_in: int + interval: int + + +class OAuth2CredentialsBase: + def __init__( + self, + issuer: str, + client_id: str, + scope: typing.Union[str, typing.Sequence[str], None] = None, + audience: typing.Optional[str] = None, + ca_file: typing.Optional[str] = None, + request_timeout: float = 10, + ): + if not issuer: + raise ValueError("OAuth 2.0 issuer must not be empty") + if not client_id: + raise ValueError("OAuth 2.0 client ID must not be empty") + if request_timeout <= 0: + raise ValueError("OAuth 2.0 request timeout must be positive") + + self._issuer = issuer.rstrip("/") + self._client_id = client_id + self._scope = self._scope_parameter(scope) + self._audience = audience + self._request_timeout = request_timeout + self._discovery_document: typing.Optional[typing.Dict[str, typing.Any]] = None + self._ssl_context = ssl.create_default_context(cafile=os.path.expanduser(ca_file) if ca_file else None) + + @staticmethod + def _scope_parameter(scope: typing.Union[str, typing.Sequence[str], None]) -> typing.Optional[str]: + if scope is None or isinstance(scope, str): + return scope + return " ".join(scope) + + @staticmethod + def _decode_json(content: bytes, url: str) -> typing.Dict[str, typing.Any]: + try: + value = json.loads(content.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + raise issues.Error("OAuth 2.0 endpoint returned invalid JSON from {}: {}".format(url, error)) + if not isinstance(value, dict): + raise issues.Error("OAuth 2.0 endpoint returned a non-object JSON response from {}".format(url)) + return value + + @staticmethod + def _raise_for_status(status: int, response: typing.Mapping[str, typing.Any]) -> None: + if 200 <= status < 300: + return + + error = response.get("error", "unknown_error") + description = response.get("error_description") + message = "OAuth 2.0 request failed: {}".format(error) + if description: + message += ": {}".format(description) + + if status in (401, 403) or error in ("invalid_client", "access_denied"): + raise issues.Unauthenticated(message) + if status >= 500: + raise issues.Unavailable(message) + if status >= 400: + raise issues.BadRequest(message) + raise issues.Error(message) + + def _process_discovery_response( + self, status: int, response: typing.Dict[str, typing.Any] + ) -> typing.Dict[str, typing.Any]: + self._raise_for_status(status, response) + discovered_issuer = response.get("issuer") + if discovered_issuer != self._issuer: + raise issues.Error( + "OIDC discovery issuer mismatch: expected {!r}, got {!r}".format(self._issuer, discovered_issuer) + ) + token_endpoint = response.get("token_endpoint") + if not isinstance(token_endpoint, str) or not token_endpoint: + raise issues.Error("OIDC discovery response does not contain a token_endpoint") + return response + + @staticmethod + def _process_token_response(response: typing.Mapping[str, typing.Any]) -> typing.Dict[str, typing.Any]: + access_token = response.get("access_token") + token_type = response.get("token_type") + expires_in = response.get("expires_in") + + if not isinstance(access_token, str) or not access_token: + raise issues.Error("OAuth 2.0 token response does not contain an access_token") + if not isinstance(token_type, str) or token_type.lower() != "bearer": + raise issues.Error("OAuth 2.0 token response contains an unsupported token_type: {!r}".format(token_type)) + if isinstance(expires_in, bool) or not isinstance(expires_in, (int, float)) or expires_in <= 0: + raise issues.Error("OAuth 2.0 token response contains an invalid expires_in: {!r}".format(expires_in)) + + return {"access_token": "Bearer " + access_token, "expires_in": expires_in} + + @staticmethod + def _client_authorization_header(client_id: str, client_secret: str) -> str: + encoded_id = quote_plus(client_id) + encoded_secret = quote_plus(client_secret) + value = base64.b64encode("{}:{}".format(encoded_id, encoded_secret).encode("utf-8")).decode("ascii") + return "Basic " + value + + def _client_credentials_data(self) -> typing.Dict[str, str]: + data = {"grant_type": "client_credentials"} + if self._scope: + data["scope"] = self._scope + if self._audience: + data["audience"] = self._audience + return data + + def _device_authorization_data(self) -> typing.Dict[str, str]: + data = {"client_id": self._client_id} + if self._scope: + data["scope"] = self._scope + if self._audience: + data["audience"] = self._audience + return data + + def _device_token_data(self, device_code: str) -> typing.Dict[str, str]: + return { + "grant_type": DEVICE_GRANT_TYPE, + "client_id": self._client_id, + "device_code": device_code, + } + + def _refresh_token_data(self, refresh_token: str) -> typing.Dict[str, str]: + data = { + "grant_type": "refresh_token", + "client_id": self._client_id, + "refresh_token": refresh_token, + } + if self._scope: + data["scope"] = self._scope + return data + + @staticmethod + def _process_device_authorization_response( + response: typing.Mapping[str, typing.Any], + ) -> typing.Tuple[str, DeviceAuthorizationInfo]: + device_code = response.get("device_code") + user_code = response.get("user_code") + verification_uri = response.get("verification_uri") + verification_uri_complete = response.get("verification_uri_complete") + expires_in = response.get("expires_in") + interval = response.get("interval", 5) + + if not isinstance(device_code, str) or not device_code: + raise issues.Error("Device Authorization response does not contain a device_code") + if not isinstance(user_code, str) or not user_code: + raise issues.Error("Device Authorization response does not contain a user_code") + if not isinstance(verification_uri, str) or not verification_uri: + raise issues.Error("Device Authorization response does not contain a verification_uri") + if verification_uri_complete is not None and not isinstance(verification_uri_complete, str): + raise issues.Error("Device Authorization response contains an invalid verification_uri_complete") + if isinstance(expires_in, bool) or not isinstance(expires_in, (int, float)) or expires_in <= 0: + raise issues.Error("Device Authorization response contains an invalid expires_in: {!r}".format(expires_in)) + if isinstance(interval, bool) or not isinstance(interval, (int, float)) or interval <= 0: + raise issues.Error("Device Authorization response contains an invalid interval: {!r}".format(interval)) + + info = DeviceAuthorizationInfo( + verification_uri=verification_uri, + user_code=user_code, + verification_uri_complete=verification_uri_complete, + expires_in=int(expires_in), + interval=int(interval), + ) + return device_code, info + + +def bearer_token(token: str) -> str: + if not isinstance(token, str) or not token: + raise ValueError("OAuth 2.0 access token must not be empty") + if token.lower().startswith("bearer "): + return token + return "Bearer " + token diff --git a/ydb/oidc/credentials.py b/ydb/oidc/credentials.py new file mode 100644 index 000000000..0530d5ccf --- /dev/null +++ b/ydb/oidc/credentials.py @@ -0,0 +1,203 @@ +# -*- coding: utf-8 -*- +import socket +import time +import typing +import urllib.error +import urllib.parse +import urllib.request + +from ydb import credentials, issues, tracing + +from ._common import DeviceAuthorizationInfo, OAuth2CredentialsBase, bearer_token + + +class OAuth2TokenCredentials(credentials.Credentials): + """Credentials for an OAuth 2.0 access token obtained outside the SDK.""" + + def __init__(self, access_token: str, tracer=None): + super(OAuth2TokenCredentials, self).__init__(tracer) + self._access_token = bearer_token(access_token) + + def auth_metadata(self): + return [(credentials.YDB_AUTH_TICKET_HEADER, self._access_token)] + + +class _OAuth2Credentials(credentials.AbstractExpiringTokenCredentials, OAuth2CredentialsBase): + def __init__( + self, + issuer: str, + client_id: str, + scope: typing.Union[str, typing.Sequence[str], None], + audience: typing.Optional[str], + ca_file: typing.Optional[str], + request_timeout: float, + tracer=None, + ): + credentials.AbstractExpiringTokenCredentials.__init__(self, tracer) + OAuth2CredentialsBase.__init__(self, issuer, client_id, scope, audience, ca_file, request_timeout) + + def _request_json( + self, + url: str, + data: typing.Optional[typing.Mapping[str, str]] = None, + headers: typing.Optional[typing.Mapping[str, str]] = None, + ) -> typing.Tuple[int, typing.Dict[str, typing.Any]]: + body = urllib.parse.urlencode(data).encode("utf-8") if data is not None else None + request_headers = dict(headers or {}) + if body is not None: + request_headers.setdefault("Content-Type", "application/x-www-form-urlencoded") + request = urllib.request.Request(url, data=body, headers=request_headers) + + try: + with urllib.request.urlopen( + request, + context=self._ssl_context, + timeout=self._request_timeout, + ) as response: + return response.status, self._decode_json(response.read(), url) + except urllib.error.HTTPError as error: + return error.code, self._decode_json(error.read(), url) + except (urllib.error.URLError, TimeoutError, socket.timeout, OSError) as error: + raise issues.Unavailable("OAuth 2.0 endpoint is unavailable at {}: {}".format(url, error)) + + def _discovery(self) -> typing.Dict[str, typing.Any]: + if self._discovery_document is None: + url = self._issuer + "/.well-known/openid-configuration" + status, response = self._request_json(url) + self._discovery_document = self._process_discovery_response(status, response) + return self._discovery_document + + +class OAuth2ClientCredentials(_OAuth2Credentials): + """OAuth 2.0 Client Credentials Grant discovered through an OIDC issuer.""" + + def __init__( + self, + issuer: str, + client_id: str, + client_secret: str, + scope: typing.Union[str, typing.Sequence[str], None] = None, + audience: typing.Optional[str] = None, + ca_file: typing.Optional[str] = None, + request_timeout: float = 10, + tracer=None, + ): + if not client_secret: + raise ValueError("OAuth 2.0 client secret must not be empty") + super(OAuth2ClientCredentials, self).__init__( + issuer, client_id, scope, audience, ca_file, request_timeout, tracer + ) + self._client_secret = client_secret + + @tracing.with_trace() + def _make_token_request(self): + token_endpoint = self._discovery()["token_endpoint"] + headers = {"Authorization": self._client_authorization_header(self._client_id, self._client_secret)} + status, response = self._request_json(token_endpoint, self._client_credentials_data(), headers) + self._raise_for_status(status, response) + return self._process_token_response(response) + + +class OAuth2DeviceCredentials(_OAuth2Credentials): + """OAuth 2.0 Device Authorization Grant discovered through an OIDC issuer.""" + + def __init__( + self, + issuer: str, + client_id: str, + device_authorization_callback: typing.Callable[[DeviceAuthorizationInfo], None], + scope: typing.Union[str, typing.Sequence[str], None] = "openid", + audience: typing.Optional[str] = None, + client_secret: typing.Optional[str] = None, + ca_file: typing.Optional[str] = None, + request_timeout: float = 10, + device_flow_timeout: typing.Optional[float] = None, + tracer=None, + ): + if not callable(device_authorization_callback): + raise ValueError("Device Authorization callback must be callable") + if device_flow_timeout is not None and device_flow_timeout <= 0: + raise ValueError("Device Authorization timeout must be positive") + super(OAuth2DeviceCredentials, self).__init__( + issuer, client_id, scope, audience, ca_file, request_timeout, tracer + ) + self._device_authorization_callback = device_authorization_callback + self._client_secret = client_secret + self._device_flow_timeout = device_flow_timeout + self._refresh_token_value: typing.Optional[str] = None + + def _client_headers(self) -> typing.Dict[str, str]: + if self._client_secret is None: + return {} + return {"Authorization": self._client_authorization_header(self._client_id, self._client_secret)} + + def _save_token_response(self, response: typing.Mapping[str, typing.Any]) -> typing.Dict[str, typing.Any]: + refresh_token = response.get("refresh_token") + if refresh_token is not None: + if not isinstance(refresh_token, str) or not refresh_token: + raise issues.Error("OAuth 2.0 token response contains an invalid refresh_token") + self._refresh_token_value = refresh_token + return self._process_token_response(response) + + def _try_refresh(self, token_endpoint: str) -> typing.Optional[typing.Dict[str, typing.Any]]: + if self._refresh_token_value is None: + return None + status, response = self._request_json( + token_endpoint, + self._refresh_token_data(self._refresh_token_value), + self._client_headers(), + ) + if status >= 400 and response.get("error") == "invalid_grant": + self._refresh_token_value = None + return None + self._raise_for_status(status, response) + return self._save_token_response(response) + + @tracing.with_trace() + def _make_token_request(self): + discovery = self._discovery() + token_endpoint = discovery["token_endpoint"] + refreshed = self._try_refresh(token_endpoint) + if refreshed is not None: + return refreshed + + device_endpoint = discovery.get("device_authorization_endpoint") + if not isinstance(device_endpoint, str) or not device_endpoint: + raise issues.Error("OIDC discovery response does not contain a device_authorization_endpoint") + + status, response = self._request_json( + device_endpoint, + self._device_authorization_data(), + self._client_headers(), + ) + self._raise_for_status(status, response) + device_code, info = self._process_device_authorization_response(response) + self._device_authorization_callback(info) + + timeout = info.expires_in + if self._device_flow_timeout is not None: + timeout = min(timeout, self._device_flow_timeout) + deadline = time.monotonic() + timeout + interval = info.interval + + while time.monotonic() < deadline: + time.sleep(interval) + status, response = self._request_json( + token_endpoint, + self._device_token_data(device_code), + self._client_headers(), + ) + if 200 <= status < 300: + return self._save_token_response(response) + + error = response.get("error") + if error == "authorization_pending": + continue + if error == "slow_down": + interval += 5 + continue + if error == "expired_token": + raise issues.Unauthenticated("OAuth 2.0 device code expired") + self._raise_for_status(status, response) + + raise issues.Unauthenticated("OAuth 2.0 Device Authorization timed out") From 51bb664b509af54481a6f0c525b0805839f2ca09 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Mon, 21 Sep 2026 13:12:45 +0300 Subject: [PATCH 2/7] test: cover OIDC credential error paths --- tests/aio/test_credentials.py | 140 ++++++++++++++++++++ tests/auth/test_credentials.py | 230 +++++++++++++++++++++++++++++++++ 2 files changed, 370 insertions(+) diff --git a/tests/aio/test_credentials.py b/tests/aio/test_credentials.py index 4e9d0f376..c56af9246 100644 --- a/tests/aio/test_credentials.py +++ b/tests/aio/test_credentials.py @@ -6,6 +6,7 @@ import os import json import asyncio +import aiohttp from unittest.mock import patch, AsyncMock, MagicMock import tests.auth.test_credentials @@ -15,6 +16,7 @@ import ydb.aio.oidc import ydb.aio.oauth2_token_exchange import ydb.oauth2_token_exchange.token_source +from ydb import issues class ServiceAccountCredentialsForTest(ydb.aio.iam.ServiceAccountCredentials): @@ -320,3 +322,141 @@ async def callback(info): assert callback_values[0].user_code == "user-code" assert sleep.await_count == 2 + + +@pytest.mark.asyncio +async def test_oauth2_async_http_requests_and_discovery_cache(): + credentials = ydb.aio.oidc.OAuth2ClientCredentials( + "https://issuer.example", + "client-id", + "client-secret", + ) + response = MagicMock(status=200) + response.read = AsyncMock( + return_value=b'{"issuer":"https://issuer.example","token_endpoint":"https://issuer.example/token"}' + ) + request_context = MagicMock() + request_context.__aenter__ = AsyncMock(return_value=response) + request_context.__aexit__ = AsyncMock(return_value=None) + session = MagicMock() + session.request.return_value = request_context + session_context = MagicMock() + session_context.__aenter__ = AsyncMock(return_value=session) + session_context.__aexit__ = AsyncMock(return_value=None) + + with patch("ydb.aio.oidc.aiohttp.ClientSession", return_value=session_context) as client_session: + first = await credentials._discovery() + second = await credentials._discovery() + + assert first is second + assert client_session.call_count == 1 + assert session.request.call_args.kwargs["ssl"] is not None + + with patch( + "ydb.aio.oidc.aiohttp.ClientSession", + side_effect=aiohttp.ClientError("unavailable"), + ): + with pytest.raises(issues.Unavailable): + await credentials._request_json("http://issuer.example/token", {"key": "value"}) + + +@pytest.mark.asyncio +async def test_oauth2_async_device_refresh_and_error_paths(): + with pytest.raises(ValueError): + ydb.aio.oidc.OAuth2ClientCredentials("https://issuer.example", "client-id", "") + with pytest.raises(ValueError): + ydb.aio.oidc.OAuth2DeviceCredentials("https://issuer.example", "client-id", None) + with pytest.raises(ValueError): + ydb.aio.oidc.OAuth2DeviceCredentials( + "https://issuer.example", + "client-id", + lambda info: None, + device_flow_timeout=0, + ) + + credentials = ydb.aio.oidc.OAuth2DeviceCredentials( + "https://issuer.example", + "client-id", + lambda info: None, + client_secret="client-secret", + ) + assert credentials._client_headers()["Authorization"].startswith("Basic ") + with pytest.raises(issues.Error, match="refresh_token"): + credentials._save_token_response( + { + "access_token": "token", + "token_type": "Bearer", + "expires_in": 300, + "refresh_token": "", + } + ) + + credentials._refresh_token_value = "refresh-token" + credentials._request_json = AsyncMock( + return_value=( + 200, + { + "access_token": "refreshed-token", + "token_type": "Bearer", + "expires_in": 300, + "refresh_token": "new-refresh-token", + }, + ) + ) + assert await credentials._try_refresh("https://issuer.example/token") == { + "access_token": "Bearer refreshed-token", + "expires_in": 300, + } + assert credentials._refresh_token_value == "new-refresh-token" + + credentials._request_json = AsyncMock(return_value=(400, {"error": "invalid_grant"})) + assert await credentials._try_refresh("https://issuer.example/token") is None + assert credentials._refresh_token_value is None + + credentials._discovery_document = {"token_endpoint": "https://issuer.example/token"} + with pytest.raises(issues.Error, match="device_authorization_endpoint"): + await credentials._make_token_request() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "token_response, expected_message", + [ + ({"error": "expired_token"}, "expired"), + (None, "timed out"), + ], +) +async def test_oauth2_async_device_expiration(token_response, expected_message): + callback_values = [] + credentials = ydb.aio.oidc.OAuth2DeviceCredentials( + "https://issuer.example", + "client-id", + callback_values.append, + device_flow_timeout=1, + ) + credentials._discovery_document = { + "token_endpoint": "https://issuer.example/token", + "device_authorization_endpoint": "https://issuer.example/device", + } + device_response = { + "device_code": "device-code", + "user_code": "user-code", + "verification_uri": "https://issuer.example/verify", + "expires_in": 600, + "interval": 1, + } + responses = [(200, device_response)] + loop = MagicMock() + loop.time.side_effect = [0, 2] + if token_response is not None: + responses.append((400, token_response)) + loop.time.side_effect = [0, 0] + credentials._request_json = AsyncMock(side_effect=responses) + + with patch("ydb.aio.oidc.asyncio.get_running_loop", return_value=loop), patch( + "ydb.aio.oidc.asyncio.sleep", new=AsyncMock() + ): + with pytest.raises(issues.Unauthenticated, match=expected_message): + await credentials._make_token_request() + + assert callback_values[0].user_code == "user-code" diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index 3207beeb3..d62b750ac 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -1,11 +1,18 @@ import jwt import concurrent.futures import grpc +import io +import pytest import time +import urllib.error from unittest.mock import patch +from unittest.mock import MagicMock import ydb.iam import ydb.oidc +from ydb import issues +from ydb.oidc._common import OAuth2CredentialsBase +from ydb.oidc._common import bearer_token from yandex.cloud.iam.v1 import iam_token_service_pb2_grpc from yandex.cloud.iam.v1 import iam_token_service_pb2 @@ -189,3 +196,226 @@ def request_json(url, data=None, headers=None): } assert requests[-1][1]["grant_type"] == "refresh_token" assert requests[-1][1]["refresh_token"] == "refresh-token" + + +@pytest.mark.parametrize( + "kwargs", + [ + {"issuer": ""}, + {"client_id": ""}, + {"request_timeout": 0}, + {"client_secret": ""}, + ], +) +def test_oauth2_client_credentials_validation(kwargs): + arguments = { + "issuer": "https://issuer.example", + "client_id": "client-id", + "client_secret": "client-secret", + } + arguments.update(kwargs) + + with pytest.raises(ValueError): + ydb.oidc.OAuth2ClientCredentials(**arguments) + + +@pytest.mark.parametrize("token", ["", None]) +def test_oauth2_token_credentials_validation(token): + with pytest.raises(ValueError): + bearer_token(token) + + +def test_oauth2_common_response_processing(): + credentials = OAuth2CredentialsBase("https://issuer.example/", "client id", audience="ydb") + + with pytest.raises(issues.Error, match="invalid JSON"): + credentials._decode_json(b"not-json", "https://issuer.example/token") + with pytest.raises(issues.Error, match="non-object"): + credentials._decode_json(b"[]", "https://issuer.example/token") + + credentials._raise_for_status(204, {}) + for status, response, error_type in ( + (401, {}, issues.Unauthenticated), + (400, {"error": "invalid_client"}, issues.Unauthenticated), + (503, {}, issues.Unavailable), + (400, {"error": "bad_request", "error_description": "details"}, issues.BadRequest), + (300, {}, issues.Error), + ): + with pytest.raises(error_type): + credentials._raise_for_status(status, response) + + with pytest.raises(issues.Error, match="issuer mismatch"): + credentials._process_discovery_response( + 200, + {"issuer": "https://other.example", "token_endpoint": "https://issuer.example/token"}, + ) + with pytest.raises(issues.Error, match="token_endpoint"): + credentials._process_discovery_response(200, {"issuer": "https://issuer.example"}) + + for response, message in ( + ({}, "access_token"), + ({"access_token": "token", "token_type": "Basic", "expires_in": 60}, "token_type"), + ({"access_token": "token", "token_type": "Bearer", "expires_in": True}, "expires_in"), + ): + with pytest.raises(issues.Error, match=message): + credentials._process_token_response(response) + + assert credentials._client_credentials_data() == { + "grant_type": "client_credentials", + "audience": "ydb", + } + assert credentials._device_authorization_data() == { + "client_id": "client id", + "audience": "ydb", + } + assert credentials._refresh_token_data("refresh-token") == { + "grant_type": "refresh_token", + "client_id": "client id", + "refresh_token": "refresh-token", + } + assert ( + credentials._client_authorization_header("client id", "secret/value") + == "Basic Y2xpZW50K2lkOnNlY3JldCUyRnZhbHVl" + ) + + +@pytest.mark.parametrize( + "invalid_value, message", + [ + ({"device_code": ""}, "device_code"), + ({"user_code": ""}, "user_code"), + ({"verification_uri": ""}, "verification_uri"), + ({"verification_uri_complete": 42}, "verification_uri_complete"), + ({"expires_in": 0}, "expires_in"), + ({"interval": 0}, "interval"), + ], +) +def test_oauth2_device_authorization_response_validation(invalid_value, message): + response = { + "device_code": "device-code", + "user_code": "user-code", + "verification_uri": "https://issuer.example/verify", + "verification_uri_complete": "https://issuer.example/verify?user_code=user-code", + "expires_in": 600, + "interval": 5, + } + response.update(invalid_value) + + with pytest.raises(issues.Error, match=message): + OAuth2CredentialsBase._process_device_authorization_response(response) + + +def test_oauth2_sync_http_requests_and_discovery_cache(): + credentials = ydb.oidc.OAuth2ClientCredentials( + "https://issuer.example", + "client-id", + "client-secret", + ) + response = MagicMock(status=200) + response.read.return_value = b'{"issuer":"https://issuer.example","token_endpoint":"https://issuer.example/token"}' + response_context = MagicMock() + response_context.__enter__.return_value = response + + with patch("ydb.oidc.credentials.urllib.request.urlopen", return_value=response_context) as urlopen: + first = credentials._discovery() + second = credentials._discovery() + + assert first is second + assert urlopen.call_count == 1 + + http_error = urllib.error.HTTPError( + "https://issuer.example/token", + 400, + "Bad Request", + {}, + io.BytesIO(b'{"error":"invalid_request"}'), + ) + with patch("ydb.oidc.credentials.urllib.request.urlopen", side_effect=http_error): + assert credentials._request_json("https://issuer.example/token", {"key": "value"}) == ( + 400, + {"error": "invalid_request"}, + ) + + with patch( + "ydb.oidc.credentials.urllib.request.urlopen", + side_effect=urllib.error.URLError("unavailable"), + ): + with pytest.raises(issues.Unavailable): + credentials._request_json("https://issuer.example/token") + + +def test_oauth2_sync_device_error_paths(): + with pytest.raises(ValueError): + ydb.oidc.OAuth2DeviceCredentials("https://issuer.example", "client-id", None) + with pytest.raises(ValueError): + ydb.oidc.OAuth2DeviceCredentials( + "https://issuer.example", + "client-id", + lambda info: None, + device_flow_timeout=0, + ) + + credentials = ydb.oidc.OAuth2DeviceCredentials( + "https://issuer.example", + "client-id", + lambda info: None, + client_secret="client-secret", + ) + assert credentials._client_headers()["Authorization"].startswith("Basic ") + with pytest.raises(issues.Error, match="refresh_token"): + credentials._save_token_response( + { + "access_token": "token", + "token_type": "Bearer", + "expires_in": 300, + "refresh_token": "", + } + ) + + credentials._refresh_token_value = "refresh-token" + credentials._request_json = MagicMock(return_value=(400, {"error": "invalid_grant"})) + assert credentials._try_refresh("https://issuer.example/token") is None + assert credentials._refresh_token_value is None + + credentials._discovery_document = {"token_endpoint": "https://issuer.example/token"} + with pytest.raises(issues.Error, match="device_authorization_endpoint"): + credentials._make_token_request() + + +@pytest.mark.parametrize( + "token_response, expected_message", + [ + ({"error": "expired_token"}, "expired"), + (None, "timed out"), + ], +) +def test_oauth2_sync_device_expiration(token_response, expected_message): + credentials = ydb.oidc.OAuth2DeviceCredentials( + "https://issuer.example", + "client-id", + lambda info: None, + device_flow_timeout=1, + ) + credentials._discovery_document = { + "token_endpoint": "https://issuer.example/token", + "device_authorization_endpoint": "https://issuer.example/device", + } + device_response = { + "device_code": "device-code", + "user_code": "user-code", + "verification_uri": "https://issuer.example/verify", + "expires_in": 600, + "interval": 1, + } + responses = [(200, device_response)] + monotonic_values = [0, 2] + if token_response is not None: + responses.append((400, token_response)) + monotonic_values = [0, 0] + credentials._request_json = MagicMock(side_effect=responses) + + with patch("ydb.oidc.credentials.time.monotonic", side_effect=monotonic_values), patch( + "ydb.oidc.credentials.time.sleep" + ): + with pytest.raises(issues.Unauthenticated, match=expected_message): + credentials._make_token_request() From 27cf89a3c5027b9d1d6e03d18d74b1d438055c1a Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Mon, 21 Sep 2026 13:45:21 +0300 Subject: [PATCH 3/7] test: reach full OIDC coverage --- tests/aio/test_credentials.py | 7 +++++-- tests/auth/test_credentials.py | 8 ++++++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/tests/aio/test_credentials.py b/tests/aio/test_credentials.py index c56af9246..e5fda9cc9 100644 --- a/tests/aio/test_credentials.py +++ b/tests/aio/test_credentials.py @@ -313,6 +313,7 @@ async def callback(info): }, ), (400, {"error": "authorization_pending"}), + (400, {"error": "slow_down"}), (200, {"access_token": "access-token", "token_type": "Bearer", "expires_in": 300}), ] ) @@ -321,7 +322,7 @@ async def callback(info): assert await credentials.get_auth_token() == "Bearer access-token" assert callback_values[0].user_code == "user-code" - assert sleep.await_count == 2 + assert sleep.await_count == 3 @pytest.mark.asyncio @@ -403,7 +404,8 @@ async def test_oauth2_async_device_refresh_and_error_paths(): }, ) ) - assert await credentials._try_refresh("https://issuer.example/token") == { + credentials._discovery_document = {"token_endpoint": "https://issuer.example/token"} + assert await credentials._make_token_request() == { "access_token": "Bearer refreshed-token", "expires_in": 300, } @@ -423,6 +425,7 @@ async def test_oauth2_async_device_refresh_and_error_paths(): "token_response, expected_message", [ ({"error": "expired_token"}, "expired"), + ({"error": "access_denied"}, "access_denied"), (None, "timed out"), ], ) diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index d62b750ac..63b754f5f 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -371,6 +371,13 @@ def test_oauth2_sync_device_error_paths(): "refresh_token": "", } ) + assert credentials._save_token_response( + { + "access_token": "token", + "token_type": "Bearer", + "expires_in": 300, + } + ) == {"access_token": "Bearer token", "expires_in": 300} credentials._refresh_token_value = "refresh-token" credentials._request_json = MagicMock(return_value=(400, {"error": "invalid_grant"})) @@ -386,6 +393,7 @@ def test_oauth2_sync_device_error_paths(): "token_response, expected_message", [ ({"error": "expired_token"}, "expired"), + ({"error": "access_denied"}, "access_denied"), (None, "timed out"), ], ) From 25ffa7a6ae94ff3b9defb21dabe8c181f885c846 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Mon, 21 Sep 2026 15:49:10 +0300 Subject: [PATCH 4/7] fix: align OIDC providers with C++ SDK --- docs/driver.rst | 6 ++-- tests/aio/test_credentials.py | 6 ++-- tests/auth/test_credentials.py | 45 +++++++++++++++++++++++-- ydb/aio/oidc.py | 2 +- ydb/oidc/_common.py | 61 ++++++++++++++++++++++++++-------- ydb/oidc/credentials.py | 2 +- 6 files changed, 99 insertions(+), 23 deletions(-) diff --git a/docs/driver.rst b/docs/driver.rst index 4ade89a0d..903ef139b 100644 --- a/docs/driver.rst +++ b/docs/driver.rst @@ -223,14 +223,16 @@ code to the user: issuer="https://identity.example.com/realms/example", client_id="public-device-client", device_authorization_callback=show_device_authorization, - scope=["openid", "offline_access"], + scope=["offline_access"], ca_file="/path/to/idp-ca.pem", ) Client Credentials obtains a new access token when needed. Device Authorization uses a returned refresh token for subsequent refreshes and starts a new user interaction if the refresh token is no longer valid. Access tokens are sent to YDB using the -``Bearer`` authentication scheme. +``Bearer`` authentication scheme. The providers require HTTPS issuer and endpoint URLs, +require the discovered issuer to match the configured value exactly, and add the +``openid`` scope when it is absent. Non-blocking counterparts are available as ``ydb.aio.oidc.OAuth2ClientCredentials`` and ``ydb.aio.oidc.OAuth2DeviceCredentials``. The asynchronous Device Authorization diff --git a/tests/aio/test_credentials.py b/tests/aio/test_credentials.py index e5fda9cc9..8969646e9 100644 --- a/tests/aio/test_credentials.py +++ b/tests/aio/test_credentials.py @@ -281,6 +281,7 @@ async def test_oauth2_client_credentials(): assert await credentials.get_auth_token() == "Bearer access-token" assert credentials._request_json.await_count == 2 + assert credentials._request_json.await_args_list[1].args[1]["scope"] == "openid" @pytest.mark.asyncio @@ -328,13 +329,13 @@ async def callback(info): @pytest.mark.asyncio async def test_oauth2_async_http_requests_and_discovery_cache(): credentials = ydb.aio.oidc.OAuth2ClientCredentials( - "https://issuer.example", + "https://issuer.example/", "client-id", "client-secret", ) response = MagicMock(status=200) response.read = AsyncMock( - return_value=b'{"issuer":"https://issuer.example","token_endpoint":"https://issuer.example/token"}' + return_value=b'{"issuer":"https://issuer.example/","token_endpoint":"https://issuer.example/token"}' ) request_context = MagicMock() request_context.__aenter__ = AsyncMock(return_value=response) @@ -351,6 +352,7 @@ async def test_oauth2_async_http_requests_and_discovery_cache(): assert first is second assert client_session.call_count == 1 + assert session.request.call_args.args[1] == "https://issuer.example/.well-known/openid-configuration" assert session.request.call_args.kwargs["ssl"] is not None with patch( diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index 63b754f5f..3ec2b740f 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -105,7 +105,7 @@ def test_oauth2_client_credentials(): issuer, "client-id", "client-secret", - scope=["openid", "profile"], + scope=["profile"], audience="ydb", ) @@ -120,7 +120,7 @@ def request_json(url, data=None, headers=None): assert requests[1][0] == issuer + "/token" assert requests[1][1] == { "grant_type": "client_credentials", - "scope": "openid profile", + "scope": "profile openid", "audience": "ydb", } assert requests[1][2]["Authorization"].startswith("Basic ") @@ -202,6 +202,9 @@ def request_json(url, data=None, headers=None): "kwargs", [ {"issuer": ""}, + {"issuer": "http://issuer.example"}, + {"issuer": "https://issuer.example?tenant=test"}, + {"issuer": "https://issuer.example:invalid"}, {"client_id": ""}, {"request_timeout": 0}, {"client_secret": ""}, @@ -228,6 +231,12 @@ def test_oauth2_token_credentials_validation(token): def test_oauth2_common_response_processing(): credentials = OAuth2CredentialsBase("https://issuer.example/", "client id", audience="ydb") + assert credentials._issuer == "https://issuer.example/" + assert credentials._scope == "openid" + assert OAuth2CredentialsBase._scope_parameter("profile") == "profile openid" + assert OAuth2CredentialsBase._scope_parameter(["openid", "profile"]) == "openid profile" + assert not OAuth2CredentialsBase._is_https_url("https://issuer.example/\n") + with pytest.raises(issues.Error, match="invalid JSON"): credentials._decode_json(b"not-json", "https://issuer.example/token") with pytest.raises(issues.Error, match="non-object"): @@ -250,7 +259,21 @@ def test_oauth2_common_response_processing(): {"issuer": "https://other.example", "token_endpoint": "https://issuer.example/token"}, ) with pytest.raises(issues.Error, match="token_endpoint"): - credentials._process_discovery_response(200, {"issuer": "https://issuer.example"}) + credentials._process_discovery_response(200, {"issuer": "https://issuer.example/"}) + with pytest.raises(issues.Error, match="token_endpoint URL"): + credentials._process_discovery_response( + 200, + {"issuer": "https://issuer.example/", "token_endpoint": "http://issuer.example/token"}, + ) + with pytest.raises(issues.Error, match="device_authorization_endpoint URL"): + credentials._process_discovery_response( + 200, + { + "issuer": "https://issuer.example/", + "token_endpoint": "https://issuer.example/token", + "device_authorization_endpoint": "http://issuer.example/device", + }, + ) for response, message in ( ({}, "access_token"), @@ -262,22 +285,38 @@ def test_oauth2_common_response_processing(): assert credentials._client_credentials_data() == { "grant_type": "client_credentials", + "scope": "openid", "audience": "ydb", } assert credentials._device_authorization_data() == { "client_id": "client id", + "scope": "openid", "audience": "ydb", } assert credentials._refresh_token_data("refresh-token") == { "grant_type": "refresh_token", "client_id": "client id", "refresh_token": "refresh-token", + "scope": "openid", } assert ( credentials._client_authorization_header("client id", "secret/value") == "Basic Y2xpZW50K2lkOnNlY3JldCUyRnZhbHVl" ) + device_response = { + "device_code": "device-code", + "user_code": "user-code", + "verification_uri": "http://issuer.example/verify", + "expires_in": 60, + } + with pytest.raises(issues.Error, match="verification_uri URL"): + credentials._process_device_authorization_response(device_response) + device_response["verification_uri"] = "https://issuer.example/verify" + device_response["verification_uri_complete"] = "http://issuer.example/verify?code=user-code" + with pytest.raises(issues.Error, match="verification_uri_complete URL"): + credentials._process_device_authorization_response(device_response) + @pytest.mark.parametrize( "invalid_value, message", diff --git a/ydb/aio/oidc.py b/ydb/aio/oidc.py index 0e515aed5..0ba97a4f2 100644 --- a/ydb/aio/oidc.py +++ b/ydb/aio/oidc.py @@ -47,7 +47,7 @@ async def _request_json( async def _discovery(self) -> typing.Dict[str, typing.Any]: if self._discovery_document is None: - url = self._issuer + "/.well-known/openid-configuration" + url = self._issuer.rstrip("/") + "/.well-known/openid-configuration" status, response = await self._request_json(url) self._discovery_document = self._process_discovery_response(status, response) return self._discovery_document diff --git a/ydb/oidc/_common.py b/ydb/oidc/_common.py index a1cc2f2c0..e8c2bfeaf 100644 --- a/ydb/oidc/_common.py +++ b/ydb/oidc/_common.py @@ -5,7 +5,7 @@ import ssl import typing from dataclasses import dataclass -from urllib.parse import quote_plus +from urllib.parse import quote_plus, urlsplit from ydb import issues @@ -39,7 +39,10 @@ def __init__( if request_timeout <= 0: raise ValueError("OAuth 2.0 request timeout must be positive") - self._issuer = issuer.rstrip("/") + if not self._is_https_url(issuer, issuer=True): + raise ValueError("OAuth 2.0 issuer must be an absolute HTTPS URL without userinfo, query, or fragment") + + self._issuer = issuer self._client_id = client_id self._scope = self._scope_parameter(scope) self._audience = audience @@ -48,10 +51,34 @@ def __init__( self._ssl_context = ssl.create_default_context(cafile=os.path.expanduser(ca_file) if ca_file else None) @staticmethod - def _scope_parameter(scope: typing.Union[str, typing.Sequence[str], None]) -> typing.Optional[str]: - if scope is None or isinstance(scope, str): - return scope - return " ".join(scope) + def _scope_parameter(scope: typing.Union[str, typing.Sequence[str], None]) -> str: + if scope is None: + scopes = [] + elif isinstance(scope, str): + scopes = scope.split() + else: + scopes = list(scope) + if "openid" not in scopes: + scopes.append("openid") + return " ".join(scopes) + + @staticmethod + def _is_https_url(value: str, issuer: bool = False) -> bool: + if not isinstance(value, str) or any(ord(character) <= 0x20 or ord(character) == 0x7F for character in value): + return False + try: + parsed = urlsplit(value) + parsed.port + except ValueError: + return False + return ( + parsed.scheme == "https" + and parsed.hostname is not None + and parsed.username is None + and parsed.password is None + and not parsed.fragment + and (not issuer or not parsed.query) + ) @staticmethod def _decode_json(content: bytes, url: str) -> typing.Dict[str, typing.Any]: @@ -94,6 +121,13 @@ def _process_discovery_response( token_endpoint = response.get("token_endpoint") if not isinstance(token_endpoint, str) or not token_endpoint: raise issues.Error("OIDC discovery response does not contain a token_endpoint") + if not self._is_https_url(token_endpoint): + raise issues.Error("OIDC discovery response contains an invalid token_endpoint URL") + device_endpoint = response.get("device_authorization_endpoint") + if device_endpoint is not None and ( + not isinstance(device_endpoint, str) or not self._is_https_url(device_endpoint) + ): + raise issues.Error("OIDC discovery response contains an invalid device_authorization_endpoint URL") return response @staticmethod @@ -119,17 +153,13 @@ def _client_authorization_header(client_id: str, client_secret: str) -> str: return "Basic " + value def _client_credentials_data(self) -> typing.Dict[str, str]: - data = {"grant_type": "client_credentials"} - if self._scope: - data["scope"] = self._scope + data = {"grant_type": "client_credentials", "scope": self._scope} if self._audience: data["audience"] = self._audience return data def _device_authorization_data(self) -> typing.Dict[str, str]: - data = {"client_id": self._client_id} - if self._scope: - data["scope"] = self._scope + data = {"client_id": self._client_id, "scope": self._scope} if self._audience: data["audience"] = self._audience return data @@ -146,9 +176,8 @@ def _refresh_token_data(self, refresh_token: str) -> typing.Dict[str, str]: "grant_type": "refresh_token", "client_id": self._client_id, "refresh_token": refresh_token, + "scope": self._scope, } - if self._scope: - data["scope"] = self._scope return data @staticmethod @@ -168,8 +197,12 @@ def _process_device_authorization_response( raise issues.Error("Device Authorization response does not contain a user_code") if not isinstance(verification_uri, str) or not verification_uri: raise issues.Error("Device Authorization response does not contain a verification_uri") + if not OAuth2CredentialsBase._is_https_url(verification_uri): + raise issues.Error("Device Authorization response contains an invalid verification_uri URL") if verification_uri_complete is not None and not isinstance(verification_uri_complete, str): raise issues.Error("Device Authorization response contains an invalid verification_uri_complete") + if verification_uri_complete is not None and not OAuth2CredentialsBase._is_https_url(verification_uri_complete): + raise issues.Error("Device Authorization response contains an invalid verification_uri_complete URL") if isinstance(expires_in, bool) or not isinstance(expires_in, (int, float)) or expires_in <= 0: raise issues.Error("Device Authorization response contains an invalid expires_in: {!r}".format(expires_in)) if isinstance(interval, bool) or not isinstance(interval, (int, float)) or interval <= 0: diff --git a/ydb/oidc/credentials.py b/ydb/oidc/credentials.py index 0530d5ccf..d85d804b5 100644 --- a/ydb/oidc/credentials.py +++ b/ydb/oidc/credentials.py @@ -62,7 +62,7 @@ def _request_json( def _discovery(self) -> typing.Dict[str, typing.Any]: if self._discovery_document is None: - url = self._issuer + "/.well-known/openid-configuration" + url = self._issuer.rstrip("/") + "/.well-known/openid-configuration" status, response = self._request_json(url) self._discovery_document = self._process_discovery_response(status, response) return self._discovery_document From e9ff4114693d29111bbaf768f46561f22ccbc1e5 Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Mon, 21 Sep 2026 17:04:02 +0300 Subject: [PATCH 5/7] fix: harden OIDC credential edge cases --- tests/aio/test_credentials.py | 32 +++++++++++++++--------- tests/auth/test_credentials.py | 45 ++++++++++++++++++++++------------ ydb/aio/oidc.py | 24 ++++++++++++------ ydb/oidc/_common.py | 12 ++++++--- ydb/oidc/credentials.py | 19 ++++++++++---- 5 files changed, 89 insertions(+), 43 deletions(-) diff --git a/tests/aio/test_credentials.py b/tests/aio/test_credentials.py index 8969646e9..88cf021d2 100644 --- a/tests/aio/test_credentials.py +++ b/tests/aio/test_credentials.py @@ -324,6 +324,7 @@ async def callback(info): assert callback_values[0].user_code == "user-code" assert sleep.await_count == 3 + assert credentials._request_json.await_args_list[2].kwargs["request_timeout"] == 10 @pytest.mark.asyncio @@ -349,10 +350,12 @@ async def test_oauth2_async_http_requests_and_discovery_cache(): with patch("ydb.aio.oidc.aiohttp.ClientSession", return_value=session_context) as client_session: first = await credentials._discovery() second = await credentials._discovery() + await credentials._request_json("https://issuer.example/token", request_timeout=0.5) assert first is second - assert client_session.call_count == 1 - assert session.request.call_args.args[1] == "https://issuer.example/.well-known/openid-configuration" + assert client_session.call_count == 2 + assert client_session.call_args.kwargs["timeout"].total == 0.5 + assert session.request.call_args_list[0].args[1] == "https://issuer.example/.well-known/openid-configuration" assert session.request.call_args.kwargs["ssl"] is not None with patch( @@ -424,19 +427,26 @@ async def test_oauth2_async_device_refresh_and_error_paths(): @pytest.mark.asyncio @pytest.mark.parametrize( - "token_response, expected_message", + "token_response, expected_message, monotonic_values", [ - ({"error": "expired_token"}, "expired"), - ({"error": "access_denied"}, "access_denied"), - (None, "timed out"), + ({"error": "expired_token"}, "expired", [0, 0, 0]), + ({"error": "access_denied"}, "access_denied", [0, 0, 0]), + (None, "timed out", [0, 2]), + (None, "timed out", [0, 0, 2]), ], ) -async def test_oauth2_async_device_expiration(token_response, expected_message): +async def test_oauth2_async_device_expiration(token_response, expected_message, monotonic_values): callback_values = [] + loop = MagicMock() + + def callback(info): + callback_values.append(info) + assert loop.time.call_count == 1 + credentials = ydb.aio.oidc.OAuth2DeviceCredentials( "https://issuer.example", "client-id", - callback_values.append, + callback, device_flow_timeout=1, ) credentials._discovery_document = { @@ -451,11 +461,9 @@ async def test_oauth2_async_device_expiration(token_response, expected_message): "interval": 1, } responses = [(200, device_response)] - loop = MagicMock() - loop.time.side_effect = [0, 2] + loop.time.side_effect = monotonic_values if token_response is not None: responses.append((400, token_response)) - loop.time.side_effect = [0, 0] credentials._request_json = AsyncMock(side_effect=responses) with patch("ydb.aio.oidc.asyncio.get_running_loop", return_value=loop), patch( @@ -465,3 +473,5 @@ async def test_oauth2_async_device_expiration(token_response, expected_message): await credentials._make_token_request() assert callback_values[0].user_code == "user-code" + if token_response is not None: + assert credentials._request_json.await_args_list[1].kwargs["request_timeout"] == 1 diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index 3ec2b740f..dc6ee1934 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -109,7 +109,7 @@ def test_oauth2_client_credentials(): audience="ydb", ) - def request_json(url, data=None, headers=None): + def request_json(url, data=None, headers=None, request_timeout=None): requests.append((url, data, headers)) return next(responses) @@ -179,8 +179,8 @@ def test_oauth2_device_credentials_poll_and_refresh(): callback_values.append, ) - def request_json(url, data=None, headers=None): - requests.append((url, data, headers)) + def request_json(url, data=None, headers=None, request_timeout=None): + requests.append((url, data, headers, request_timeout)) return next(responses) credentials._request_json = request_json @@ -190,6 +190,7 @@ def request_json(url, data=None, headers=None): assert callback_values[0].user_code == "user-code" assert [value.args[0] for value in sleep.call_args_list] == [1, 1, 6] + assert requests[2][3] == 10 assert credentials._make_token_request() == { "access_token": "Bearer refreshed-access-token", "expires_in": 300, @@ -222,7 +223,7 @@ def test_oauth2_client_credentials_validation(kwargs): ydb.oidc.OAuth2ClientCredentials(**arguments) -@pytest.mark.parametrize("token", ["", None]) +@pytest.mark.parametrize("token", ["", " ", "Bearer ", "bearer ", None]) def test_oauth2_token_credentials_validation(token): with pytest.raises(ValueError): bearer_token(token) @@ -326,7 +327,9 @@ def test_oauth2_common_response_processing(): ({"verification_uri": ""}, "verification_uri"), ({"verification_uri_complete": 42}, "verification_uri_complete"), ({"expires_in": 0}, "expires_in"), + ({"expires_in": 0.5}, "expires_in"), ({"interval": 0}, "interval"), + ({"interval": 0.5}, "interval"), ], ) def test_oauth2_device_authorization_response_validation(invalid_value, message): @@ -358,9 +361,11 @@ def test_oauth2_sync_http_requests_and_discovery_cache(): with patch("ydb.oidc.credentials.urllib.request.urlopen", return_value=response_context) as urlopen: first = credentials._discovery() second = credentials._discovery() + credentials._request_json("https://issuer.example/token", request_timeout=0.5) assert first is second - assert urlopen.call_count == 1 + assert urlopen.call_count == 2 + assert urlopen.call_args.kwargs["timeout"] == 0.5 http_error = urllib.error.HTTPError( "https://issuer.example/token", @@ -429,18 +434,25 @@ def test_oauth2_sync_device_error_paths(): @pytest.mark.parametrize( - "token_response, expected_message", + "token_response, expected_message, monotonic_values", [ - ({"error": "expired_token"}, "expired"), - ({"error": "access_denied"}, "access_denied"), - (None, "timed out"), + ({"error": "expired_token"}, "expired", [0, 0, 0]), + ({"error": "access_denied"}, "access_denied", [0, 0, 0]), + (None, "timed out", [0, 2]), + (None, "timed out", [0, 0, 2]), ], ) -def test_oauth2_sync_device_expiration(token_response, expected_message): +def test_oauth2_sync_device_expiration(token_response, expected_message, monotonic_values): + callback_clock_calls = [] + monotonic = MagicMock() + + def callback(info): + callback_clock_calls.append(monotonic.call_count) + credentials = ydb.oidc.OAuth2DeviceCredentials( "https://issuer.example", "client-id", - lambda info: None, + callback, device_flow_timeout=1, ) credentials._discovery_document = { @@ -455,14 +467,15 @@ def test_oauth2_sync_device_expiration(token_response, expected_message): "interval": 1, } responses = [(200, device_response)] - monotonic_values = [0, 2] if token_response is not None: responses.append((400, token_response)) - monotonic_values = [0, 0] credentials._request_json = MagicMock(side_effect=responses) - with patch("ydb.oidc.credentials.time.monotonic", side_effect=monotonic_values), patch( - "ydb.oidc.credentials.time.sleep" - ): + monotonic.side_effect = monotonic_values + with patch("ydb.oidc.credentials.time.monotonic", monotonic), patch("ydb.oidc.credentials.time.sleep"): with pytest.raises(issues.Unauthenticated, match=expected_message): credentials._make_token_request() + + assert callback_clock_calls == [1] + if token_response is not None: + assert credentials._request_json.call_args_list[1].kwargs["request_timeout"] == 1 diff --git a/ydb/aio/oidc.py b/ydb/aio/oidc.py index 0ba97a4f2..5b0e54b57 100644 --- a/ydb/aio/oidc.py +++ b/ydb/aio/oidc.py @@ -29,8 +29,9 @@ async def _request_json( url: str, data: typing.Optional[typing.Mapping[str, str]] = None, headers: typing.Optional[typing.Mapping[str, str]] = None, + request_timeout: typing.Optional[float] = None, ) -> typing.Tuple[int, typing.Dict[str, typing.Any]]: - timeout = aiohttp.ClientTimeout(total=self._request_timeout) + timeout = aiohttp.ClientTimeout(total=self._request_timeout if request_timeout is None else request_timeout) ssl_context = self._ssl_context if url.startswith("https://") else None try: async with aiohttp.ClientSession(timeout=timeout) as session: @@ -151,22 +152,31 @@ async def _make_token_request(self): ) self._raise_for_status(status, response) device_code, info = self._process_device_authorization_response(response) + timeout = info.expires_in + if self._device_flow_timeout is not None: + timeout = min(timeout, self._device_flow_timeout) + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + callback_result = self._device_authorization_callback(info) if inspect.isawaitable(callback_result): await callback_result - timeout = info.expires_in - if self._device_flow_timeout is not None: - timeout = min(timeout, self._device_flow_timeout) - deadline = asyncio.get_running_loop().time() + timeout interval = info.interval - while asyncio.get_running_loop().time() < deadline: - await asyncio.sleep(interval) + while True: + remaining = deadline - loop.time() + if remaining <= 0: + break + await asyncio.sleep(min(interval, remaining)) + remaining = deadline - loop.time() + if remaining <= 0: + break status, response = await self._request_json( token_endpoint, self._device_token_data(device_code), self._client_headers(), + request_timeout=min(self._request_timeout, remaining), ) if 200 <= status < 300: return self._save_token_response(response) diff --git a/ydb/oidc/_common.py b/ydb/oidc/_common.py index e8c2bfeaf..960bc8f9c 100644 --- a/ydb/oidc/_common.py +++ b/ydb/oidc/_common.py @@ -203,17 +203,17 @@ def _process_device_authorization_response( raise issues.Error("Device Authorization response contains an invalid verification_uri_complete") if verification_uri_complete is not None and not OAuth2CredentialsBase._is_https_url(verification_uri_complete): raise issues.Error("Device Authorization response contains an invalid verification_uri_complete URL") - if isinstance(expires_in, bool) or not isinstance(expires_in, (int, float)) or expires_in <= 0: + if isinstance(expires_in, bool) or not isinstance(expires_in, int) or expires_in <= 0: raise issues.Error("Device Authorization response contains an invalid expires_in: {!r}".format(expires_in)) - if isinstance(interval, bool) or not isinstance(interval, (int, float)) or interval <= 0: + if isinstance(interval, bool) or not isinstance(interval, int) or interval <= 0: raise issues.Error("Device Authorization response contains an invalid interval: {!r}".format(interval)) info = DeviceAuthorizationInfo( verification_uri=verification_uri, user_code=user_code, verification_uri_complete=verification_uri_complete, - expires_in=int(expires_in), - interval=int(interval), + expires_in=expires_in, + interval=interval, ) return device_code, info @@ -222,5 +222,9 @@ def bearer_token(token: str) -> str: if not isinstance(token, str) or not token: raise ValueError("OAuth 2.0 access token must not be empty") if token.lower().startswith("bearer "): + if not token[7:].strip(): + raise ValueError("OAuth 2.0 access token must not be empty") return token + if not token.strip(): + raise ValueError("OAuth 2.0 access token must not be empty") return "Bearer " + token diff --git a/ydb/oidc/credentials.py b/ydb/oidc/credentials.py index d85d804b5..3273d3dc0 100644 --- a/ydb/oidc/credentials.py +++ b/ydb/oidc/credentials.py @@ -41,6 +41,7 @@ def _request_json( url: str, data: typing.Optional[typing.Mapping[str, str]] = None, headers: typing.Optional[typing.Mapping[str, str]] = None, + request_timeout: typing.Optional[float] = None, ) -> typing.Tuple[int, typing.Dict[str, typing.Any]]: body = urllib.parse.urlencode(data).encode("utf-8") if data is not None else None request_headers = dict(headers or {}) @@ -52,7 +53,7 @@ def _request_json( with urllib.request.urlopen( request, context=self._ssl_context, - timeout=self._request_timeout, + timeout=self._request_timeout if request_timeout is None else request_timeout, ) as response: return response.status, self._decode_json(response.read(), url) except urllib.error.HTTPError as error: @@ -172,20 +173,28 @@ def _make_token_request(self): ) self._raise_for_status(status, response) device_code, info = self._process_device_authorization_response(response) - self._device_authorization_callback(info) - timeout = info.expires_in if self._device_flow_timeout is not None: timeout = min(timeout, self._device_flow_timeout) deadline = time.monotonic() + timeout + + self._device_authorization_callback(info) + interval = info.interval - while time.monotonic() < deadline: - time.sleep(interval) + while True: + remaining = deadline - time.monotonic() + if remaining <= 0: + break + time.sleep(min(interval, remaining)) + remaining = deadline - time.monotonic() + if remaining <= 0: + break status, response = self._request_json( token_endpoint, self._device_token_data(device_code), self._client_headers(), + request_timeout=min(self._request_timeout, remaining), ) if 200 <= status < 300: return self._save_token_response(response) From 22f15807cf550c8b7b3cbbae269d1f8567fa3f0c Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Mon, 21 Sep 2026 17:37:23 +0300 Subject: [PATCH 6/7] fix: correct OIDC token lifetime handling --- tests/aio/test_credentials.py | 26 ++++++++++++++++++++++++-- tests/auth/test_credentials.py | 20 ++++++++++++++++++++ ydb/aio/credentials.py | 8 ++------ ydb/aio/oidc.py | 2 +- ydb/credentials.py | 12 ++++++------ 5 files changed, 53 insertions(+), 15 deletions(-) diff --git a/tests/aio/test_credentials.py b/tests/aio/test_credentials.py index 88cf021d2..0d884cd2c 100644 --- a/tests/aio/test_credentials.py +++ b/tests/aio/test_credentials.py @@ -284,6 +284,27 @@ async def test_oauth2_client_credentials(): assert credentials._request_json.await_args_list[1].args[1]["scope"] == "openid" +@pytest.mark.asyncio +async def test_oauth2_async_token_lifetime_starts_after_request_and_preserves_short_lived_token(): + credentials = ydb.aio.oidc.OAuth2ClientCredentials( + "https://issuer.example", + "client-id", + "client-secret", + ) + now = [1000] + + async def make_token_request(): + now[0] = 1120 + return {"access_token": "Bearer access-token", "expires_in": 10} + + credentials._make_token_request = make_token_request + + with patch("ydb.aio.credentials.time.time", side_effect=lambda: now[0]): + assert await credentials.get_auth_token() == "Bearer access-token" + + assert credentials._expires_in == 1129 + + @pytest.mark.asyncio async def test_oauth2_device_credentials(): issuer = "https://issuer.example" @@ -351,10 +372,11 @@ async def test_oauth2_async_http_requests_and_discovery_cache(): first = await credentials._discovery() second = await credentials._discovery() await credentials._request_json("https://issuer.example/token", request_timeout=0.5) + await credentials._request_json("HTTPS://issuer.example/token") assert first is second - assert client_session.call_count == 2 - assert client_session.call_args.kwargs["timeout"].total == 0.5 + assert client_session.call_count == 3 + assert client_session.call_args_list[1].kwargs["timeout"].total == 0.5 assert session.request.call_args_list[0].args[1] == "https://issuer.example/.well-known/openid-configuration" assert session.request.call_args.kwargs["ssl"] is not None diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index dc6ee1934..e40434c7e 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -126,6 +126,26 @@ def request_json(url, data=None, headers=None, request_timeout=None): assert requests[1][2]["Authorization"].startswith("Basic ") +def test_oauth2_token_lifetime_starts_after_request_and_preserves_short_lived_token(): + credentials = ydb.oidc.OAuth2ClientCredentials( + "https://issuer.example", + "client-id", + "client-secret", + ) + now = [1000] + + def make_token_request(): + now[0] = 1120 + return {"access_token": "Bearer access-token", "expires_in": 10} + + credentials._make_token_request = make_token_request + + with patch("ydb.credentials.time.time", side_effect=lambda: now[0]): + assert credentials.get_auth_token() == "Bearer access-token" + + assert credentials._expires_in == 1129 + + def test_oauth2_device_credentials_poll_and_refresh(): issuer = "https://issuer.example" callback_values = [] diff --git a/ydb/aio/credentials.py b/ydb/aio/credentials.py index 48784e8ee..bce161765 100644 --- a/ydb/aio/credentials.py +++ b/ydb/aio/credentials.py @@ -47,15 +47,11 @@ async def get_auth_token(self) -> str: # type: ignore[override] return "" async def _refresh_token(self, should_raise=False): - current_time = time.time() - try: - self.logger.debug( - "Refreshing token async, current_time: %s, expires_in: %s", current_time, self._expires_in - ) + self.logger.debug("Refreshing token async, expires_in: %s", self._expires_in) token_response = await self._make_token_request() - self._update_token_info(token_response, current_time) + self._update_token_info(token_response, time.time()) self.logger.info("Token refreshed successfully async, expires_in: %s", self._expires_in) self.last_error = None diff --git a/ydb/aio/oidc.py b/ydb/aio/oidc.py index 5b0e54b57..ed3f36de7 100644 --- a/ydb/aio/oidc.py +++ b/ydb/aio/oidc.py @@ -32,7 +32,7 @@ async def _request_json( request_timeout: typing.Optional[float] = None, ) -> typing.Tuple[int, typing.Dict[str, typing.Any]]: timeout = aiohttp.ClientTimeout(total=self._request_timeout if request_timeout is None else request_timeout) - ssl_context = self._ssl_context if url.startswith("https://") else None + ssl_context = self._ssl_context if self._is_https_url(url) else None try: async with aiohttp.ClientSession(timeout=timeout) as session: async with session.request( diff --git a/ydb/credentials.py b/ydb/credentials.py index 28ea8b52f..2be2790ac 100644 --- a/ydb/credentials.py +++ b/ydb/credentials.py @@ -100,18 +100,18 @@ def _should_refresh(self): return time.time() >= self._refresh_in def _update_token_info(self, token_response, current_time): - self._refresh_in = current_time + min(self._hour / 2, token_response["expires_in"] / 10) - self._expires_in = current_time + token_response["expires_in"] - self._time_shift_protection_seconds + expires_in = token_response["expires_in"] + self._refresh_in = current_time + min(self._hour / 2, expires_in / 10) + safety_margin = min(self._time_shift_protection_seconds, expires_in / 10) + self._expires_in = current_time + expires_in - safety_margin self._cached_token = token_response["access_token"] def _refresh_token(self, should_raise=False): - current_time = time.time() - try: - self.logger.debug("Refreshing token, current_time: %s, expires_in: %s", current_time, self._expires_in) + self.logger.debug("Refreshing token, expires_in: %s", self._expires_in) token_response = self._make_token_request() - self._update_token_info(token_response, current_time) + self._update_token_info(token_response, time.time()) self.logger.info("Token refreshed successfully, expires_in: %s", self._expires_in) self.last_error = None From 13d0bcd1bf8cddbdbf5ee42d7658a8e00483ab7f Mon Sep 17 00:00:00 2001 From: Oleg Ovcharuk Date: Mon, 21 Sep 2026 17:51:51 +0300 Subject: [PATCH 7/7] fix: reject redirects from OAuth endpoints --- tests/aio/test_credentials.py | 1 + tests/auth/test_credentials.py | 38 +++++++++++++++++++--------------- ydb/aio/oidc.py | 1 + ydb/oidc/credentials.py | 12 +++++++++-- 4 files changed, 33 insertions(+), 19 deletions(-) diff --git a/tests/aio/test_credentials.py b/tests/aio/test_credentials.py index 0d884cd2c..866fbe881 100644 --- a/tests/aio/test_credentials.py +++ b/tests/aio/test_credentials.py @@ -379,6 +379,7 @@ async def test_oauth2_async_http_requests_and_discovery_cache(): assert client_session.call_args_list[1].kwargs["timeout"].total == 0.5 assert session.request.call_args_list[0].args[1] == "https://issuer.example/.well-known/openid-configuration" assert session.request.call_args.kwargs["ssl"] is not None + assert session.request.call_args.kwargs["allow_redirects"] is False with patch( "ydb.aio.oidc.aiohttp.ClientSession", diff --git a/tests/auth/test_credentials.py b/tests/auth/test_credentials.py index e40434c7e..8b9fcde26 100644 --- a/tests/auth/test_credentials.py +++ b/tests/auth/test_credentials.py @@ -10,6 +10,7 @@ import ydb.iam import ydb.oidc +import ydb.oidc.credentials from ydb import issues from ydb.oidc._common import OAuth2CredentialsBase from ydb.oidc._common import bearer_token @@ -378,14 +379,20 @@ def test_oauth2_sync_http_requests_and_discovery_cache(): response_context = MagicMock() response_context.__enter__.return_value = response - with patch("ydb.oidc.credentials.urllib.request.urlopen", return_value=response_context) as urlopen: - first = credentials._discovery() - second = credentials._discovery() - credentials._request_json("https://issuer.example/token", request_timeout=0.5) + credentials._opener.open = MagicMock(return_value=response_context) + first = credentials._discovery() + second = credentials._discovery() + credentials._request_json("https://issuer.example/token", request_timeout=0.5) + + redirect_handler = ydb.oidc.credentials._NoRedirectHandler() + assert ( + redirect_handler.redirect_request(None, None, 307, "Temporary Redirect", {}, "http://other.example/token") + is None + ) assert first is second - assert urlopen.call_count == 2 - assert urlopen.call_args.kwargs["timeout"] == 0.5 + assert credentials._opener.open.call_count == 2 + assert credentials._opener.open.call_args.kwargs["timeout"] == 0.5 http_error = urllib.error.HTTPError( "https://issuer.example/token", @@ -394,18 +401,15 @@ def test_oauth2_sync_http_requests_and_discovery_cache(): {}, io.BytesIO(b'{"error":"invalid_request"}'), ) - with patch("ydb.oidc.credentials.urllib.request.urlopen", side_effect=http_error): - assert credentials._request_json("https://issuer.example/token", {"key": "value"}) == ( - 400, - {"error": "invalid_request"}, - ) + credentials._opener.open.side_effect = http_error + assert credentials._request_json("https://issuer.example/token", {"key": "value"}) == ( + 400, + {"error": "invalid_request"}, + ) - with patch( - "ydb.oidc.credentials.urllib.request.urlopen", - side_effect=urllib.error.URLError("unavailable"), - ): - with pytest.raises(issues.Unavailable): - credentials._request_json("https://issuer.example/token") + credentials._opener.open.side_effect = urllib.error.URLError("unavailable") + with pytest.raises(issues.Unavailable): + credentials._request_json("https://issuer.example/token") def test_oauth2_sync_device_error_paths(): diff --git a/ydb/aio/oidc.py b/ydb/aio/oidc.py index ed3f36de7..9ea6821f3 100644 --- a/ydb/aio/oidc.py +++ b/ydb/aio/oidc.py @@ -41,6 +41,7 @@ async def _request_json( data=data, headers=headers, ssl=ssl_context, + allow_redirects=False, ) as response: return response.status, self._decode_json(await response.read(), url) except (aiohttp.ClientError, asyncio.TimeoutError, OSError) as error: diff --git a/ydb/oidc/credentials.py b/ydb/oidc/credentials.py index 3273d3dc0..aa652a64c 100644 --- a/ydb/oidc/credentials.py +++ b/ydb/oidc/credentials.py @@ -11,6 +11,11 @@ from ._common import DeviceAuthorizationInfo, OAuth2CredentialsBase, bearer_token +class _NoRedirectHandler(urllib.request.HTTPRedirectHandler): + def redirect_request(self, request, file_pointer, code, message, headers, new_url): + return None + + class OAuth2TokenCredentials(credentials.Credentials): """Credentials for an OAuth 2.0 access token obtained outside the SDK.""" @@ -35,6 +40,10 @@ def __init__( ): credentials.AbstractExpiringTokenCredentials.__init__(self, tracer) OAuth2CredentialsBase.__init__(self, issuer, client_id, scope, audience, ca_file, request_timeout) + self._opener = urllib.request.build_opener( + urllib.request.HTTPSHandler(context=self._ssl_context), + _NoRedirectHandler(), + ) def _request_json( self, @@ -50,9 +59,8 @@ def _request_json( request = urllib.request.Request(url, data=body, headers=request_headers) try: - with urllib.request.urlopen( + with self._opener.open( request, - context=self._ssl_context, timeout=self._request_timeout if request_timeout is None else request_timeout, ) as response: return response.status, self._decode_json(response.read(), url)