Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions google/auth/credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@ def __init__(self):

self._use_non_blocking_refresh = False
self._refresh_worker = RefreshThreadManager()
self._custom_headers = {}
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not used?


@property
def expired(self):
Expand Down Expand Up @@ -185,6 +186,7 @@ def apply(self, headers, token=None):
self._apply(headers, token)
if self.quota_project_id:
headers["x-goog-user-project"] = self.quota_project_id
headers.update(self._custom_headers)

def _blocking_refresh(self, request):
if not self.valid:
Expand Down Expand Up @@ -233,6 +235,38 @@ def before_request(self, request, method, url, headers):
def with_non_blocking_refresh(self):
self._use_non_blocking_refresh = True

def with_headers(self, headers):
"""Returns a copy of these credentials with additional custom headers.

Args:
headers (Mapping[str, str]): The custom headers to add.

Returns:
google.auth.credentials.Credentials: A new credentials instance.

Raises:
ValueError: If a protected header is included in the input headers.
"""
import copy

PROTECTED_HEADERS = {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a way to capture this in a more central way? I worry that we are missing some now or may miss some in the future if new ones get added

for example

_METADATA_HEADERS = {_METADATA_FLAVOR_HEADER: _METADATA_FLAVOR_VALUE}

"authorization",
"x-goog-user-project",
"x-goog-api-client",
"x-allowed-locations",
}

for key in headers:
if key.lower() in PROTECTED_HEADERS:
raise ValueError(
f"Header '{key}' is protected and cannot be set with with_headers. "
"These headers are managed by the library."
)

new_creds = copy.deepcopy(self)
new_creds._custom_headers.update(headers)
return new_creds


class CredentialsWithQuotaProject(Credentials):
"""Abstract base for credentials supporting ``with_quota_project`` factory"""
Expand Down
42 changes: 42 additions & 0 deletions tests/test_credentials.py
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,48 @@ def test_with_non_blocking_refresh():
assert c._use_non_blocking_refresh


def test_with_headers():
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: this test is testing multiple scenarios. could you break it up so that we are testing ~1 thing per test?

https://engdoc.corp.google.com/eng/doc/tott/episodes/726.md?cl=head

credentials = CredentialsImpl()
request = mock.Mock()

# 1. Add a new custom header
creds_with_header = credentials.with_headers({"X-Custom-Header": "value1"})
headers = {}
creds_with_header.before_request(request, "http://example.com", "GET", headers)
assert headers["X-Custom-Header"] == "value1"
assert "authorization" in headers # Ensure base apply logic ran
assert creds_with_header is not credentials
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure what you're trying to test here but as I understand python, this will never be true.

i.e. if you did:

credentials_a = CredentialsImpl()
credentials_b = CredentialsImpl()
assert credentials_a is not credentials_b # success 

assert not hasattr(credentials, "_custom_headers") or not credentials._custom_headers

# 2. Update an existing custom header
creds_updated = creds_with_header.with_headers({"X-Custom-Header": "value2"})
headers = {}
creds_updated.before_request(request, "http://example.com", "GET", headers)
assert headers["X-Custom-Header"] == "value2"

# 3. Chaining with_headers calls
creds_chained = credentials.with_headers({"X-Header-1": "v1"}).with_headers(
{"X-Header-2": "v2"}
)
headers = {}
creds_chained.before_request(request, "http://example.com", "GET", headers)
assert headers["X-Header-1"] == "v1"
assert headers["X-Header-2"] == "v2"

# 4. Ensure protected headers cannot be set
with pytest.raises(ValueError):
credentials.with_headers({"Authorization": "Bearer token"})
with pytest.raises(ValueError):
credentials.with_headers({"X-Goog-User-Project": "test"})
with pytest.raises(ValueError):
credentials.with_headers({"authorization": "Bearer token"}) # Case-insensitive

# 5. Check original credentials are not modified
headers = {}
credentials.before_request(request, "http://example.com", "GET", headers)
assert "X-Custom-Header" not in headers


def test_expired_and_valid():
credentials = CredentialsImpl()
credentials.token = "token"
Expand Down