Skip to content

Commit

Permalink
fix: Retry sign blob call with exponential backoff (#1452)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
lqiu96 authored Aug 6, 2024
1 parent c83a71f commit d42f30a
Show file tree
Hide file tree
Showing 6 changed files with 161 additions and 27 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -158,11 +158,10 @@ static GoogleAuthException createWithTokenEndpointIOException(
if (message == null) {
// TODO: temporarily setting retry Count to service account default to remove a direct
// dependency, to be reverted after release
return new GoogleAuthException(
true, ServiceAccountCredentials.DEFAULT_NUMBER_OF_RETRIES, ioException);
return new GoogleAuthException(true, OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES, ioException);
} else {
return new GoogleAuthException(
true, ServiceAccountCredentials.DEFAULT_NUMBER_OF_RETRIES, message, ioException);
true, OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES, message, ioException);
}
}

Expand Down
48 changes: 38 additions & 10 deletions oauth2_http/java/com/google/auth/oauth2/IamUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -32,21 +32,28 @@
package com.google.auth.oauth2;

import com.google.api.client.http.GenericUrl;
import com.google.api.client.http.HttpBackOffIOExceptionHandler;
import com.google.api.client.http.HttpBackOffUnsuccessfulResponseHandler;
import com.google.api.client.http.HttpRequest;
import com.google.api.client.http.HttpRequestFactory;
import com.google.api.client.http.HttpResponse;
import com.google.api.client.http.HttpStatusCodes;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.http.json.JsonHttpContent;
import com.google.api.client.json.GenericJson;
import com.google.api.client.json.JsonObjectParser;
import com.google.api.client.util.ExponentialBackOff;
import com.google.api.client.util.GenericData;
import com.google.auth.Credentials;
import com.google.auth.ServiceAccountSigner;
import com.google.auth.http.HttpCredentialsAdapter;
import com.google.common.io.BaseEncoding;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

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

// Following guidance for IAM retries:
// https://cloud.google.com/iam/docs/retry-strategy#errors-to-retry
static final Set<Integer> IAM_RETRYABLE_STATUS_CODES =
new HashSet<>(Arrays.asList(500, 502, 503, 504));

