Skip to content

Commit 16c728d

Browse files
TimurSadykovgcf-owl-bot[bot]
TimurSadykov
and
gcf-owl-bot[bot]
authored
feat: making iam endpoint universe-aware (#1604)
* feat: making iam endpoint universe-aware * feat: make sign and idtoken endpooints universe aware * add universe_domain parameter for the iam request * fix: test updates --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent f070de0 commit 16c728d

File tree

8 files changed

+178
-21
lines changed

8 files changed

+178
-21
lines changed

‎google/auth/iam.py

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,22 +35,19 @@
3535
http_client.GATEWAY_TIMEOUT,
3636
}
3737

38-
3938
_IAM_SCOPE = ["https://www.googleapis.com/auth/iam"]
4039

4140
_IAM_ENDPOINT = (
42-
"https://iamcredentials.googleapis.com/v1/projects/-"
41+
"https://iamcredentials.{}/v1/projects/-"
4342
+ "/serviceAccounts/{}:generateAccessToken"
4443
)
4544

4645
_IAM_SIGN_ENDPOINT = (
47-
"https://iamcredentials.googleapis.com/v1/projects/-"
48-
+ "/serviceAccounts/{}:signBlob"
46+
"https://iamcredentials.{}/v1/projects/-" + "/serviceAccounts/{}:signBlob"
4947
)
5048

5149
_IAM_IDTOKEN_ENDPOINT = (
52-
"https://iamcredentials.googleapis.com/v1/"
53-
+ "projects/-/serviceAccounts/{}:generateIdToken"
50+
"https://iamcredentials.{}/v1/" + "projects/-/serviceAccounts/{}:generateIdToken"
5451
)
5552

5653

@@ -90,7 +87,9 @@ def _make_signing_request(self, message):
9087
message = _helpers.to_bytes(message)
9188

9289
method = "POST"
93-
url = _IAM_SIGN_ENDPOINT.format(self._service_account_email)
90+
url = _IAM_SIGN_ENDPOINT.format(
91+
self._credentials.universe_domain, self._service_account_email
92+
)
9493
headers = {"Content-Type": "application/json"}
9594
body = json.dumps(
9695
{"payload": base64.b64encode(message).decode("utf-8")}

‎google/auth/impersonated_credentials.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@
4646

4747

4848
def _make_iam_token_request(
49-
request, principal, headers, body, iam_endpoint_override=None
49+
request, principal, headers, body, universe_domain, iam_endpoint_override=None
5050
):
5151
"""Makes a request to the Google Cloud IAM service for an access token.
5252
Args:
@@ -67,7 +67,9 @@ def _make_iam_token_request(
6767
`iamcredentials.googleapis.com` is not enabled or the
6868
`Service Account Token Creator` is not assigned
6969
"""
70-
iam_endpoint = iam_endpoint_override or iam._IAM_ENDPOINT.format(principal)
70+
iam_endpoint = iam_endpoint_override or iam._IAM_ENDPOINT.format(
71+
universe_domain, principal
72+
)
7173

7274
body = json.dumps(body).encode("utf-8")
7375

@@ -219,6 +221,8 @@ def __init__(
219221
and self._source_credentials._always_use_jwt_access
220222
):
221223
self._source_credentials._create_self_signed_jwt(None)
224+
225+
self._universe_domain = source_credentials.universe_domain
222226
self._target_principal = target_principal
223227
self._target_scopes = target_scopes
224228
self._delegates = delegates
@@ -271,13 +275,16 @@ def _update_token(self, request):
271275
principal=self._target_principal,
272276
headers=headers,
273277
body=body,
278+
universe_domain=self.universe_domain,
274279
iam_endpoint_override=self._iam_endpoint_override,
275280
)
276281

277282
def sign_bytes(self, message):
278283
from google.auth.transport.requests import AuthorizedSession
279284

280-
iam_sign_endpoint = iam._IAM_SIGN_ENDPOINT.format(self._target_principal)
285+
iam_sign_endpoint = iam._IAM_SIGN_ENDPOINT.format(
286+
self.universe_domain, self._target_principal
287+
)
281288

