Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Enable Retries For Auth Requests #1791

Merged
merged 5 commits into from
Apr 12, 2024
Merged
Show file tree
Hide file tree
Changes from 4 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
18 changes: 18 additions & 0 deletions src/auth/authclient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -259,4 +259,22 @@ export abstract class AuthClient
}
return headers;
}

/**
* Retry config for Auth-related requests.
*
* @remarks
*
* This is not a part of the default {@link AuthClient.transporter transporter/gaxios}
* config as some downstream APIs would prefer if customers explicitly enable retries,
* such as GCS.
*/
protected static get RETRY_CONFIG(): GaxiosOptions {
return {
retry: true,
retryConfig: {
httpMethodsToRetry: ['GET', 'PUT', 'POST', 'HEAD', 'OPTIONS', 'DELETE'],
},
};
}
}
5 changes: 5 additions & 0 deletions src/auth/awsclient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -192,6 +192,7 @@ export class AwsClient extends BaseExternalAccountClient {
// Generate signed request to AWS STS GetCallerIdentity API.
// Use the required regional endpoint. Otherwise, the request will fail.
const options = await this.awsRequestSigner.getRequestOptions({
...AwsClient.RETRY_CONFIG,
url: this.regionalCredVerificationUrl.replace('{region}', this.region),
method: 'POST',
});
Expand Down Expand Up @@ -240,6 +241,7 @@ export class AwsClient extends BaseExternalAccountClient {
*/
private async getImdsV2SessionToken(): Promise<string> {
const opts: GaxiosOptions = {
...AwsClient.RETRY_CONFIG,
url: this.imdsV2SessionTokenUrl,
method: 'PUT',
responseType: 'text',
Expand All @@ -266,6 +268,7 @@ export class AwsClient extends BaseExternalAccountClient {
);
}
const opts: GaxiosOptions = {
...AwsClient.RETRY_CONFIG,
url: this.regionUrl,
method: 'GET',
responseType: 'text',
Expand All @@ -290,6 +293,7 @@ export class AwsClient extends BaseExternalAccountClient {
);
}
const opts: GaxiosOptions = {
...AwsClient.RETRY_CONFIG,
url: this.securityCredentialsUrl,
method: 'GET',
responseType: 'text',
Expand All @@ -313,6 +317,7 @@ export class AwsClient extends BaseExternalAccountClient {
): Promise<AwsSecurityCredentialsResponse> {
const response =
await this.transporter.request<AwsSecurityCredentialsResponse>({
...AwsClient.RETRY_CONFIG,
url: `${this.securityCredentialsUrl}/${roleName}`,
responseType: 'json',
headers: headers,
Expand Down
8 changes: 5 additions & 3 deletions src/auth/baseexternalclient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -381,6 +381,7 @@ export abstract class BaseExternalAccountClient extends AuthClient {
// Preferable not to use request() to avoid retrial policies.
const headers = await this.getRequestHeaders();
const response = await this.transporter.request<ProjectInfo>({
...BaseExternalAccountClient.RETRY_CONFIG,
headers,
url: `${this.cloudResourceManagerURL.toString()}${projectNumber}`,
responseType: 'json',
Expand All @@ -395,12 +396,12 @@ export abstract class BaseExternalAccountClient extends AuthClient {
* Authenticates the provided HTTP request, processes it and resolves with the
* returned response.
* @param opts The HTTP request options.
* @param retry Whether the current attempt is a retry after a failed attempt.
* @param reAuthRetried Whether the current attempt is a retry after a failed attempt due to an auth failure.
* @return A promise that resolves with the successful response.
*/
protected async requestAsync<T>(
opts: GaxiosOptions,
retry = false
reAuthRetried = false
): Promise<GaxiosResponse<T>> {
let response: GaxiosResponse;
try {
Expand All @@ -426,7 +427,7 @@ export abstract class BaseExternalAccountClient extends AuthClient {
const isReadableStream = res.config.data instanceof stream.Readable;
const isAuthErr = statusCode === 401 || statusCode === 403;
if (
!retry &&
!reAuthRetried &&
isAuthErr &&
!isReadableStream &&
this.forceRefreshOnFailure
Expand Down Expand Up @@ -554,6 +555,7 @@ export abstract class BaseExternalAccountClient extends AuthClient {
token: string
): Promise<CredentialsWithResponse> {
const opts: GaxiosOptions = {
...BaseExternalAccountClient.RETRY_CONFIG,
url: this.serviceAccountImpersonationUrl!,
method: 'POST',
headers: {
Expand Down
6 changes: 3 additions & 3 deletions src/auth/downscopedclient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -250,12 +250,12 @@ export class DownscopedClient extends AuthClient {
* Authenticates the provided HTTP request, processes it and resolves with the
* returned response.
* @param opts The HTTP request options.
* @param retry Whether the current attempt is a retry after a failed attempt.
* @param reAuthRetried Whether the current attempt is a retry after a failed attempt due to an auth failure
* @return A promise that resolves with the successful response.
*/
protected async requestAsync<T>(
opts: GaxiosOptions,
retry = false
reAuthRetried = false
): Promise<GaxiosResponse<T>> {
let response: GaxiosResponse;
try {
Expand All @@ -281,7 +281,7 @@ export class DownscopedClient extends AuthClient {
const isReadableStream = res.config.data instanceof stream.Readable;
const isAuthErr = statusCode === 401 || statusCode === 403;
if (
!retry &&
!reAuthRetried &&
isAuthErr &&
!isReadableStream &&
this.forceRefreshOnFailure
Expand Down
7 changes: 4 additions & 3 deletions src/auth/externalAccountAuthorizedUserClient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,7 @@ class ExternalAccountAuthorizedUserHandler extends OAuthClientAuthHandler {
};

const opts: GaxiosOptions = {
...ExternalAccountAuthorizedUserHandler.RETRY_CONFIG,
url: this.url,
method: 'POST',
headers,
Expand Down Expand Up @@ -248,12 +249,12 @@ export class ExternalAccountAuthorizedUserClient extends AuthClient {
* Authenticates the provided HTTP request, processes it and resolves with the
* returned response.
* @param opts The HTTP request options.
* @param retry Whether the current attempt is a retry after a failed attempt.
* @param reAuthRetried Whether the current attempt is a retry after a failed attempt due to an auth failure.
* @return A promise that resolves with the successful response.
*/
protected async requestAsync<T>(
opts: GaxiosOptions,
retry = false
reAuthRetried = false
): Promise<GaxiosResponse<T>> {
let response: GaxiosResponse;
try {
Expand All @@ -279,7 +280,7 @@ export class ExternalAccountAuthorizedUserClient extends AuthClient {
const isReadableStream = res.config.data instanceof stream.Readable;
const isAuthErr = statusCode === 401 || statusCode === 403;
if (
!retry &&
!reAuthRetried &&
isAuthErr &&
!isReadableStream &&
this.forceRefreshOnFailure
Expand Down
4 changes: 4 additions & 0 deletions src/auth/googleauth.ts
Original file line number Diff line number Diff line change
Expand Up @@ -1130,6 +1130,10 @@ export class GoogleAuth<T extends AuthClient = JSONClient> {
data: {
payload: crypto.encodeBase64StringUtf8(data),
},
retry: true,
retryConfig: {
httpMethodsToRetry: ['POST'],
},
});

return res.data.signedBlob;
Expand Down
1 change: 1 addition & 0 deletions src/auth/identitypoolclient.ts
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,7 @@ export class IdentityPoolClient extends BaseExternalAccountClient {
method: 'GET',
headers,
responseType: formatType,
retry: true,
danielbankhead marked this conversation as resolved.
Show resolved Hide resolved
};
let subjectToken: string | undefined;
if (formatType === 'text') {
Expand Down
8 changes: 4 additions & 4 deletions src/auth/impersonated.ts
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,7 @@ export class Impersonated extends OAuth2Client implements IdTokenProvider {
payload: Buffer.from(blobToSign).toString('base64'),
};
const res = await this.sourceClient.request<SignBlobResponse>({
...Impersonated.RETRY_CONFIG,
url: u,
data: body,
method: 'POST',
Expand All @@ -157,11 +158,8 @@ export class Impersonated extends OAuth2Client implements IdTokenProvider {

/**
* Refreshes the access token.
* @param refreshToken Unused parameter
*/
protected async refreshToken(
refreshToken?: string | null
Copy link
Contributor

Choose a reason for hiding this comment

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

Is this a breaking change?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Technically, yes, however since its introduction a few years ago I think this would cause more confusion for customers rather than cleaning it up:

I included this change in the current PR as it was causing a lint issue.

): Promise<GetTokenResponse> {
protected async refreshToken(): Promise<GetTokenResponse> {
try {
await this.sourceClient.getAccessToken();
const name = 'projects/-/serviceAccounts/' + this.targetPrincipal;
Expand All @@ -172,6 +170,7 @@ export class Impersonated extends OAuth2Client implements IdTokenProvider {
lifetime: this.lifetime + 's',
};
const res = await this.sourceClient.request<TokenResponse>({
...Impersonated.RETRY_CONFIG,
url: u,
data: body,
method: 'POST',
Expand Down Expand Up @@ -227,6 +226,7 @@ export class Impersonated extends OAuth2Client implements IdTokenProvider {
includeEmail: options?.includeEmail ?? true,
};
const res = await this.sourceClient.request<FetchIdTokenResponse>({
...Impersonated.RETRY_CONFIG,
url: u,
data: body,
method: 'POST',
Expand Down
25 changes: 20 additions & 5 deletions src/auth/oauth2client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -669,6 +669,7 @@ export class OAuth2Client extends AuthClient {
code_verifier: options.codeVerifier,
};
const res = await this.transporter.request<CredentialRequest>({
...OAuth2Client.RETRY_CONFIG,
method: 'POST',
url,
data: querystring.stringify(values),
Expand Down Expand Up @@ -733,6 +734,7 @@ export class OAuth2Client extends AuthClient {
try {
// request for new token
res = await this.transporter.request<CredentialRequest>({
...OAuth2Client.RETRY_CONFIG,
method: 'POST',
url,
data: querystring.stringify(data),
Expand Down Expand Up @@ -956,6 +958,7 @@ export class OAuth2Client extends AuthClient {
callback?: BodyResponseCallback<RevokeCredentialsResult>
): GaxiosPromise<RevokeCredentialsResult> | void {
const opts: GaxiosOptions = {
...OAuth2Client.RETRY_CONFIG,
url: this.getRevokeTokenURL(token).toString(),
method: 'POST',
};
Expand Down Expand Up @@ -1024,7 +1027,7 @@ export class OAuth2Client extends AuthClient {

protected async requestAsync<T>(
opts: GaxiosOptions,
retry = false
reAuthRetried = false
): Promise<GaxiosResponse<T>> {
let r2: GaxiosResponse;
try {
Expand Down Expand Up @@ -1078,11 +1081,16 @@ export class OAuth2Client extends AuthClient {
this.refreshHandler;
const isReadableStream = res.config.data instanceof stream.Readable;
const isAuthErr = statusCode === 401 || statusCode === 403;
if (!retry && isAuthErr && !isReadableStream && mayRequireRefresh) {
if (
!reAuthRetried &&
isAuthErr &&
!isReadableStream &&
mayRequireRefresh
) {
await this.refreshAccessTokenAsync();
return this.requestAsync<T>(opts, true);
} else if (
!retry &&
!reAuthRetried &&
isAuthErr &&
!isReadableStream &&
mayRequireRefreshWithNoRefreshToken
Expand Down Expand Up @@ -1157,6 +1165,7 @@ export class OAuth2Client extends AuthClient {
*/
async getTokenInfo(accessToken: string): Promise<TokenInfo> {
const {data} = await this.transporter.request<TokenInfoRequest>({
...OAuth2Client.RETRY_CONFIG,
method: 'POST',
headers: {
'Content-Type': 'application/x-www-form-urlencoded',
Expand Down Expand Up @@ -1222,7 +1231,10 @@ export class OAuth2Client extends AuthClient {
throw new Error(`Unsupported certificate format ${format}`);
}
try {
res = await this.transporter.request({url});
res = await this.transporter.request({
...OAuth2Client.RETRY_CONFIG,
url,
});
} catch (e) {
if (e instanceof Error) {
e.message = `Failed to retrieve verification certificates: ${e.message}`;
Expand Down Expand Up @@ -1290,7 +1302,10 @@ export class OAuth2Client extends AuthClient {
const url = this.endpoints.oauth2IapPublicKeyUrl.toString();

try {
res = await this.transporter.request({url});
res = await this.transporter.request({
...OAuth2Client.RETRY_CONFIG,
url,
});
} catch (e) {
if (e instanceof Error) {
e.message = `Failed to retrieve verification certificates: ${e.message}`;
Expand Down
18 changes: 18 additions & 0 deletions src/auth/oauth2common.ts
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,24 @@ export abstract class OAuthClientAuthHandler {
}
}
}

/**
* Retry config for Auth-related requests.
*
* @remarks
*
* This is not a part of the default {@link AuthClient.transporter transporter/gaxios}
* config as some downstream APIs would prefer if customers explicitly enable retries,
* such as GCS.
*/
protected static get RETRY_CONFIG(): GaxiosOptions {
return {
retry: true,
retryConfig: {
httpMethodsToRetry: ['GET', 'PUT', 'POST', 'HEAD', 'OPTIONS', 'DELETE'],
},
};
}
}

/**
Expand Down
1 change: 1 addition & 0 deletions src/auth/stscredentials.ts
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,7 @@ export class StsCredentials extends OAuthClientAuthHandler {
Object.assign(headers, additionalHeaders || {});

const opts: GaxiosOptions = {
...StsCredentials.RETRY_CONFIG,
url: this.tokenExchangeEndpoint.toString(),
method: 'POST',
headers,
Expand Down
Loading
Loading