diff --git a/src/openhound_github/auth.py b/src/openhound_github/auth.py index c899def..bbc8a10 100644 --- a/src/openhound_github/auth.py +++ b/src/openhound_github/auth.py @@ -2,6 +2,8 @@ from datetime import datetime, timedelta, timezone from threading import Lock from typing import Iterator +from urllib.parse import urlparse +from weakref import WeakKeyDictionary import requests from dlt.common.configuration import configspec @@ -17,6 +19,19 @@ logger = logging.getLogger(__name__) +def _normalized_http_origin(url: str) -> tuple[str, str, int | None]: + parsed = urlparse(url) + scheme = parsed.scheme.lower() + if scheme != "https" or not parsed.hostname: + raise ValueError("GitHub API URI must be an absolute HTTPS URL") + + port = parsed.port + if scheme == "https" and port == 443: + port = None + + return scheme, parsed.hostname.lower(), port + + class AccountConfig(BaseModel): id: int login: str | None = None @@ -63,7 +78,8 @@ def __init__( private_key_path: str, api_uri: str = "https://api.github.com/", ): - self.api_uri = api_uri + _normalized_http_origin(api_uri) + self.api_uri = f"{api_uri.rstrip('/')}/" self.jwt_issuer = jwt_issuer self.private_key_path = private_key_path self.client = RESTClient( @@ -158,11 +174,28 @@ def __init__( self, installation: GithubInstallation, refresh_margin_seconds: int = 300, + api_uri: str | None = None, ): self.installation = installation self.refresh_margin_seconds = refresh_margin_seconds + installation_api_uri = getattr( + installation, "api_uri", "https://api.github.com/" + ) + selected_api_uri = api_uri or installation_api_uri + installation_api_origin = _normalized_http_origin(installation_api_uri) + selected_api_origin = _normalized_http_origin(selected_api_uri) + if selected_api_origin != installation_api_origin: + raise ValueError( + "GitHub App auth API URI origin must match installation API URI origin" + ) + + self.api_uri = selected_api_uri + self._api_origin = selected_api_origin self.access_token: str | None = None self.expires_at: datetime | None = None + self._response_refreshed_requests: WeakKeyDictionary[ + requests.PreparedRequest, str + ] = WeakKeyDictionary() self._token_lock = Lock() def _should_refresh(self) -> bool: @@ -172,6 +205,14 @@ def _should_refresh(self) -> bool: refresh_at = self.expires_at - timedelta(seconds=self.refresh_margin_seconds) return datetime.now(timezone.utc) >= refresh_at + def _refresh_token(self) -> None: + logger.info( + f"Refreshing access token for {self.installation.installation_id}" + ) + get_token = self.installation.token + self.access_token = get_token.token + self.expires_at = get_token.expires_at + def token(self, force_refresh: bool = False) -> str | None: if ( not force_refresh @@ -182,15 +223,62 @@ def token(self, force_refresh: bool = False) -> str | None: with self._token_lock: if (force_refresh or self._should_refresh()) or self.access_token is None: - logger.info( - f"Refreshing access token for {self.installation.installation_id}" - ) - get_token = self.installation.token - self.access_token = get_token.token - self.expires_at = get_token.expires_at + self._refresh_token() return self.access_token + def refresh_request(self, request: requests.PreparedRequest) -> bool: + """Repair a rejected same-origin request without stampeding token issuance.""" + try: + request_origin = _normalized_http_origin(request.url or "") + except ValueError: + return False + + if request_origin != self._api_origin: + return False + + request_authorization = request.headers.get("Authorization") + if ( + not request_authorization + or not request_authorization.startswith("Bearer ") + or not request_authorization.removeprefix("Bearer ").strip() + ): + return False + + with self._token_lock: + current_authorization = ( + f"Bearer {self.access_token}" if self.access_token is not None else None + ) + should_refresh = self._should_refresh() + repaired_authorization = self._response_refreshed_requests.get(request) + if ( + not should_refresh + and request_authorization == current_authorization + and request_authorization == repaired_authorization + ): + return False + + if ( + self.access_token is None + or should_refresh + or request_authorization == current_authorization + ): + try: + self._refresh_token() + except Exception: + logger.warning( + "Failed to refresh GitHub App installation token for " + "installation %s during request retry", + self.installation.installation_id, + ) + return False + + replacement_authorization = f"Bearer {self.access_token}" + request.headers["Authorization"] = replacement_authorization + self._response_refreshed_requests[request] = replacement_authorization + + return True + def __call__(self, request: requests.PreparedRequest) -> requests.PreparedRequest: request.headers["Authorization"] = f"Bearer {self.token()}" return request diff --git a/src/openhound_github/helpers.py b/src/openhound_github/helpers.py index 36bd036..55c5295 100644 --- a/src/openhound_github/helpers.py +++ b/src/openhound_github/helpers.py @@ -10,6 +10,8 @@ ) from requests import Request +from openhound_github.auth import GitHubAppInstallationAuth + logger = logging.getLogger(__name__) @@ -160,6 +162,22 @@ def retry_policy( headers = response.headers now = int(time.time()) + message = _response_message(response).lower() + + # DLT retries the same prepared request after long Retry-After sleeps. + if ( + response.status_code == 401 + and "bad credentials" in message + and isinstance(auth, GitHubAppInstallationAuth) + and response.request is not None + ): + if not auth.refresh_request(response.request): + return False + logger.warning( + "GitHub App installation token rejected, retrying request with refreshed token" + ) + return True + if ( response.status_code == 200 and headers.get("x-ratelimit-resource") == "graphql" @@ -178,7 +196,6 @@ def retry_policy( return True return False - message = _response_message(response).lower() if response.status_code not in (403, 429): return False diff --git a/src/openhound_github/source.py b/src/openhound_github/source.py index bd08782..4ec8466 100644 --- a/src/openhound_github/source.py +++ b/src/openhound_github/source.py @@ -172,6 +172,7 @@ def token_client(token: str) -> RESTClient: github_app_session = GithubApp( jwt_issuer=jwt_issuer, private_key_path=credentials.key_path, + api_uri=host, ) for installation in github_app_session.installations: if installation.target_type == "Organization": @@ -179,12 +180,16 @@ def token_client(token: str) -> RESTClient: installation_id=installation.id, jwt_issuer=jwt_issuer, private_key_path=credentials.key_path, + api_uri=host, ) ctx.organizations.append( OrgContext( org_name=installation.account.login, client=client( - GitHubAppInstallationAuth(installation=org_installation) + GitHubAppInstallationAuth( + installation=org_installation, + api_uri=host, + ) ), enterprise_name=credentials.enterprise_name, github_deployment_id=github_deployment_id, @@ -196,9 +201,13 @@ def token_client(token: str) -> RESTClient: installation_id=installation.id, jwt_issuer=jwt_issuer, private_key_path=credentials.key_path, + api_uri=host, ) ctx.client = client( - GitHubAppInstallationAuth(installation=es_installation) + GitHubAppInstallationAuth( + installation=es_installation, + api_uri=host, + ) ) return (*enterprise_resources(ctx), *organization_resources(ctx)) @@ -213,11 +222,17 @@ def token_client(token: str) -> RESTClient: installation_id=credentials.install_id, jwt_issuer=credentials.client_id, private_key_path=credentials.key_path, + api_uri=host, ) ctx.organizations.append( OrgContext( org_name=credentials.org_name, - client=client(GitHubAppInstallationAuth(installation=org_installation)), + client=client( + GitHubAppInstallationAuth( + installation=org_installation, + api_uri=host, + ) + ), github_deployment_id=github_deployment_id, github_web_origin=github_web_origin, ) diff --git a/tests/test_app_auth.py b/tests/test_app_auth.py index ddfb35c..236f274 100644 --- a/tests/test_app_auth.py +++ b/tests/test_app_auth.py @@ -9,6 +9,7 @@ from openhound_github import auth from openhound_github.auth import ( AccountConfig, + GitHubAppInstallationAuth, GithubSession, InstallationResponse, resolve_github_app_jwt_issuer, @@ -97,6 +98,28 @@ def test_legacy_installation_response_does_not_require_client_id() -> None: assert installation.app_id == 123456 +def test_github_app_installation_auth_rejects_mismatched_api_origins() -> None: + installation = SimpleNamespace( + installation_id="12345", + api_uri="https://ghe.example/api/v3/", + ) + + with pytest.raises(ValueError, match="must match installation API URI origin"): + GitHubAppInstallationAuth( + installation=installation, + api_uri="https://api.github.com/", + ) + + +def test_github_session_rejects_plaintext_api_uri() -> None: + with pytest.raises(ValueError, match="absolute HTTPS URL"): + GithubSession( + jwt_issuer="123456", + private_key_path="/tmp/github-app.pem", + api_uri="http://ghe.example/api/v3/", + ) + + def test_enterprise_source_reuses_selected_issuer_for_installation_tokens( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -104,7 +127,9 @@ def test_enterprise_source_reuses_selected_issuer_for_installation_tokens( captured_issuers: list[str] = [] class FakeGithubApp: - def __init__(self, jwt_issuer: str, private_key_path: str) -> None: + def __init__( + self, jwt_issuer: str, private_key_path: str, api_uri: str + ) -> None: captured_issuers.append(jwt_issuer) self.installations = ( SimpleNamespace( @@ -125,6 +150,7 @@ def __init__( installation_id: int, jwt_issuer: str, private_key_path: str, + api_uri: str, ) -> None: captured_issuers.append(jwt_issuer) diff --git a/tests/test_github_app_retry.py b/tests/test_github_app_retry.py new file mode 100644 index 0000000..3620549 --- /dev/null +++ b/tests/test_github_app_retry.py @@ -0,0 +1,216 @@ +import logging +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta, timezone +from threading import Event, Lock + +import requests +from dlt.sources.helpers.rest_client.auth import BearerTokenAuth + +from openhound_github.auth import GitHubAppInstallationAuth, TokenResponse +from openhound_github.helpers import github_retry_policy + + +class FakeInstallation: + installation_id = "12345" + api_uri = "https://api.github.com/" + + def __init__(self, *tokens: str) -> None: + self._tokens = iter(tokens) + self.token_calls = 0 + + @property + def token(self) -> TokenResponse: + self.token_calls += 1 + return TokenResponse( + token=next(self._tokens), + expires_at=datetime.now(timezone.utc) + timedelta(hours=1), + ) + + +class BlockingFakeInstallation(FakeInstallation): + def __init__(self, *tokens: str) -> None: + super().__init__(*tokens) + self.token_started = Event() + self.release_token = Event() + + @property + def token(self) -> TokenResponse: + self.token_started.set() + if not self.release_token.wait(timeout=1): + raise TimeoutError("timed out waiting to release fake token refresh") + return super().token + + +class FailingFakeInstallation(FakeInstallation): + @property + def token(self) -> TokenResponse: + self.token_calls += 1 + raise RuntimeError("sensitive-token-data") + + +class TrackingLock: + def __init__(self) -> None: + self._lock = Lock() + self.waiting = Event() + + def __enter__(self): + if not self._lock.acquire(blocking=False): + self.waiting.set() + self._lock.acquire() + return self + + def __exit__(self, _exc_type, _exc_value, _traceback) -> None: + self._lock.release() + + +def prepared_request( + token: str | None, + url: str = "https://api.github.com/repos/example/repo", +) -> requests.PreparedRequest: + headers = {"Authorization": f"Bearer {token}"} if token is not None else {} + return requests.Request( + "GET", + url, + headers=headers, + ).prepare() + + +def bad_credentials_response( + request: requests.PreparedRequest, +) -> requests.Response: + response = requests.Response() + response.status_code = 401 + response._content = b'{"message":"Bad credentials"}' + response.request = request + return response + + +def test_refresh_request_refreshes_rejected_current_token() -> None: + installation = FakeInstallation("new-token") + auth = GitHubAppInstallationAuth(installation=installation) + auth.access_token = "old-token" + auth.expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + request = prepared_request("old-token") + + auth.refresh_request(request) + + assert request.headers["Authorization"] == "Bearer new-token" + assert auth.access_token == "new-token" + assert installation.token_calls == 1 + + +def test_refresh_request_reuses_token_refreshed_by_another_request() -> None: + installation = BlockingFakeInstallation("new-token") + auth = GitHubAppInstallationAuth(installation=installation) + auth.access_token = "old-token" + auth.expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + auth._token_lock = TrackingLock() + requests_to_refresh = [prepared_request("old-token"), prepared_request("old-token")] + + with ThreadPoolExecutor(max_workers=2) as executor: + first_refresh = executor.submit(auth.refresh_request, requests_to_refresh[0]) + assert installation.token_started.wait(timeout=1) + + second_refresh = executor.submit(auth.refresh_request, requests_to_refresh[1]) + assert auth._token_lock.waiting.wait(timeout=1) + + installation.release_token.set() + first_refresh.result(timeout=1) + second_refresh.result(timeout=1) + + assert all( + request.headers["Authorization"] == "Bearer new-token" + for request in requests_to_refresh + ) + assert installation.token_calls == 1 + + +def test_refresh_request_does_not_restore_authorization_on_cross_origin_request() -> None: + installation = FakeInstallation("new-token") + auth = GitHubAppInstallationAuth(installation=installation) + auth.access_token = "old-token" + auth.expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + request = prepared_request(None, url="https://attacker.example/redirected") + + repaired = auth.refresh_request(request) + + assert repaired is False + assert "Authorization" not in request.headers + assert installation.token_calls == 0 + + +def test_retry_policy_repairs_bad_credentials_for_github_app_auth() -> None: + installation = FakeInstallation("new-token") + auth = GitHubAppInstallationAuth(installation=installation) + auth.access_token = "old-token" + auth.expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + request = prepared_request("old-token") + + should_retry = github_retry_policy(auth)(bad_credentials_response(request), None) + + assert should_retry is True + assert request.headers["Authorization"] == "Bearer new-token" + assert installation.token_calls == 1 + + +def test_retry_policy_does_not_retry_replacement_token_bad_credentials() -> None: + installation = FakeInstallation("new-token", "unused-token") + auth = GitHubAppInstallationAuth(installation=installation) + auth.access_token = "old-token" + auth.expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + request = prepared_request("old-token") + retry_policy = github_retry_policy(auth) + + assert retry_policy(bad_credentials_response(request), None) is True + assert request.headers["Authorization"] == "Bearer new-token" + + assert retry_policy(bad_credentials_response(request), None) is False + assert request.headers["Authorization"] == "Bearer new-token" + assert installation.token_calls == 1 + + +def test_retry_policy_allows_fresh_request_to_refresh_replacement_token() -> None: + installation = FakeInstallation("new-token", "newer-token") + auth = GitHubAppInstallationAuth(installation=installation) + auth.access_token = "old-token" + auth.expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + retry_policy = github_retry_policy(auth) + recovery_request = prepared_request("old-token") + + assert retry_policy(bad_credentials_response(recovery_request), None) is True + assert recovery_request.headers["Authorization"] == "Bearer new-token" + + independent_request = prepared_request("new-token") + + assert retry_policy(bad_credentials_response(independent_request), None) is True + assert independent_request.headers["Authorization"] == "Bearer newer-token" + assert installation.token_calls == 2 + + +def test_retry_policy_does_not_retry_when_token_refresh_fails(caplog) -> None: + installation = FailingFakeInstallation() + auth = GitHubAppInstallationAuth(installation=installation) + auth.access_token = "old-token" + auth.expires_at = datetime.now(timezone.utc) + timedelta(hours=1) + request = prepared_request("old-token") + + with caplog.at_level(logging.WARNING, logger="openhound_github.auth"): + should_retry = github_retry_policy(auth)(bad_credentials_response(request), None) + + assert should_retry is False + assert request.headers["Authorization"] == "Bearer old-token" + assert installation.token_calls == 1 + assert "Failed to refresh GitHub App installation token" in caplog.text + assert "sensitive-token-data" not in caplog.text + + +def test_retry_policy_does_not_repair_bad_credentials_for_bearer_token_auth() -> None: + request = prepared_request("static-token") + + should_retry = github_retry_policy(BearerTokenAuth(token="static-token"))( + bad_credentials_response(request), + None, + ) + + assert should_retry is False + assert request.headers["Authorization"] == "Bearer static-token"