Skip to content

Commit 73a0301

Browse files
macastelazmacastelaz
committed
feat(oauth2): implement IAM impersonation mTLS transport pinning and 401 recovery
1 parent e6a4797 commit 73a0301

9 files changed

Lines changed: 500 additions & 29 deletions

File tree

‎google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/AwsCredentials.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,11 @@ public class AwsCredentials extends ExternalAccountCredentials {
119119

120120
@Override
121121
public AccessToken refreshAccessToken() throws IOException {
122+
return refreshAccessToken(this.transportFactory);
123+
}
124+
125+
@Override
126+
public AccessToken refreshAccessToken(HttpTransportFactory transportFactory) throws IOException {
122127
StsTokenExchangeRequest.Builder stsTokenExchangeRequest =
123128
StsTokenExchangeRequest.newBuilder(retrieveSubjectToken(), getSubjectTokenType())
124129
.setAudience(getAudience());
@@ -129,7 +134,8 @@ public AccessToken refreshAccessToken() throws IOException {
129134
stsTokenExchangeRequest.setScopes(new ArrayList<>(scopes));
130135
}
131136

132-
return exchangeExternalCredentialForAccessToken(stsTokenExchangeRequest.build());
137+
return exchangeExternalCredentialForAccessToken(
138+
stsTokenExchangeRequest.build(), transportFactory);
133139
}
134140

135141
@Override

‎google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/ExternalAccountCredentials.java‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -525,6 +525,19 @@ private boolean shouldBuildImpersonatedCredential() {
525525
return this.serviceAccountImpersonationUrl != null && this.impersonatedCredentials == null;
526526
}
527527

528+
/**
529+
* Refreshes the access token using the specified transport factory. Default implementation
530+
* delegates to {@link #refreshAccessToken()}. Subclasses should override this method if they
531+
* support transport pinning per refresh cycle.
532+
*
533+
* @param transportFactory the HTTP transport factory to use for this refresh cycle
534+
* @return the refreshed access token
535+
* @throws IOException if the token refresh fails
536+
*/
537+
public AccessToken refreshAccessToken(HttpTransportFactory transportFactory) throws IOException {
538+
return refreshAccessToken();
539+
}
540+
528541
/**
529542
* Exchanges the external credential for a Google Cloud access token.
530543
*
@@ -555,7 +568,7 @@ protected AccessToken exchangeExternalCredentialForAccessToken(
555568
this.impersonatedCredentials = this.buildImpersonatedCredentials();
556569
}
557570
if (this.impersonatedCredentials != null) {
558-
return this.impersonatedCredentials.refreshAccessToken();
571+
return this.impersonatedCredentials.refreshAccessToken(cycleTransportFactory);
559572
}
560573

561574
StsRequestHandler.Builder requestHandler =

‎google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/IdentityPoolCredentials.java‎

Lines changed: 25 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -230,7 +230,17 @@ public AccessToken refreshAccessToken() throws IOException {
230230
KeyStore pinnedKeyStore = this.x509Provider.getKeyStore();
231231
cycleTransportFactory = new MtlsHttpTransportFactory(pinnedKeyStore);
232232
}
233+
return refreshWithRetry(cycleTransportFactory, true);
234+
}
235+
236+
@Override
237+
public AccessToken refreshAccessToken(HttpTransportFactory cycleTransportFactory)
238+
throws IOException {
239+
return refreshWithRetry(cycleTransportFactory, false);
240+
}
233241

242+
private AccessToken refreshWithRetry(
243+
HttpTransportFactory cycleTransportFactory, boolean allowRetry) throws IOException {
234244
// Read subject and actor tokens, atomically if from the same file supplier.
235245
String subjectToken;
236246
String actorToken = null;
@@ -264,20 +274,30 @@ public AccessToken refreshAccessToken() throws IOException {
264274
try {
265275
return exchangeExternalCredentialForAccessToken(
266276
stsTokenExchangeRequest.build(), cycleTransportFactory);
267-
} catch (OAuthException e) {
268-
if (e.getHttpStatusCode() == 401 && this.x509Provider != null) {
277+
} catch (Exception e) {
278+
if (allowRetry && OAuth2Utils.isUnauthorizedException(e) && this.x509Provider != null) {
269279
try {
270280
// On 401, re-read from X509Provider for fresh certs and retry once.
271281
KeyStore freshKeyStore = this.x509Provider.getKeyStore();
272282
HttpTransportFactory retryTransportFactory = new MtlsHttpTransportFactory(freshKeyStore);
273-
return exchangeExternalCredentialForAccessToken(
274-
stsTokenExchangeRequest.build(), retryTransportFactory);
283+
return refreshWithRetry(retryTransportFactory, false);
275284
} catch (IOException retryException) {
276285
retryException.addSuppressed(e);
277286
throw retryException;
287+
} catch (Exception retryException) {
288+
IOException ioException =
289+
new IOException("Failed to reload certificate on retry", retryException);
290+
ioException.addSuppressed(e);
291+
throw ioException;
278292
}
279293
}
280-
throw e;
294+
if (e instanceof IOException) {
295+
throw (IOException) e;
296+
}
297+
if (e instanceof RuntimeException) {
298+
throw (RuntimeException) e;
299+
}
300+
throw new IOException(e);
281301
}
282302
}
283303

‎google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/ImpersonatedCredentials.java‎

Lines changed: 38 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import com.google.api.client.util.GenericData;
4747
import com.google.api.core.ObsoleteApi;
4848
import com.google.auth.CredentialTypeForMetrics;
49+
import com.google.auth.Credentials;
4950
import com.google.auth.ServiceAccountSigner;
5051
import com.google.auth.http.HttpCredentialsAdapter;
5152
import com.google.auth.http.HttpTransportFactory;
@@ -531,7 +532,7 @@ public ImpersonatedCredentials createWithCustomCalendar(Calendar calendar) {
531532
.build();
532533
}
533534

534-
private ImpersonatedCredentials(Builder builder) throws IOException {
535+
protected ImpersonatedCredentials(Builder builder) throws IOException {
535536
super(builder);
536537
this.sourceCredentials = builder.getSourceCredentials();
537538
this.targetPrincipal = builder.getTargetPrincipal();
@@ -580,31 +581,47 @@ public String getUniverseDomain() throws IOException {
580581

581582
@Override
582583
public AccessToken refreshAccessToken() throws IOException {
583-
if (this.sourceCredentials.getAccessToken() == null) {
584-
// Apply the `CLOUD_PLATFORM_SCOPE` to access the iamcredentials endpoint
585-
this.sourceCredentials =
586-
this.sourceCredentials.createScoped(
587-
Collections.singletonList(OAuth2Utils.CLOUD_PLATFORM_SCOPE));
588-
}
589-
590-
// skip for SA with SSJ flow because it uses self-signed JWT
591-
// and will get refreshed at initialize request step
592-
// run for other source credential types or SA with GDU assert flow
593-
if (!(this.sourceCredentials instanceof ServiceAccountCredentials)
594-
|| (isDefaultUniverseDomain()
595-
&& ((ServiceAccountCredentials) this.sourceCredentials)
596-
.shouldUseAssertionFlowForGdu())) {
597-
try {
598-
this.sourceCredentials.refreshIfExpired();
599-
} catch (IOException e) {
600-
throw new IOException("Unable to refresh sourceCredentials", e);
584+
return refreshAccessToken(this.transportFactory);
585+
}
586+
587+
AccessToken refreshAccessToken(HttpTransportFactory transportFactory) throws IOException {
588+
HttpCredentialsAdapter adapter;
589+
if (this.sourceCredentials instanceof ExternalAccountCredentials) {
590+
AccessToken intermediateAccessToken =
591+
((ExternalAccountCredentials) this.sourceCredentials)
592+
.refreshAccessToken(transportFactory);
593+
Credentials authCredentials =
594+
intermediateAccessToken != null
595+
? OAuth2Credentials.create(intermediateAccessToken)
596+
: this.sourceCredentials;
597+
adapter = new HttpCredentialsAdapter(authCredentials);
598+
} else {
599+
if (this.sourceCredentials.getAccessToken() == null) {
600+
// Apply the `CLOUD_PLATFORM_SCOPE` to access the iamcredentials endpoint
601+
this.sourceCredentials =
602+
this.sourceCredentials.createScoped(
603+
Collections.singletonList(OAuth2Utils.CLOUD_PLATFORM_SCOPE));
604+
}
605+
606+
// skip for SA with SSJ flow because it uses self-signed JWT
607+
// and will get refreshed at initialize request step
608+
// run for other source credential types or SA with GDU assert flow
609+
if (!(this.sourceCredentials instanceof ServiceAccountCredentials)
610+
|| (isDefaultUniverseDomain()
611+
&& ((ServiceAccountCredentials) this.sourceCredentials)
612+
.shouldUseAssertionFlowForGdu())) {
613+
try {
614+
this.sourceCredentials.refreshIfExpired();
615+
} catch (IOException e) {
616+
throw new IOException("Unable to refresh sourceCredentials", e);
617+
}
601618
}
619+
adapter = new HttpCredentialsAdapter(sourceCredentials);
602620
}
603621

604-
HttpTransport httpTransport = this.transportFactory.create();
622+
HttpTransport httpTransport = transportFactory.create();
605623
JsonObjectParser parser = new JsonObjectParser(OAuth2Utils.JSON_FACTORY);
606624

607-
HttpCredentialsAdapter adapter = new HttpCredentialsAdapter(sourceCredentials);
608625
HttpRequestFactory requestFactory = httpTransport.createRequestFactory();
609626

610627
String endpointUrl =

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

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
package com.google.auth.oauth2;
3333

3434
import com.google.api.client.http.HttpHeaders;
35+
import com.google.api.client.http.HttpResponseException;
3536
import com.google.api.client.http.HttpTransport;
3637
import com.google.api.client.http.javanet.NetHttpTransport;
3738
import com.google.api.client.json.GenericJson;
@@ -324,5 +325,28 @@ static String generateBasicAuthHeader(String username, String password) {
324325
return "Basic " + encodedCredentials;
325326
}
326327

328+
/**
329+
* Returns whether the given throwable or any exception in its causal chain represents a 401
330+
* Unauthorized error (either an {@link OAuthException} or {@link HttpResponseException} with
331+
* status code 401).
332+
*/
333+
static boolean isUnauthorizedException(@Nullable Throwable t) {
334+
while (t != null) {
335+
if (t instanceof OAuthException && ((OAuthException) t).getHttpStatusCode() == 401) {
336+
return true;
337+
}
338+
if (t instanceof HttpResponseException
339+
&& ((HttpResponseException) t).getStatusCode() == 401) {
340+
return true;
341+
}
342+
Throwable cause = t.getCause();
343+
if (cause == t) {
344+
break;
345+
}
346+
t = cause;
347+
}
348+
return false;
349+
}
350+
327351
private OAuth2Utils() {}
328352
}

‎google-auth-library-java/oauth2_http/java/com/google/auth/oauth2/PluggableAuthCredentials.java‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,11 @@ public class PluggableAuthCredentials extends ExternalAccountCredentials {
121121

122122
@Override
123123
public AccessToken refreshAccessToken() throws IOException {
124+
return refreshAccessToken(this.transportFactory);
125+
}
126+
127+
@Override
128+
public AccessToken refreshAccessToken(HttpTransportFactory transportFactory) throws IOException {
124129
String credential = retrieveSubjectToken();
125130
StsTokenExchangeRequest.Builder stsTokenExchangeRequest =
126131
StsTokenExchangeRequest.newBuilder(credential, getSubjectTokenType())
@@ -130,7 +135,8 @@ public AccessToken refreshAccessToken() throws IOException {
130135
if (scopes != null && !scopes.isEmpty()) {
131136
stsTokenExchangeRequest.setScopes(new ArrayList<>(scopes));
132137
}
133-
return exchangeExternalCredentialForAccessToken(stsTokenExchangeRequest.build());
138+
return exchangeExternalCredentialForAccessToken(
139+
stsTokenExchangeRequest.build(), transportFactory);
134140
}
135141

136142
/**

0 commit comments

Comments
 (0)