Skip to content

Commit

Permalink
chore: Refactor to use MockIAMCredentialsServiceTransportFactory
Browse files Browse the repository at this point in the history
  • Loading branch information
lqiu96 committed Jul 11, 2024
1 parent d2b7dc7 commit 0056961
Show file tree
Hide file tree
Showing 7 changed files with 82 additions and 85 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,6 @@
import com.google.auth.http.HttpTransportFactory;
import com.google.auth.oauth2.ExternalAccountAuthorizedUserCredentialsTest.MockExternalAccountAuthorizedUserCredentialsTransportFactory;
import com.google.auth.oauth2.IdentityPoolCredentialsTest.MockExternalAccountCredentialsTransportFactory;
import com.google.auth.oauth2.ImpersonatedCredentialsTest.MockIAMCredentialsServiceTransportFactory;
import com.google.common.collect.ImmutableList;
import java.io.ByteArrayInputStream;
import java.io.IOException;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,6 @@
import static org.junit.Assert.fail;

import com.google.api.client.http.HttpStatusCodes;
import com.google.api.client.http.HttpTransport;
import com.google.api.client.json.GenericJson;
import com.google.api.client.json.JsonFactory;
import com.google.api.client.json.JsonGenerator;
Expand All @@ -51,7 +50,6 @@
import com.google.api.client.util.Clock;
import com.google.auth.ServiceAccountSigner.SigningException;
import com.google.auth.TestUtils;
import com.google.auth.http.HttpTransportFactory;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import java.io.ByteArrayOutputStream;
Expand Down Expand Up @@ -140,16 +138,6 @@ public class ImpersonatedCredentialsTest extends BaseSerializationTest {
public static final List<String> DELEGATES =
Arrays.asList("[email protected]", "[email protected]");

static class MockIAMCredentialsServiceTransportFactory implements HttpTransportFactory {

MockIAMCredentialsServiceTransport transport = new MockIAMCredentialsServiceTransport();

@Override
public HttpTransport create() {
return transport;
}
}

private GoogleCredentials sourceCredentials;
private MockIAMCredentialsServiceTransportFactory mockTransportFactory;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -47,11 +47,11 @@
public class MockIAMCredentialsServiceTransport extends MockHttpTransport {

private static final String DEFAULT_IAM_ACCESS_TOKEN_ENDPOINT =
"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:generateAccessToken";
"https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:generateAccessToken";
private static final String IAM_ID_TOKEN_ENDPOINT =
"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:generateIdToken";
"https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:generateIdToken";
private static final String IAM_SIGN_ENDPOINT =
"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/%s:signBlob";
"https://iamcredentials.%s/v1/projects/-/serviceAccounts/%s:signBlob";
private Integer tokenResponseErrorCode;
private String tokenResponseErrorContent;
private String targetPrincipal;
Expand All @@ -65,9 +65,13 @@ public class MockIAMCredentialsServiceTransport extends MockHttpTransport {

private String idToken;

private String universeDomain;

private MockLowLevelHttpRequest request;

public MockIAMCredentialsServiceTransport() {}
MockIAMCredentialsServiceTransport(String universeDomain) {
this.universeDomain = universeDomain;
}

public void setTokenResponseErrorCode(Integer tokenResponseErrorCode) {
this.tokenResponseErrorCode = tokenResponseErrorCode;
Expand Down Expand Up @@ -112,13 +116,15 @@ public MockLowLevelHttpRequest getRequest() {

@Override
public LowLevelHttpRequest buildRequest(String method, String url) throws IOException {

String iamAccesssTokenformattedUrl =
iamAccessTokenEndpoint != null
? iamAccessTokenEndpoint
: String.format(DEFAULT_IAM_ACCESS_TOKEN_ENDPOINT, this.targetPrincipal);
String iamSignBlobformattedUrl = String.format(IAM_SIGN_ENDPOINT, this.targetPrincipal);
String iamIdTokenformattedUrl = String.format(IAM_ID_TOKEN_ENDPOINT, this.targetPrincipal);
: String.format(
DEFAULT_IAM_ACCESS_TOKEN_ENDPOINT, universeDomain, this.targetPrincipal);
String iamSignBlobformattedUrl =
String.format(IAM_SIGN_ENDPOINT, universeDomain, this.targetPrincipal);
String iamIdTokenformattedUrl =
String.format(IAM_ID_TOKEN_ENDPOINT, universeDomain, this.targetPrincipal);
if (url.equals(iamAccesssTokenformattedUrl)) {
this.request =
new MockLowLevelHttpRequest(url) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,52 @@
/*
* 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 com.google.api.client.http.HttpTransport;
import com.google.auth.Credentials;
import com.google.auth.http.HttpTransportFactory;

public class MockIAMCredentialsServiceTransportFactory implements HttpTransportFactory {
MockIAMCredentialsServiceTransport transport;

MockIAMCredentialsServiceTransportFactory() {
this(Credentials.GOOGLE_DEFAULT_UNIVERSE);
}

MockIAMCredentialsServiceTransportFactory(String universeDomain) {
this.transport = new MockIAMCredentialsServiceTransport(universeDomain);
}

@Override
public HttpTransport create() {
return transport;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -52,8 +52,6 @@
import java.util.Queue;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/** Mock transport to simulate providing Google OAuth2 access tokens */
public class MockTokenServerTransport extends MockHttpTransport {
Expand All @@ -69,22 +67,12 @@ public class MockTokenServerTransport extends MockHttpTransport {
final Map<String, String> codes = new HashMap<String, String>();
final Map<String, Map<String, String>> additionalParameters =
new HashMap<String, Map<String, String>>();
private URI tokenServerUri;
URI tokenServerUri = OAuth2Utils.TOKEN_SERVER_URI;
private IOException error;
private final Queue<Future<LowLevelHttpResponse>> responseSequence = new ArrayDeque<>();
private int expiresInSeconds = 3600;

private Pattern iamEndpointRegex =
Pattern.compile(
"https://iamcredentials.(.*).com/v1/projects/-/serviceAccounts/(.*):generateIdToken");

public MockTokenServerTransport() {
this(OAuth2Utils.TOKEN_SERVER_URI);
}

MockTokenServerTransport(URI tokenServerUri) {
this.tokenServerUri = tokenServerUri;
}
public MockTokenServerTransport() {}

public URI getTokenServerUri() {
return tokenServerUri;
Expand Down Expand Up @@ -186,9 +174,6 @@ public LowLevelHttpResponse execute() throws IOException {
};
}

Matcher matcher = iamEndpointRegex.matcher(tokenServerUri.toString());
boolean useIamEndpoint = matcher.matches();

if (urlWithoutQuery.equals(tokenServerUri.toString())) {
return new MockLowLevelHttpRequest(url) {
@Override
Expand Down Expand Up @@ -305,18 +290,6 @@ public LowLevelHttpResponse execute() throws IOException {
} else {
throw new IOException("Service Account Email not found as issuer.");
}
} else if (useIamEndpoint) {
// Group 1 is the universe domain and group 2 is the client email
String universeDomain = matcher.group(1);
if (universeDomain.equals("googleapis.com")) {
throw new IOException(
"IAM Token Endpoint is in GDU. Iam Token Endpoint flow should not be invoked");
}
String clientEmail = matcher.group(2);
if (!serviceAccounts.containsKey(clientEmail)) {
throw new IOException("Client Email " + clientEmail + " does not exist");
}
generateAccessToken = false;
} else {
throw new IOException("Unknown token type.");
}
Expand All @@ -337,11 +310,7 @@ public LowLevelHttpResponse execute() throws IOException {
}
}
if (isUserEmailScope || !generateAccessToken) {
String fieldName = "id_token";
if (useIamEndpoint) {
fieldName = "token";
}
responseContents.put(fieldName, ServiceAccountCredentialsTest.DEFAULT_ID_TOKEN);
responseContents.put("id_token", ServiceAccountCredentialsTest.DEFAULT_ID_TOKEN);
}
String refreshText = responseContents.toPrettyString();

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,7 @@

public class MockTokenServerTransportFactory implements HttpTransportFactory {

public MockTokenServerTransport transport;

public MockTokenServerTransportFactory() {
this(new MockTokenServerTransport());
}

MockTokenServerTransportFactory(MockTokenServerTransport transport) {
this.transport = transport;
}
MockTokenServerTransport transport = new MockTokenServerTransport();

@Override
public HttpTransport create() {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -123,7 +123,6 @@ public class ServiceAccountCredentialsTest extends BaseSerializationTest {
private static final int DEFAULT_LIFETIME_IN_SECONDS = 3600;
private static final int INVALID_LIFETIME = 43210;
private static final String JWT_ACCESS_PREFIX = "Bearer ";
private static final String GOOGLE_DEFAULT_UNIVERSE = "googleapis.com";

private ServiceAccountCredentials.Builder createDefaultBuilderWithToken(String accessToken)
throws IOException {
Expand Down Expand Up @@ -210,7 +209,7 @@ public void createdScoped_clones() throws IOException {
assertArrayEquals(newScopes.toArray(), newCredentials.getScopes().toArray());
assertEquals(USER, newCredentials.getServiceAccountUser());
assertEquals(PROJECT_ID, newCredentials.getProjectId());
assertEquals(GOOGLE_DEFAULT_UNIVERSE, newCredentials.getUniverseDomain());
assertEquals(Credentials.GOOGLE_DEFAULT_UNIVERSE, newCredentials.getUniverseDomain());

assertArrayEquals(
SCOPES.toArray(), ((ServiceAccountCredentials) credentials).getScopes().toArray());
Expand Down Expand Up @@ -380,7 +379,7 @@ public void createdScoped_withAud_noUniverse_jwtWithScopesDisabled_accessToken()

GoogleCredentials scopedCredentials = credentials.createScoped(SCOPES);
assertEquals(false, credentials.isExplicitUniverseDomain());
assertEquals(GOOGLE_DEFAULT_UNIVERSE, credentials.getUniverseDomain());
assertEquals(Credentials.GOOGLE_DEFAULT_UNIVERSE, credentials.getUniverseDomain());
Map<String, List<String>> metadata = scopedCredentials.getRequestMetadata(CALL_URI);
TestUtils.assertContainsBearerToken(metadata, ACCESS_TOKEN);
}
Expand Down Expand Up @@ -515,7 +514,7 @@ public void fromJSON_getProjectId() throws IOException {
ServiceAccountCredentials credentials =
ServiceAccountCredentials.fromJson(json, new MockTokenServerTransportFactory());
assertEquals(PROJECT_ID, credentials.getProjectId());
assertEquals(GOOGLE_DEFAULT_UNIVERSE, credentials.getUniverseDomain());
assertEquals(Credentials.GOOGLE_DEFAULT_UNIVERSE, credentials.getUniverseDomain());
}

@Test
Expand Down Expand Up @@ -620,7 +619,7 @@ public void getRequestMetadata_customTokenServer_hasAccessToken() throws IOExcep
@Test
public void getUniverseDomain_defaultUniverse() throws IOException {
ServiceAccountCredentials credentials = createDefaultBuilder().build();
assertEquals(GOOGLE_DEFAULT_UNIVERSE, credentials.getUniverseDomain());
assertEquals(Credentials.GOOGLE_DEFAULT_UNIVERSE, credentials.getUniverseDomain());
}

@Test
Expand Down Expand Up @@ -901,21 +900,17 @@ public void idTokenWithAudience_oauthFlow_incorrect() throws IOException {
@Test
public void idTokenWithAudience_iamFlow_correct() throws IOException {
String nonGDU = "test.com";
String accessToken1 = "1/MkSJoj1xsli0AccessToken_NKPY2";
MockTokenServerTransport transport =
new MockTokenServerTransport(
URI.create(String.format(OAuth2Utils.IAM_ID_TOKEN_URI_FORMAT, nonGDU, CLIENT_EMAIL)));
MockTokenServerTransportFactory transportFactory =
new MockTokenServerTransportFactory(transport);
MockIAMCredentialsServiceTransportFactory transportFactory =
new MockIAMCredentialsServiceTransportFactory(nonGDU);
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
transportFactory.transport.setIdToken(DEFAULT_ID_TOKEN);
ServiceAccountCredentials credentials =
createDefaultBuilder()
.setScopes(SCOPES)
.setHttpTransportFactory(transportFactory)
.setUniverseDomain(nonGDU)
.build();

transport.addServiceAccount(CLIENT_EMAIL, accessToken1);

String targetAudience = "https://foo.bar";
IdTokenCredentials tokenCredential =
IdTokenCredentials.newBuilder()
Expand All @@ -935,21 +930,17 @@ public void idTokenWithAudience_iamFlow_correct() throws IOException {
@Test
public void idTokenWithAudience_iamFlow_incorrect() throws IOException {
String nonGDU = "test.com";
String accessToken1 = "1/MkSJoj1xsli0AccessToken_NKPY2";
MockTokenServerTransport transport =
new MockTokenServerTransport(
URI.create(String.format(OAuth2Utils.IAM_ID_TOKEN_URI_FORMAT, nonGDU, CLIENT_EMAIL)));
MockTokenServerTransportFactory transportFactory =
new MockTokenServerTransportFactory(transport);
MockIAMCredentialsServiceTransportFactory transportFactory =
new MockIAMCredentialsServiceTransportFactory(nonGDU);
transportFactory.transport.setTargetPrincipal(CLIENT_EMAIL);
transportFactory.transport.setIdToken(DEFAULT_ID_TOKEN);
ServiceAccountCredentials credentials =
createDefaultBuilder()
.setScopes(SCOPES)
.setHttpTransportFactory(transportFactory)
.setUniverseDomain(nonGDU)
.build();

transport.addServiceAccount(CLIENT_EMAIL, accessToken1);

String targetAudience = "https://bar";
IdTokenCredentials tokenCredential =
IdTokenCredentials.newBuilder()
Expand Down

0 comments on commit 0056961

Please sign in to comment.