/**
* Returns a signature for the provided bytes.
*
Expand All @@ -78,11 +90,12 @@ static byte[] sign(
byte[] toSign,
Map<String, ?> additionalFields) {
BaseEncoding base64 = BaseEncoding.base64();
HttpRequestFactory factory =
transport.createRequestFactory(new HttpCredentialsAdapter(credentials));
String signature;
try {
signature =
getSignature(
serviceAccountEmail, credentials, transport, base64.encode(toSign), additionalFields);
getSignature(serviceAccountEmail, base64.encode(toSign), additionalFields, factory);
} catch (IOException ex) {
throw new ServiceAccountSigner.SigningException("Failed to sign the provided bytes", ex);
}
Expand All @@ -91,10 +104,9 @@ static byte[] sign(

private static String getSignature(
String serviceAccountEmail,
Credentials credentials,
HttpTransport transport,
String bytes,
Map<String, ?> additionalFields)
Map<String, ?> additionalFields,
HttpRequestFactory factory)
throws IOException {
String signBlobUrl = String.format(SIGN_BLOB_URL_FORMAT, serviceAccountEmail);
GenericUrl genericUrl = new GenericUrl(signBlobUrl);
Expand All @@ -106,13 +118,27 @@ private static String getSignature(
}
JsonHttpContent signContent = new JsonHttpContent(OAuth2Utils.JSON_FACTORY, signRequest);

HttpCredentialsAdapter adapter = new HttpCredentialsAdapter(credentials);
HttpRequest request =
transport.createRequestFactory(adapter).buildPostRequest(genericUrl, signContent);
HttpRequest request = factory.buildPostRequest(genericUrl, signContent);

JsonObjectParser parser = new JsonObjectParser(OAuth2Utils.JSON_FACTORY);
request.setParser(parser);
request.setThrowExceptionOnExecuteError(false);
request.setNumberOfRetries(OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES);

ExponentialBackOff backoff =
new ExponentialBackOff.Builder()
.setInitialIntervalMillis(OAuth2Utils.INITIAL_RETRY_INTERVAL_MILLIS)
.setRandomizationFactor(OAuth2Utils.RETRY_RANDOMIZATION_FACTOR)
.setMultiplier(OAuth2Utils.RETRY_MULTIPLIER)
.build();

// Retry on 500, 502, 503, and 503 status codes
request.setUnsuccessfulResponseHandler(
new HttpBackOffUnsuccessfulResponseHandler(backoff)
.setBackOffRequired(
response ->
IamUtils.IAM_RETRYABLE_STATUS_CODES.contains(response.getStatusCode())));
request.setIOExceptionHandler(new HttpBackOffIOExceptionHandler(backoff));

HttpResponse response = request.execute();
int statusCode = response.getStatusCode();
Expand All @@ -125,6 +151,8 @@ private static String getSignature(
String.format(
"Error code %s trying to sign provided bytes: %s", statusCode, errorMessage));
}

// Request will have retried a 5xx error 3 times and is still receiving a 5xx error code
if (statusCode != HttpStatusCodes.STATUS_CODE_OK) {
throw new IOException(
String.format(
Expand Down Expand Up @@ -152,8 +180,8 @@ private static String getSignature(
* @param additionalFields additional fields to send in the IAM call
* @return IdToken issed to the serviceAccount
* @throws IOException if the IdToken cannot be issued.
* @see
* https://cloud.google.com/iam/credentials/reference/rest/v1/projects.serviceAccounts/generateIdToken
* @see <a
* href="https://cloud.google.com/iam/credentials/reference/rest/v1/projects.serviceAccounts/generateIdToken">...</a>
*/
static IdToken getIdToken(
String serviceAccountEmail,
Expand Down
5 changes: 5 additions & 0 deletions oauth2_http/java/com/google/auth/oauth2/OAuth2Utils.java
Original file line number Diff line number Diff line change
Expand Up @@ -92,6 +92,11 @@ class OAuth2Utils {

static final String TOKEN_RESPONSE_SCOPE = "scope";

static final int INITIAL_RETRY_INTERVAL_MILLIS = 1000;
static final double RETRY_RANDOMIZATION_FACTOR = 0.1;
static final double RETRY_MULTIPLIER = 2;
static final int DEFAULT_NUMBER_OF_RETRIES = 3;

// Includes expected server errors from Google token endpoint
// Other 5xx codes are either not used or retries are unlikely to succeed
public static final Set<Integer> TOKEN_ENDPOINT_RETRYABLE_STATUS_CODES =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -89,10 +89,6 @@ public class ServiceAccountCredentials extends GoogleCredentials
private static final String PARSE_ERROR_PREFIX = "Error parsing token refresh response. ";
private static final int TWELVE_HOURS_IN_SECONDS = 43200;
private static final int DEFAULT_LIFETIME_IN_SECONDS = 3600;
private static final int INITIAL_RETRY_INTERVAL_MILLIS = 1000;
private static final double RETRY_RANDOMIZATION_FACTOR = 0.1;
private static final double RETRY_MULTIPLIER = 2;
static final int DEFAULT_NUMBER_OF_RETRIES = 3;

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

if (this.defaultRetriesEnabled) {
request.setNumberOfRetries(DEFAULT_NUMBER_OF_RETRIES);
request.setNumberOfRetries(OAuth2Utils.DEFAULT_NUMBER_OF_RETRIES);
} else {
request.setNumberOfRetries(0);
}
request.setParser(new JsonObjectParser(jsonFactory));

ExponentialBackOff backoff =
new ExponentialBackOff.Builder()
.setInitialIntervalMillis(INITIAL_RETRY_INTERVAL_MILLIS)
.setRandomizationFactor(RETRY_RANDOMIZATION_FACTOR)
.setMultiplier(RETRY_MULTIPLIER)
.setInitialIntervalMillis(OAuth2Utils.INITIAL_RETRY_INTERVAL_MILLIS)
.setRandomizationFactor(OAuth2Utils.RETRY_RANDOMIZATION_FACTOR)
.setMultiplier(OAuth2Utils.RETRY_MULTIPLIER)
.build();

request.setUnsuccessfulResponseHandler(
Expand Down
9 changes: 3 additions & 6 deletions oauth2_http/java/com/google/auth/oauth2/TokenVerifier.java
Original file line number Diff line number Diff line change
Expand Up @@ -279,9 +279,6 @@ public TokenVerifier build() {
/** Custom CacheLoader for mapping certificate urls to the contained public keys. */
static class PublicKeyLoader extends CacheLoader<String, Map<String, PublicKey>> {
private static final int DEFAULT_NUMBER_OF_RETRIES = 2;
private static final int INITIAL_RETRY_INTERVAL_MILLIS = 1000;
private static final double RETRY_RANDOMIZATION_FACTOR = 0.1;
private static final double RETRY_MULTIPLIER = 2;
private final HttpTransportFactory httpTransportFactory;

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

ExponentialBackOff backoff =
new ExponentialBackOff.Builder()
.setInitialIntervalMillis(INITIAL_RETRY_INTERVAL_MILLIS)
.setRandomizationFactor(RETRY_RANDOMIZATION_FACTOR)
.setMultiplier(RETRY_MULTIPLIER)
.setInitialIntervalMillis(OAuth2Utils.INITIAL_RETRY_INTERVAL_MILLIS)
.setRandomizationFactor(OAuth2Utils.RETRY_RANDOMIZATION_FACTOR)
.setMultiplier(OAuth2Utils.RETRY_MULTIPLIER)
.build();

request.setUnsuccessfulResponseHandler(
Expand Down
109 changes: 109 additions & 0 deletions oauth2_http/javatests/com/google/auth/oauth2/IamUtilsTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
/*
* Copyright 2024, Google Inc. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are
* met:
*
* * Redistributions of source code must retain the above copyright
* notice, this list of conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above
* copyright notice, this list of conditions and the following disclaimer
* in the documentation and/or other materials provided with the
* distribution.
*
* * Neither the name of Google Inc. nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/
package com.google.auth.oauth2;

import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;

import com.google.api.client.http.HttpStatusCodes;
import com.google.auth.ServiceAccountSigner;
import com.google.common.collect.ImmutableMap;
import java.io.IOException;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.Mockito;

@RunWith(JUnit4.class)
public class IamUtilsTest {

private static final String CLIENT_EMAIL =
"36680232662-vrd7ji19qe3nelgchd0ah2csanun6bnr@developer.gserviceaccount.com";

@Test
public void sign_noRetry() throws IOException {
byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};

// Mock this call because signing requires an access token. The call is initialized with
// HttpCredentialsAdapter which will make a call to get the access token
ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());

ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
transportFactory.transport.setSignedBlob(expectedSignature);
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);

byte[] signature =
IamUtils.sign(
CLIENT_EMAIL,
credentials,
transportFactory.transport,
expectedSignature,
ImmutableMap.of());
assertArrayEquals(expectedSignature, signature);
}

@Test
public void sign_4xxServerError_exception() throws IOException {
byte[] expectedSignature = {0xD, 0xE, 0xA, 0xD};

// Mock this call because signing requires an access token. The call is initialized with
// HttpCredentialsAdapter which will make a call to get the access token
ServiceAccountCredentials credentials = Mockito.mock(ServiceAccountCredentials.class);
Mockito.when(credentials.getRequestMetadata(Mockito.any())).thenReturn(ImmutableMap.of());

ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory transportFactory =
new ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory();
transportFactory.transport.setSignedBlob(expectedSignature);
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
transportFactory.transport.setErrorResponseCodeAndMessage(
HttpStatusCodes.STATUS_CODE_UNAUTHORIZED, "Failed to sign the provided bytes");

ServiceAccountSigner.SigningException exception =
assertThrows(
ServiceAccountSigner.SigningException.class,
() ->
IamUtils.sign(
CLIENT_EMAIL,
credentials,
transportFactory.transport,
expectedSignature,
ImmutableMap.of()));
assertTrue(exception.getMessage().contains("Failed to sign the provided bytes"));
assertTrue(
exception
.getCause()
.getMessage()
.contains("Error code 401 trying to sign provided bytes:"));
}
}

0 comments on commit d42f30a

Please sign in to comment.