Skip to content

Commit d42f30a

Browse files
authoredAug 6, 2024··
fix: Retry sign blob call with exponential backoff (#1452)
* fix: Retry sign blob call with exponential backoff * chore: Add header for IamUtilsTest * chore: Add comments for IamUtilsTest * chore: Move IAM Retry status codes to IamUtils
1 parent c83a71f commit d42f30a

File tree

6 files changed

+161
-27
lines changed

6 files changed

+161
-27
lines changed
 

‎oauth2_http/java/com/google/auth/oauth2/GoogleAuthException.java

+2-3
Original file line numberDiff line numberDiff line change
@@ -158,11 +158,10 @@ static GoogleAuthException createWithTokenEndpointIOException(
158158
if (message == null) {
159159
// TODO: temporarily setting retry Count to service account default to remove a direct
160160
// dependency, to be reverted after release
161-
return new GoogleAuthException(
162-
true, ServiceAccountCredentials.DEFAULT_NUMBER_OF_RETRIES, ioException);
161+
return new GoogleAuthException(true, OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES, ioException);
163162
} else {
164163
return new GoogleAuthException(
165-
true, ServiceAccountCredentials.DEFAULT_NUMBER_OF_RETRIES, message, ioException);
164+
true, OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES, message, ioException);
166165
}
167166
}
168167

‎oauth2_http/java/com/google/auth/oauth2/IamUtils.java

+38-10
Original file line numberDiff line numberDiff line change
@@ -32,21 +32,28 @@
3232
package com.google.auth.oauth2;
3333

3434
import com.google.api.client.http.GenericUrl;
35+
import com.google.api.client.http.HttpBackOffIOExceptionHandler;
36+
import com.google.api.client.http.HttpBackOffUnsuccessfulResponseHandler;
3537
import com.google.api.client.http.HttpRequest;
38+
import com.google.api.client.http.HttpRequestFactory;
3639
import com.google.api.client.http.HttpResponse;
3740
import com.google.api.client.http.HttpStatusCodes;
3841
import com.google.api.client.http.HttpTransport;
3942
import com.google.api.client.http.json.JsonHttpContent;
4043
import com.google.api.client.json.GenericJson;
4144
import com.google.api.client.json.JsonObjectParser;
45+
import com.google.api.client.util.ExponentialBackOff;
4246
import com.google.api.client.util.GenericData;
4347
import com.google.auth.Credentials;
4448
import com.google.auth.ServiceAccountSigner;
4549
import com.google.auth.http.HttpCredentialsAdapter;
4650
import com.google.common.io.BaseEncoding;
4751
import java.io.IOException;
4852
import java.io.InputStream;
53+
import java.util.Arrays;
54+
import java.util.HashSet;
4955
import java.util.Map;
56+
import java.util.Set;
5057

