Skip to content

Commit 484c8db

Browse files
authored
feat: IAM signblob retries (#1600)
Add exponential backoff w/ jitter retries to IAM signBlob calls.
1 parent 63f6571 commit 484c8db

File tree

5 files changed

+67
-18
lines changed

5 files changed

+67
-18
lines changed

google/auth/iam.py

Lines changed: 22 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,18 @@
2323
import http.client as http_client
2424
import json
2525

26+
from google.auth import _exponential_backoff
2627
from google.auth import _helpers
2728
from google.auth import crypt
2829
from google.auth import exceptions
2930

31+
IAM_RETRY_CODES = {
32+
http_client.INTERNAL_SERVER_ERROR,
33+
http_client.BAD_GATEWAY,
34+
http_client.SERVICE_UNAVAILABLE,
35+
http_client.GATEWAY_TIMEOUT,
36+
}
37+
3038

3139
_IAM_SCOPE = ["https://www.googleapis.com/auth/iam"]
3240

@@ -88,15 +96,22 @@ def _make_signing_request(self, message):
8896
{"payload": base64.b64encode(message).decode("utf-8")}
8997
).encode("utf-8")
9098

91-
self._credentials.before_request(self._request, method, url, headers)
92-
response = self._request(url=url, method=method, body=body, headers=headers)
99+
retries = _exponential_backoff.ExponentialBackoff()
100+
for _ in retries:
101+
self._credentials.before_request(self._request, method, url, headers)
102+
103+
response = self._request(url=url, method=method, body=body, headers=headers)
104+
105+
if response.status in IAM_RETRY_CODES:
106+
continue
93107

94-
if response.status != http_client.OK:
95-
raise exceptions.TransportError(
96-
"Error calling the IAM signBlob API: {}".format(response.data)
97-
)
108+
if response.status != http_client.OK:
109+
raise exceptions.TransportError(
110+
"Error calling the IAM signBlob API: {}".format(response.data)
111+
)
98112

99-
return json.loads(response.data.decode("utf-8"))
113+
return json.loads(response.data.decode("utf-8"))
114+
raise exceptions.TransportError("exhausted signBlob endpoint retries")
100115

101116
@property
102117
def key_id(self):

google/auth/impersonated_credentials.py

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import http.client as http_client
3232
import json
3333

34+
from google.auth import _exponential_backoff
3435
from google.auth import _helpers
3536
from google.auth import credentials
3637
from google.auth import exceptions
@@ -288,18 +289,22 @@ def sign_bytes(self, message):
288289
authed_session = AuthorizedSession(self._source_credentials)
289290

290291
try:
291-
response = authed_session.post(
292-
url=iam_sign_endpoint, headers=headers, json=body
293-
)
292+
retries = _exponential_backoff.ExponentialBackoff()
293+
for _ in retries:
294+
response = authed_session.post(
295+
url=iam_sign_endpoint, headers=headers, json=body
296+
)
297+
if response.status_code in iam.IAM_RETRY_CODES:
298+
continue
299+
if response.status_code != http_client.OK:
300+
raise exceptions.TransportError(
301+
"Error calling sign_bytes: {}".format(response.json())
302+
)
303+
304+
return base64.b64decode(response.json()["signedBlob"])
294305
finally:
295306
authed_session.close()
296-
297-
if response.status_code != http_client.OK:
298-
raise exceptions.TransportError(
299-
"Error calling sign_bytes: {}".format(response.json())
300-
)
301-
302-
return base64.b64decode(response.json()["signedBlob"])
307+
raise exceptions.TransportError("exhausted signBlob endpoint retries")
303308

304309
@property
305310
def signer_email(self):

system_tests/secrets.tar.enc

0 Bytes
Binary file not shown.

tests/test_iam.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ def test_sign_bytes(self):
9191
assert returned_signature == signature
9292
kwargs = request.call_args[1]
9393
assert kwargs["headers"]["Content-Type"] == "application/json"
94+
request.call_count == 1
9495

9596
def test_sign_bytes_failure(self):
9697
request = make_request(http_client.UNAUTHORIZED)
@@ -100,3 +101,15 @@ def test_sign_bytes_failure(self):
100101

101102
with pytest.raises(exceptions.TransportError):
102103
signer.sign("123")
104+
request.call_count == 1
105+
106+
@mock.patch("time.sleep", return_value=None)
107+
def test_sign_bytes_retryable_failure(self, mock_time):
108+
request = make_request(http_client.INTERNAL_SERVER_ERROR)
109+
credentials = make_credentials()
110+
111+
signer = iam.Signer(request, credentials, mock.sentinel.service_account_email)
112+
113+
with pytest.raises(exceptions.TransportError):
114+
signer.sign("123")
115+
request.call_count == 3

tests/test_impersonated_credentials.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -426,12 +426,28 @@ def test_sign_bytes_failure(self):
426426
"google.auth.transport.requests.AuthorizedSession.request", autospec=True
427427
) as auth_session:
428428
data = {"error": {"code": 403, "message": "unauthorized"}}
429-
auth_session.return_value = MockResponse(data, http_client.FORBIDDEN)
429+
mock_response = MockResponse(data, http_client.UNAUTHORIZED)
430+
auth_session.return_value = mock_response
430431

431432
with pytest.raises(exceptions.TransportError) as excinfo:
432433
credentials.sign_bytes(b"foo")
433434
assert excinfo.match("'code': 403")
434435

436+
@mock.patch("time.sleep", return_value=None)
437+
def test_sign_bytes_retryable_failure(self, mock_time):
438+
credentials = self.make_credentials(lifetime=None)
439+
440+
with mock.patch(
441+
"google.auth.transport.requests.AuthorizedSession.request", autospec=True
442+
) as auth_session:
443+
data = {"error": {"code": 500, "message": "internal_failure"}}
444+
mock_response = MockResponse(data, http_client.INTERNAL_SERVER_ERROR)
445+
auth_session.return_value = mock_response
446+
447+
with pytest.raises(exceptions.TransportError) as excinfo:
448+
credentials.sign_bytes(b"foo")
449+
assert excinfo.match("exhausted signBlob endpoint retries")
450+
435451
def test_with_quota_project(self):
436452
credentials = self.make_credentials()
437453

0 commit comments

Comments
 (0)