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.

Sign up for GitHub

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

Jump to bottom

feat: Enable Retries For Auth Requests #1791

Merged
merged 5 commits into from
Apr 12, 2024
Merged

feat: Enable Retries For Auth Requests #1791

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
0c1e767
feat: Enable Retries For Auth Requests
d-goog Apr 10, 2024
c2e20dc
chore: lint
d-goog Apr 10, 2024
8740939
docs: `reAuthRetried` clarification
d-goog Apr 10, 2024
5a6308f
chore: lint
d-goog Apr 10, 2024
409da97
refactor: Use `IdentityPoolClient.RETRY_CONFIG`
d-goog Apr 11, 2024
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
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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.
* @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
View file
Open in desktop
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.
* @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
View file
Open in desktop
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.
* @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
View file
Open in desktop
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
View file
Open in desktop
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
3 changes: 3 additions & 0 deletions src/auth/impersonated.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,7 @@
payload: Buffer.from(blobToSign).toString('base64'),
};
const res = await this.sourceClient.request<SignBlobResponse>({
...Impersonated.RETRY_CONFIG,
url: u,
data: body,
method: 'POST',
Expand All @@ -160,7 +161,7 @@
* @param refreshToken Unused parameter
*/
protected async refreshToken(
refreshToken?: string | null

Check warning on line 164 in src/auth/impersonated.ts

View workflow job for this annotation

GitHub Actions / lint 

'refreshToken' is defined but never used
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> {
try {
await this.sourceClient.getAccessToken();
Expand All @@ -172,6 +173,7 @@
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 +229,7 @@
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
View file
Open in desktop
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
View file
Open in desktop
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
View file
Open in desktop
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
20 changes: 10 additions & 10 deletions test/test.awsclient.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -401,12 +401,12 @@ describe('AwsClient', () => {
// Simulate error during region retrieval.
const scope = nock(metadataBaseUrl)
.get('/latest/meta-data/placement/availability-zone')
.reply(500);
.reply(400);
Copy link
Contributor

Choose a reason for hiding this comment

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

Why were the status codes here all changed?

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 500s are retryable and tests would fail with an unexpected error. As the tests want a error on the first call, I changed it to 400.


const client = new AwsClient(awsOptions);

await assert.rejects(client.retrieveSubjectToken(), {
status: 500,
status: 400,
});
scope.done();
});
Expand Down Expand Up @@ -435,12 +435,12 @@ describe('AwsClient', () => {
.get('/latest/meta-data/iam/security-credentials')
.reply(200, awsRole)
.get(`/latest/meta-data/iam/security-credentials/${awsRole}`)
.reply(408);
.reply(404);

const client = new AwsClient(awsOptions);

await assert.rejects(client.retrieveSubjectToken(), {
status: 408,
status: 404,
});
scope.done();
});
Expand Down Expand Up @@ -602,12 +602,12 @@ describe('AwsClient', () => {
.get('/latest/meta-data/placement/availability-zone')
.reply(200, `${awsRegion}b`)
.get('/latest/meta-data/iam/security-credentials')
.reply(500);
.reply(400);

const client = new AwsClient(awsOptions);

await assert.rejects(client.getAccessToken(), {
status: 500,
status: 400,
});
scope.done();
});
Expand Down Expand Up @@ -704,12 +704,12 @@ describe('AwsClient', () => {
// Simulate error during region retrieval.
const scope = nock(metadataBaseUrl)
.get('/latest/meta-data/placement/availability-zone')
.reply(500);
.reply(400);

const client = new AwsClient(awsOptions);

await assert.rejects(client.retrieveSubjectToken(), {
status: 500,
status: 400,
});
scope.done();
});
Expand Down Expand Up @@ -982,12 +982,12 @@ describe('AwsClient', () => {

const scope = nock(metadataBaseUrl)
.get('/latest/meta-data/placement/availability-zone')
.reply(500);
.reply(400);

const client = new AwsClient(awsOptions);

await assert.rejects(client.getAccessToken(), {
status: 500,
status: 400,
});
scope.done();
});
Expand Down
Loading
Loading