5158
/**
5259
* This internal class provides shared utilities for interacting with the IAM API for common
@@ -60,6 +67,11 @@ class IamUtils {
6067
private static final String PARSE_ERROR_MESSAGE = "Error parsing error message response. ";
6168
private static final String PARSE_ERROR_SIGNATURE = "Error parsing signature response. ";
6269

70+
// Following guidance for IAM retries:
71+
// https://cloud.google.com/iam/docs/retry-strategy#errors-to-retry
72+
static final Set<Integer> IAM_RETRYABLE_STATUS_CODES =
73+
new HashSet<>(Arrays.asList(500, 502, 503, 504));
74+
6375
/**
6476
* Returns a signature for the provided bytes.
6577
*
@@ -78,11 +90,12 @@ static byte[] sign(
7890
byte[] toSign,
7991
Map<String, ?> additionalFields) {
8092
BaseEncoding base64 = BaseEncoding.base64();
93+
HttpRequestFactory factory =
94+
transport.createRequestFactory(new HttpCredentialsAdapter(credentials));
8195
String signature;
8296
try {
8397
signature =
84-
getSignature(
85-
serviceAccountEmail, credentials, transport, base64.encode(toSign), additionalFields);
98+
getSignature(serviceAccountEmail, base64.encode(toSign), additionalFields, factory);
8699
} catch (IOException ex) {
87100
throw new ServiceAccountSigner.SigningException("Failed to sign the provided bytes", ex);
88101
}
@@ -91,10 +104,9 @@ static byte[] sign(
91104

92105
private static String getSignature(
93106
String serviceAccountEmail,
94-
Credentials credentials,
95-
HttpTransport transport,
96107
String bytes,
97-
Map<String, ?> additionalFields)
108+
Map<String, ?> additionalFields,
109+
HttpRequestFactory factory)
98110
throws IOException {
99111
String signBlobUrl = String.format(SIGN_BLOB_URL_FORMAT, serviceAccountEmail);
100112
GenericUrl genericUrl = new GenericUrl(signBlobUrl);
@@ -106,13 +118,27 @@ private static String getSignature(
106118
}
107119
JsonHttpContent signContent = new JsonHttpContent(OAuth2Utils.JSON_FACTORY, signRequest);
108120

109-
HttpCredentialsAdapter adapter = new HttpCredentialsAdapter(credentials);
110-
HttpRequest request =
111-
transport.createRequestFactory(adapter).buildPostRequest(genericUrl, signContent);
121+
HttpRequest request = factory.buildPostRequest(genericUrl, signContent);
112122

113123
JsonObjectParser parser = new JsonObjectParser(OAuth2Utils.JSON_FACTORY);
114124
request.setParser(parser);
115125
request.setThrowExceptionOnExecuteError(false);
126+
request.setNumberOfRetries(OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES);
127+
128+
ExponentialBackOff backoff =
129+
new ExponentialBackOff.Builder()
130+
.setInitialIntervalMillis(OAuth2Utils.INITIAL_RETRY_INTERVAL_MILLIS)
131+
.setRandomizationFactor(OAuth2Utils.RETRY_RANDOMIZATION_FACTOR)
132+
.setMultiplier(OAuth2Utils.RETRY_MULTIPLIER)
133+
.build();
134+
135+
// Retry on 500, 502, 503, and 503 status codes
136+
request.setUnsuccessfulResponseHandler(
137+
new HttpBackOffUnsuccessfulResponseHandler(backoff)
138+
.setBackOffRequired(
139+
response ->
140+
IamUtils.IAM_RETRYABLE_STATUS_CODES.contains(response.getStatusCode())));
141+
request.setIOExceptionHandler(new HttpBackOffIOExceptionHandler(backoff));
116142

117143
HttpResponse response = request.execute();
118144
int statusCode = response.getStatusCode();
@@ -125,6 +151,8 @@ private static String getSignature(
125151
String.format(
126152
"Error code %s trying to sign provided bytes: %s", statusCode, errorMessage));
127153
}
154+
155+
// Request will have retried a 5xx error 3 times and is still receiving a 5xx error code
128156
if (statusCode != HttpStatusCodes.STATUS_CODE_OK) {
129157
throw new IOException(
130158
String.format(
@@ -152,8 +180,8 @@ private static String getSignature(
152180
* @param additionalFields additional fields to send in the IAM call
153181
* @return IdToken issed to the serviceAccount
154182
* @throws IOException if the IdToken cannot be issued.
155-
* @see
156-
* https://cloud.google.com/iam/credentials/reference/rest/v1/projects.serviceAccounts/generateIdToken
183+
* @see <a
184+
* href="https://cloud.google.com/iam/credentials/reference/rest/v1/projects.serviceAccounts/generateIdToken">...</a>
157185
*/
158186
static IdToken getIdToken(
159187
String serviceAccountEmail,

‎oauth2_http/java/com/google/auth/oauth2/OAuth2Utils.java

+5
Original file line numberDiff line numberDiff line change
@@ -92,6 +92,11 @@ class OAuth2Utils {
9292

9393
static final String TOKEN_RESPONSE_SCOPE = "scope";
9494

95+
static final int INITIAL_RETRY_INTERVAL_MILLIS = 1000;
96+
static final double RETRY_RANDOMIZATION_FACTOR = 0.1;
97+
static final double RETRY_MULTIPLIER = 2;
98+
static final int DEFAULT_NUMBER_OF_RETRIES = 3;
99+
95100
// Includes expected server errors from Google token endpoint
96101
// Other 5xx codes are either not used or retries are unlikely to succeed
97102
public static final Set<Integer> TOKEN_ENDPOINT_RETRYABLE_STATUS_CODES =

‎oauth2_http/java/com/google/auth/oauth2/ServiceAccountCredentials.java

+4-8
Original file line numberDiff line numberDiff line change
@@ -89,10 +89,6 @@ public class ServiceAccountCredentials extends GoogleCredentials
8989
private static final String PARSE_ERROR_PREFIX = "Error parsing token refresh response. ";
9090
private static final int TWELVE_HOURS_IN_SECONDS = 43200;
9191
private static final int DEFAULT_LIFETIME_IN_SECONDS = 3600;
92-
private static final int INITIAL_RETRY_INTERVAL_MILLIS = 1000;
93-
private static final double RETRY_RANDOMIZATION_FACTOR = 0.1;
94-
private static final double RETRY_MULTIPLIER = 2;
95-
static final int DEFAULT_NUMBER_OF_RETRIES = 3;
9692

9793
private final String clientId;
9894
private final String clientEmail;
@@ -505,17 +501,17 @@ public AccessToken refreshAccessToken() throws IOException {
505501
HttpRequest request = requestFactory.buildPostRequest(new GenericUrl(tokenServerUri), content);
506502

507503
if (this.defaultRetriesEnabled) {
508-
request.setNumberOfRetries(DEFAULT_NUMBER_OF_RETRIES);
504+
request.setNumberOfRetries(OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES);
509505
} else {
510506
request.setNumberOfRetries(0);
511507
}
512508
request.setParser(new JsonObjectParser(jsonFactory));
513509

514510
ExponentialBackOff backoff =
515511
new ExponentialBackOff.Builder()
516-
.setInitialIntervalMillis(INITIAL_RETRY_INTERVAL_MILLIS)
517-
.setRandomizationFactor(RETRY_RANDOMIZATION_FACTOR)
518-
.setMultiplier(RETRY_MULTIPLIER)
512+
.setInitialIntervalMillis(OAuth2Utils.INITIAL_RETRY_INTERVAL_MILLIS)
513+
.setRandomizationFactor(OAuth2Utils.RETRY_RANDOMIZATION_FACTOR)
514+
.setMultiplier(OAuth2Utils.RETRY_MULTIPLIER)
519515
.build();
520516

521517
request.setUnsuccessfulResponseHandler(

‎oauth2_http/java/com/google/auth/oauth2/TokenVerifier.java

+3-6
Original file line numberDiff line numberDiff line change
@@ -279,9 +279,6 @@ public TokenVerifier build() {
279279
/** Custom CacheLoader for mapping certificate urls to the contained public keys. */
280280
static class PublicKeyLoader extends CacheLoader<String, Map<String, PublicKey>> {
281281
private static final int DEFAULT_NUMBER_OF_RETRIES = 2;
282-
private static final int INITIAL_RETRY_INTERVAL_MILLIS = 1000;
283-
private static final double RETRY_RANDOMIZATION_FACTOR = 0.1;
284-
private static final double RETRY_MULTIPLIER = 2;
285282
private final HttpTransportFactory httpTransportFactory;
286283

287284
/**
@@ -330,9 +327,9 @@ public Map<String, PublicKey> load(String certificateUrl) throws Exception {
330327

331328
ExponentialBackOff backoff =
332329
new ExponentialBackOff.Builder()
333-
.setInitialIntervalMillis(INITIAL_RETRY_INTERVAL_MILLIS)
334-
.setRandomizationFactor(RETRY_RANDOMIZATION_FACTOR)
335-
.setMultiplier(RETRY_MULTIPLIER)
330+
.setInitialIntervalMillis(OAuth2Utils.INITIAL_RETRY_INTERVAL_MILLIS)
331+
.setRandomizationFactor(OAuth2Utils.RETRY_RANDOMIZATION_FACTOR)
332+
.setMultiplier(OAuth2Utils.RETRY_MULTIPLIER)
336333
.build();
337334

338335
request.setUnsuccessfulResponseHandler(
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
/*
2+
* Copyright 2024, Google Inc. All rights reserved.
3+
*
4+
* Redistribution and use in source and binary forms, with or without
5+
* modification, are permitted provided that the following conditions are
6+
* met:
7+
*
8+
* * Redistributions of source code must retain the above copyright
9+
* notice, this list of conditions and the following disclaimer.
10+
* * Redistributions in binary form must reproduce the above
11+
* copyright notice, this list of conditions and the following disclaimer
12+
* in the documentation and/or other materials provided with the
13+
* distribution.
14+
*
15+
* * Neither the name of Google Inc. nor the names of its
16+
* contributors may be used to endorse or promote products derived from
17+
* this software without specific prior written permission.
18+
*
19+
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
20+
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
21+
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
22+
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
23+
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
24+
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
25+
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
26+
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
27+
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
28+
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29+
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30+
*/
31+
package com.google.auth.oauth2;
32+
33+
import static org.junit.Assert.assertArrayEquals;
34+
import static org.junit.Assert.assertThrows;
35+
import static org.junit.Assert.assertTrue;
36+
37+
import com.google.api.client.http.HttpStatusCodes;
38+
import com.google.auth.ServiceAccountSigner;
39+
import com.google.common.collect.ImmutableMap;
40+
import java.io.IOException;
41+
import org.junit.Test;
42+
import org.junit.runner.RunWith;
43+
import org.junit.runners.JUnit4;
44+
import org.mockito.Mockito;
45+
46+
@RunWith(JUnit4.class)
47+
public class IamUtilsTest {
48+
49+
private static final String CLIENT_EMAIL =
50+
"36680232662-vrd7ji19qe3nelgchd0ah2csanun6bnr@developer.gserviceaccount.com";
51+
52+
@Test
53+
public void sign_noRetry() throws IOException {
54+
byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
55+
56+
// Mock this call because signing requires an access token. The call is initialized with
57+
// HttpCredentialsAdapter which will make a call to get the access token
58+
ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
59+
Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
60+
61+
ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
62+
new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
63+
transportFactory.transport.setSignedBlob(expectedSignature);
64+
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
65+
66+
byte[] signature =
67+
IamUtils.sign(
68+
CLIENT_EMAIL,
69+
credentials,
70+
transportFactory.transport,
71+
expectedSignature,
72+
ImmutableMap.of());
73+
assertArrayEquals(expectedSignature, signature);
74+
}
75+
76+
@Test
77+
public void sign_4xxServerError_exception() throws IOException {
78+
byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};
79+
80+
// Mock this call because signing requires an access token. The call is initialized with
81+
// HttpCredentialsAdapter which will make a call to get the access token
82+
ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
83+
Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());
84+
85+
ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
86+
new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
87+
transportFactory.transport.setSignedBlob(expectedSignature);
88+
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
89+
transportFactory.transport.setErrorResponseCodeAndMessage(
90+
HttpStatusCodes.STATUS_CODE_UNAUTHORIZED, "Failed to sign the provided bytes");
91+
92+
ServiceAccountSigner.SigningException exception =
93+
assertThrows(
94+
ServiceAccountSigner.SigningException.class,
95+
() ->
96+
IamUtils.sign(
97+
CLIENT_EMAIL,
98+
credentials,
99+
transportFactory.transport,
100+
expectedSignature,
101+
ImmutableMap.of()));
102+
assertTrue(exception.getMessage().contains("Failed to sign the provided bytes"));
103+
assertTrue(
104+
exception
105+
.getCause()
106+
.getMessage()
107+
.contains("Error code 401 trying to sign provided bytes:"));
108+
}
109+
}

0 commit comments

Comments
 (0)
Please sign in to comment.