Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 35af616

Browse filesBrowse files
nbayatinbayati
authored
feat(auth): implement regional access boundary support for standalone JWT and async service accounts (#17025)
This PR implements the following changes: - Add RAB support to async service account and jwt credential types, by providing async manager and fetching methods. - Update unit tests to accept both mtls and standard lookup endpoint urls. - Refactor before_request to use a _after_refresh hook so we don't have to override the method. - Add RAb support for self signed jwt flow through jwt.py - some small enhancements for test coverage and backward compatibility
1 parent 4005e66 commit 35af616
Copy full SHA for 35af616

30 files changed

+1,769-174Lines changed: 1769 additions & 174 deletions
Expand file treeCollapse file tree
Open diff view settings
Collapse file

‎packages/google-auth/google/auth/_credentials_async.py‎

Copy file name to clipboardExpand all lines: packages/google-auth/google/auth/_credentials_async.py
+92Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import abc
1919
import inspect
2020

21+
from google.auth import _regional_access_boundary_utils
2122
from google.auth import credentials
2223

2324

@@ -64,8 +65,28 @@ async def before_request(self, request, method, url, headers):
6465
await self.refresh(request)
6566
else:
6667
self.refresh(request)
68+
69+
if inspect.iscoroutinefunction(self._after_refresh):
70+
await self._after_refresh(request, method, url, headers)
71+
else:
72+
self._after_refresh(request, method, url, headers)
73+
6774
self.apply(headers)
6875

76+
def _after_refresh(self, request, method, url, headers):
77+
"""Hook for subclasses to perform actions after refresh but before
78+
applying credentials to headers.
79+
80+
Args:
81+
request (google.auth.transport.Request): The object used to make
82+
HTTP requests.
83+
method (str): The request's HTTP method or the RPC method being
84+
invoked.
85+
url (str): The request's URI or the RPC service's URI.
86+
headers (Mapping[str, str]): The request's headers.
87+
"""
88+
pass
89+
6990

7091
class CredentialsWithQuotaProject(credentials.CredentialsWithQuotaProject):
7192
"""Abstract base for credentials supporting ``with_quota_project`` factory"""
@@ -169,3 +190,74 @@ def with_scopes_if_required(credentials, scopes):
169190

170191
class Signing(credentials.Signing, metaclass=abc.ABCMeta):
171192
"""Interface for credentials that can cryptographically sign messages."""
193+
194+
195+
class CredentialsWithRegionalAccessBoundary(
196+
Credentials, credentials.CredentialsWithRegionalAccessBoundary
197+
):
198+
"""Async base for credentials supporting regional access boundary configuration."""
199+
200+
def __init__(self):
201+
super().__init__()
202+
self._rab_manager.refresh_manager = (
203+
_regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
204+
)
205+
206+
def __setstate__(self, state):
207+
super().__setstate__(state)
208+
self._rab_manager.refresh_manager = (
209+
_regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
210+
)
211+
212+
async def _after_refresh(self, request, method, url, headers):
213+
"""Triggers the Regional Access Boundary lookup asynchronously if necessary."""
214+
await self._maybe_start_regional_access_boundary_refresh_async(request, url)
215+
216+
async def _maybe_start_regional_access_boundary_refresh_async(self, request, url):
217+
"""Starts a background refresh or performs a blocking refresh asynchronously.
218+
219+
Args:
220+
request (google.auth.aio.transport.Request): The object used to make
221+
HTTP requests.
222+
url (str): The URL of the request.
223+
"""
224+
# Do not perform a lookup if the request is for a regional endpoint.
225+
if self._is_regional_endpoint(url):
226+
return
227+
228+
# A refresh is only needed if the feature is enabled.
229+
if not self._is_regional_access_boundary_lookup_required():
230+
return
231+
232+
# Trigger background or blocking refresh if needed.
233+
await self._rab_manager.maybe_start_refresh_async(self, request)
234+
235+
async def _lookup_regional_access_boundary(self, request, fail_fast=False):
236+
"""Calls the Regional Access Boundary lookup API asynchronously.
237+
238+
Args:
239+
request (google.auth.aio.transport.Request): The object used to make
240+
HTTP requests.
241+
fail_fast (bool): Whether the lookup should fail fast (short timeout, no retries).
242+
243+
Returns:
244+
Optional[Dict[str, str]]: The Regional Access Boundary information
245+
returned by the lookup API, or None if the lookup failed.
246+
"""
247+
url_builder = self._build_regional_access_boundary_lookup_url
248+
if inspect.iscoroutinefunction(url_builder):
249+
url = await url_builder(request=request)
250+
else:
251+
url = url_builder(request=request)
252+
253+
if not url:
254+
return None
255+
256+
headers = {}
257+
self._apply(headers)
258+
259+
from google.oauth2 import _client_async
260+
261+
return await _client_async._lookup_regional_access_boundary(
262+
request, url, headers=headers, fail_fast=fail_fast
263+
)
Collapse file

‎packages/google-auth/google/auth/_helpers.py‎

Copy file name to clipboardExpand all lines: packages/google-auth/google/auth/_helpers.py
+2Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@
2828
from google.auth import exceptions
2929

3030

31+
DEFAULT_UNIVERSE_DOMAIN = "googleapis.com"
32+
3133
# _BASE_LOGGER_NAME is the base logger for all google-based loggers.
3234
_BASE_LOGGER_NAME = "google"
3335

Collapse file

‎packages/google-auth/google/auth/_jwt_async.py‎

Copy file name to clipboardExpand all lines: packages/google-auth/google/auth/_jwt_async.py
+17-1Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,8 @@
4444
"""
4545

4646
from google.auth import _credentials_async
47+
from google.auth import _helpers
48+
from google.auth import _regional_access_boundary_utils
4749
from google.auth import jwt
4850

4951

@@ -91,7 +93,9 @@ def decode(token, certs=None, verify=True, audience=None):
9193

9294

9395
class Credentials(
94-
jwt.Credentials, _credentials_async.Signing, _credentials_async.Credentials
96+
jwt.Credentials,
97+
_credentials_async.Signing,
98+
_credentials_async.CredentialsWithRegionalAccessBoundary,
9599
):
96100
"""Credentials that use a JWT as the bearer token.
97101
@@ -142,6 +146,14 @@ class Credentials(
142146
new_credentials = credentials.with_claims(audience=new_audience)
143147
"""
144148

149+
def __setstate__(self, state):
150+
"""Restores the credential state and ensures the async refresh manager is attached."""
151+
super().__setstate__(state)
152+
153+
self._rab_manager.refresh_manager = (
154+
_regional_access_boundary_utils._AsyncRegionalAccessBoundaryRefreshManager()
155+
)
156+
145157

146158
class OnDemandCredentials(
147159
jwt.OnDemandCredentials, _credentials_async.Signing, _credentials_async.Credentials
@@ -162,3 +174,7 @@ class OnDemandCredentials(
162174
163175
.. _grpc: http://www.grpc.io/
164176
"""
177+
178+
@_helpers.copy_docstring(jwt.OnDemandCredentials)
179+
async def before_request(self, request, method, url, headers):
180+
super(OnDemandCredentials, self).before_request(request, method, url, headers)

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.