282289
body = {
283290
"payload": base64.b64encode(message).decode("utf-8"),
@@ -428,7 +435,8 @@ def refresh(self, request):
428435
from google.auth.transport.requests import AuthorizedSession
429436

430437
iam_sign_endpoint = iam._IAM_IDTOKEN_ENDPOINT.format(
431-
self._target_credentials.signer_email
438+
self._target_credentials.universe_domain,
439+
self._target_credentials.signer_email,
432440
)
433441

434442
body = {

‎google/oauth2/_client.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -319,7 +319,12 @@ def jwt_grant(request, token_uri, assertion, can_retry=True):
319319

320320

321321
def call_iam_generate_id_token_endpoint(
322-
request, iam_id_token_endpoint, signer_email, audience, access_token
322+
request,
323+
iam_id_token_endpoint,
324+
signer_email,
325+
audience,
326+
access_token,
327+
universe_domain,
323328
):
324329
"""Call iam.generateIdToken endpoint to get ID token.
325330
@@ -339,7 +344,7 @@ def call_iam_generate_id_token_endpoint(
339344

340345
response_data = _token_endpoint_request(
341346
request,
342-
iam_id_token_endpoint.format(signer_email),
347+
iam_id_token_endpoint.format(universe_domain, signer_email),
343348
body,
344349
access_token=access_token,
345350
use_json=True,

‎google/oauth2/service_account.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -812,6 +812,7 @@ def _refresh_with_iam_endpoint(self, request):
812812
self.signer_email,
813813
self._target_audience,
814814
jwt_credentials.token.decode(),
815+
self._universe_domain,
815816
)
816817

817818
@_helpers.copy_docstring(credentials.Credentials)

‎tests/compute_engine/test_credentials.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -487,6 +487,16 @@ def test_with_target_audience_integration(self):
487487
},
488488
)
489489

490+
# mock information about universe_domain
491+
responses.add(
492+
responses.GET,
493+
"http://metadata.google.internal/computeMetadata/v1/universe/"
494+
"universe_domain",
495+
status=200,
496+
content_type="application/json",
497+
json={},
498+
)
499+
490500
# mock token for credentials
491501
responses.add(
492502
responses.GET,
@@ -659,6 +669,16 @@ def test_with_quota_project_integration(self):
659669
},
660670
)
661671

672+
# stubby response about universe_domain
673+
responses.add(
674+
responses.GET,
675+
"http://metadata.google.internal/computeMetadata/v1/universe/"
676+
"universe_domain",
677+
status=200,
678+
content_type="application/json",
679+
json={},
680+
)
681+
662682
# mock sign blob endpoint
663683
signature = base64.b64encode(b"some-signature").decode("utf-8")
664684
responses.add(

‎tests/oauth2/test__client.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,7 @@ def test_call_iam_generate_id_token_endpoint():
324324
"fake_email",
325325
"fake_audience",
326326
"fake_access_token",
327+
"googleapis.com",
327328
)
328329

329330
assert (
@@ -361,6 +362,7 @@ def test_call_iam_generate_id_token_endpoint_no_id_token():
361362
"fake_email",
362363
"fake_audience",
363364
"fake_access_token",
365+
"googleapis.com",
364366
)
365367
assert excinfo.match("No ID token in response")
366368

‎tests/oauth2/test_service_account.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -789,7 +789,7 @@ def test_refresh_iam_flow(self, call_iam_generate_id_token_endpoint):
789789
)
790790
request = mock.Mock()
791791
credentials.refresh(request)
792-
req, iam_endpoint, signer_email, target_audience, access_token = call_iam_generate_id_token_endpoint.call_args[
792+
req, iam_endpoint, signer_email, target_audience, access_token, universe_domain = call_iam_generate_id_token_endpoint.call_args[
793793
0
794794
]
795795
assert req == request
@@ -798,6 +798,7 @@ def test_refresh_iam_flow(self, call_iam_generate_id_token_endpoint):
798798
assert target_audience == "https://example.com"
799799
decoded_access_token = jwt.decode(access_token, verify=False)
800800
assert decoded_access_token["scope"] == "https://www.googleapis.com/auth/iam"
801+
assert universe_domain == "googleapis.com"
801802

802803
@mock.patch(
803804
"google.oauth2._client.call_iam_generate_id_token_endpoint", autospec=True
@@ -811,18 +812,19 @@ def test_refresh_iam_flow_non_gdu(self, call_iam_generate_id_token_endpoint):
811812
)
812813
request = mock.Mock()
813814
credentials.refresh(request)
814-
req, iam_endpoint, signer_email, target_audience, access_token = call_iam_generate_id_token_endpoint.call_args[
815+
req, iam_endpoint, signer_email, target_audience, access_token, universe_domain = call_iam_generate_id_token_endpoint.call_args[
815816
0
816817
]
817818
assert req == request
818819
assert (
819820
iam_endpoint
820-
== "https://iamcredentials.fake-universe/v1/projects/-/serviceAccounts/{}:generateIdToken"
821+
== "https://iamcredentials.{}/v1/projects/-/serviceAccounts/{}:generateIdToken"
821822
)
822823
assert signer_email == "[email protected]"
823824
assert target_audience == "https://example.com"
824825
decoded_access_token = jwt.decode(access_token, verify=False)
825826
assert decoded_access_token["scope"] == "https://www.googleapis.com/auth/iam"
827+
assert universe_domain == "fake-universe"
826828

827829
@mock.patch("google.oauth2._client.id_token_jwt_grant", autospec=True)
828830
def test_before_request_refreshes(self, id_token_jwt_grant):

‎tests/test_impersonated_credentials.py

Lines changed: 125 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,13 @@ def test_get_cred_info(self):
146146
"principal": "[email protected]",
147147
}
148148

149+
def test_universe_domain_matching_source(self):
150+
source_credentials = service_account.Credentials(
151+
SIGNER, "[email protected]", TOKEN_URI, universe_domain="foo.bar"
152+
)
153+
credentials = self.make_credentials(source_credentials=source_credentials)
154+
assert credentials.universe_domain == "foo.bar"
155+
149156
def test__make_copy_get_cred_info(self):
150157
credentials = self.make_credentials()
151158
credentials._cred_file_path = "/path/to/file"
@@ -231,6 +238,38 @@ def test_refresh_success(self, use_data_bytes, mock_donor_credentials):
231238
== ACCESS_TOKEN_REQUEST_METRICS_HEADER_VALUE
232239
)
233240

241+
@pytest.mark.parametrize("use_data_bytes", [True, False])
242+
def test_refresh_success_nonGdu(self, use_data_bytes, mock_donor_credentials):
243+
source_credentials = service_account.Credentials(
244+
SIGNER, "[email protected]", TOKEN_URI, universe_domain="foo.bar"
245+
)
246+
credentials = self.make_credentials(
247+
lifetime=None, source_credentials=source_credentials
248+
)
249+
token = "token"
250+
251+
expire_time = (
252+
_helpers.utcnow().replace(microsecond=0) + datetime.timedelta(seconds=500)
253+
).isoformat("T") + "Z"
254+
response_body = {"accessToken": token, "expireTime": expire_time}
255+
256+
request = self.make_request(
257+
data=json.dumps(response_body),
258+
status=http_client.OK,
259+
use_data_bytes=use_data_bytes,
260+
)
261+
262+
credentials.refresh(request)
263+
264+
assert credentials.valid
265+
assert not credentials.expired
266+
# Confirm override endpoint used.
267+
request_kwargs = request.call_args[1]
268+
assert (
269+
request_kwargs["url"]
270+
== "https://iamcredentials.foo.bar/v1/projects/-/serviceAccounts/[email protected]:generateAccessToken"
271+
)
272+
234273
@pytest.mark.parametrize("use_data_bytes", [True, False])
235274
def test_refresh_success_iam_endpoint_override(
236275
self, use_data_bytes, mock_donor_credentials
@@ -397,6 +436,38 @@ def test_service_account_email(self):
397436

398437
def test_sign_bytes(self, mock_donor_credentials, mock_authorizedsession_sign):
399438
credentials = self.make_credentials(lifetime=None)
439+
expected_url = "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/[email protected]:signBlob"
440+
self._sign_bytes_helper(
441+
credentials,
442+
mock_donor_credentials,
443+
mock_authorizedsession_sign,
444+
expected_url,
445+
)
446+
447+
def test_sign_bytes_nonGdu(
448+
self, mock_donor_credentials, mock_authorizedsession_sign
449+
):
450+
source_credentials = service_account.Credentials(
451+
SIGNER, "[email protected]", TOKEN_URI, universe_domain="foo.bar"
452+
)
453+
credentials = self.make_credentials(
454+
lifetime=None, source_credentials=source_credentials
455+
)
456+
expected_url = "https://iamcredentials.foo.bar/v1/projects/-/serviceAccounts/[email protected]:signBlob"
457+
self._sign_bytes_helper(
458+
credentials,
459+
mock_donor_credentials,
460+
mock_authorizedsession_sign,
461+
expected_url,
462+
)
463+
464+
def _sign_bytes_helper(
465+
self,
466+
credentials,
467+
mock_donor_credentials,
468+
mock_authorizedsession_sign,
469+
expected_url,
470+
):
400471
token = "token"
401472

402473
expire_time = (
@@ -412,11 +483,19 @@ def test_sign_bytes(self, mock_donor_credentials, mock_authorizedsession_sign):
412483
request.return_value = response
413484

414485
credentials.refresh(request)
415-
416486
assert credentials.valid
417487
assert not credentials.expired
418488

419489
signature = credentials.sign_bytes(b"signed bytes")
490+
mock_authorizedsession_sign.assert_called_with(
491+
mock.ANY,
492+
"POST",
493+
expected_url,
494+
None,
495+
json={"payload": "c2lnbmVkIGJ5dGVz", "delegates": []},
496+
headers={"Content-Type": "application/json"},
497+
)
498+
420499
assert signature == b"signature"
421500

422501
def test_sign_bytes_failure(self):
@@ -563,6 +642,45 @@ def test_id_token_from_credential(
563642
self, mock_donor_credentials, mock_authorizedsession_idtoken
564643
):
565644
credentials = self.make_credentials(lifetime=None)
645+
target_credentials = self.make_credentials(lifetime=None)
646+
expected_url = "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/[email protected]:generateIdToken"
647+
self._test_id_token_helper(
648+
credentials,
649+
target_credentials,
650+
mock_donor_credentials,
651+
mock_authorizedsession_idtoken,
652+
expected_url,
653+
)
654+
655+
def test_id_token_from_credential_nonGdu(
656+
self, mock_donor_credentials, mock_authorizedsession_idtoken
657+
):
658+
source_credentials = service_account.Credentials(
659+
SIGNER, "[email protected]", TOKEN_URI, universe_domain="foo.bar"
660+
)
661+
credentials = self.make_credentials(
662+
lifetime=None, source_credentials=source_credentials
663+
)
664+
target_credentials = self.make_credentials(
665+
lifetime=None, source_credentials=source_credentials
666+
)
667+
expected_url = "https://iamcredentials.foo.bar/v1/projects/-/serviceAccounts/[email protected]:generateIdToken"
668+
self._test_id_token_helper(
669+
credentials,
670+
target_credentials,
671+
mock_donor_credentials,
672+
mock_authorizedsession_idtoken,
673+
expected_url,
674+
)
675+
676+
def _test_id_token_helper(
677+
self,
678+
credentials,
679+
target_credentials,
680+
mock_donor_credentials,
681+
mock_authorizedsession_idtoken,
682+
expected_url,
683+
):
566684
token = "token"
567685
target_audience = "https://foo.bar"
568686

@@ -580,17 +698,19 @@ def test_id_token_from_credential(
580698
assert credentials.valid
581699
assert not credentials.expired
582700

583-
new_credentials = self.make_credentials(lifetime=None)
584-
585701
id_creds = impersonated_credentials.IDTokenCredentials(
586702
credentials, target_audience=target_audience, include_email=True
587703
)
588-
id_creds = id_creds.from_credentials(target_credentials=new_credentials)
704+
id_creds = id_creds.from_credentials(target_credentials=target_credentials)
589705
id_creds.refresh(request)
590706

707+
args = mock_authorizedsession_idtoken.call_args.args
708+
709+
assert args[2] == expected_url
710+
591711
assert id_creds.token == ID_TOKEN_DATA
592712
assert id_creds._include_email is True
593-
assert id_creds._target_credentials is new_credentials
713+
assert id_creds._target_credentials is target_credentials
594714

595715
def test_id_token_with_target_audience(
596716
self, mock_donor_credentials, mock_authorizedsession_idtoken

0 commit comments

Comments
 (0)