Skip to content

Commit 0a8412f

Browse files
chingor13chingor13
authored
fix: quotaProjectId should be applied for cached getRequestMetadata(URI, Executor, RequestMetadataCallback) (#509)
* test: add failing test for request metadata with callback * refactor: use protected getAdditionalHeaders to provide quotaProjectId and other headers * docs: add javadoc for new protected getAdditionalHeaders() * test: add more tests for quotaProjectId * docs: fix nit javadoc return format
1 parent 3b50627 commit 0a8412f

6 files changed

Lines changed: 276 additions & 12 deletions

File tree

‎oauth2_http/java/com/google/auth/oauth2/OAuth2Credentials.java‎

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import com.google.common.annotations.VisibleForTesting;
3939
import com.google.common.base.MoreObjects;
4040
import com.google.common.base.Preconditions;
41+
import com.google.common.collect.ImmutableMap;
4142
import com.google.common.collect.Iterables;
4243
import java.io.IOException;
4344
import java.io.ObjectInputStream;
@@ -56,6 +57,7 @@ public class OAuth2Credentials extends Credentials {
5657

5758
private static final long serialVersionUID = 4556936364828217687L;
5859
private static final long MINIMUM_TOKEN_MILLISECONDS = 60000L * 5L;
60+
private static final Map<String, List<String>> EMPTY_EXTRA_HEADERS = Collections.emptyMap();
5961

6062
// byte[] is serializable, so the lock variable can be final
6163
private final Object lock = new byte[0];
@@ -89,7 +91,7 @@ protected OAuth2Credentials() {
8991
*/
9092
protected OAuth2Credentials(AccessToken accessToken) {
9193
if (accessToken != null) {
92-
useAccessToken(accessToken);
94+
useAccessToken(accessToken, EMPTY_EXTRA_HEADERS);
9395
}
9496
}
9597

@@ -154,7 +156,9 @@ public void refresh() throws IOException {
154156
synchronized (lock) {
155157
requestMetadata = null;
156158
temporaryAccess = null;
157-
useAccessToken(Preconditions.checkNotNull(refreshAccessToken(), "new access token"));
159+
useAccessToken(
160+
Preconditions.checkNotNull(refreshAccessToken(), "new access token"),
161+
getAdditionalHeaders());
158162
if (changeListeners != null) {
159163
for (CredentialsChangedListener listener : changeListeners) {
160164
listener.onChanged(this);
@@ -163,6 +167,15 @@ public void refresh() throws IOException {
163167
}
164168
}
165169

170+
/**
171+
* Provide additional headers to return as request metadata.
172+
*
173+
* @return additional headers
174+
*/
175+
protected Map<String, List<String>> getAdditionalHeaders() {
176+
return EMPTY_EXTRA_HEADERS;
177+
}
178+
166179
/**
167180
* Refresh these credentials only if they have expired or are expiring imminently.
168181
*
@@ -177,12 +190,15 @@ public void refreshIfExpired() throws IOException {
177190
}
178191

179192
// Must be called under lock
180-
private void useAccessToken(AccessToken token) {
193+
private void useAccessToken(AccessToken token, Map<String, List<String>> additionalHeaders) {
181194
this.temporaryAccess = token;
182195
this.requestMetadata =
183-
Collections.singletonMap(
184-
AuthHttpConstants.AUTHORIZATION,
185-
Collections.singletonList(OAuth2Utils.BEARER_PREFIX + token.getTokenValue()));
196+
ImmutableMap.<String, List<String>>builder()
197+
.put(
198+
AuthHttpConstants.AUTHORIZATION,
199+
Collections.singletonList(OAuth2Utils.BEARER_PREFIX + token.getTokenValue()))
200+
.putAll(additionalHeaders)
201+
.build();
186202
}
187203

188204
// Must be called under lock

‎oauth2_http/java/com/google/auth/oauth2/ServiceAccountCredentials.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -599,9 +599,12 @@ public JwtCredentials jwtWithClaims(JwtClaims newClaims) {
599599
}
600600

601601
@Override
602-
public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException {
603-
Map<String, List<String>> requestMetadata = super.getRequestMetadata(uri);
604-
return addQuotaProjectIdToRequestMetadata(quotaProjectId, requestMetadata);
602+
protected Map<String, List<String>> getAdditionalHeaders() {
603+
Map<String, List<String>> headers = super.getAdditionalHeaders();
604+
if (quotaProjectId != null) {
605+
return addQuotaProjectIdToRequestMetadata(quotaProjectId, headers);
606+
}
607+
return headers;
605608
}
606609

607610
@Override

‎oauth2_http/java/com/google/auth/oauth2/UserCredentials.java‎

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -277,9 +277,12 @@ public void save(String filePath) throws IOException {
277277
}
278278

279279
@Override
280-
public Map<String, List<String>> getRequestMetadata(URI uri) throws IOException {
281-
Map<String, List<String>> requestMetadata = super.getRequestMetadata(uri);
282-
return addQuotaProjectIdToRequestMetadata(quotaProjectId, requestMetadata);
280+
protected Map<String, List<String>> getAdditionalHeaders() {
281+
Map<String, List<String>> headers = super.getAdditionalHeaders();
282+
if (quotaProjectId != null) {
283+
return addQuotaProjectIdToRequestMetadata(quotaProjectId, headers);
284+
}
285+
return headers;
283286
}
284287

285288
@Override

‎oauth2_http/javatests/com/google/auth/oauth2/ServiceAccountCredentialsTest.java‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import com.google.api.client.testing.http.MockLowLevelHttpResponse;
5050
import com.google.api.client.util.Clock;
5151
import com.google.api.client.util.Joiner;
52+
import com.google.auth.RequestMetadataCallback;
5253
import com.google.auth.TestUtils;
5354
import com.google.auth.http.HttpTransportFactory;
5455
import com.google.auth.oauth2.GoogleCredentialsTest.MockHttpTransportFactory;
@@ -68,6 +69,7 @@
6869
import java.util.Collections;
6970
import java.util.List;
7071
import java.util.Map;
72+
import java.util.concurrent.atomic.AtomicBoolean;
7173
import org.junit.Test;
7274
import org.junit.runner.RunWith;
7375
import org.junit.runners.JUnit4;
@@ -1032,6 +1034,97 @@ public void fromStream_noPrivateKeyId_throws() throws IOException {
10321034
testFromStreamException(serviceAccountStream, "private_key_id");
10331035
}
10341036

1037+
@Test
1038+
public void getRequestMetadataSetsQuotaProjectId() throws IOException {
1039+
MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
1040+
transportFactory.transport.addClient(CLIENT_ID, "unused-client-secret");
1041+
transportFactory.transport.addServiceAccount(CLIENT_EMAIL, ACCESS_TOKEN);
1042+
1043+
PrivateKey privateKey = ServiceAccountCredentials.privateKeyFromPkcs8(PRIVATE_KEY_PKCS8);
1044+
GoogleCredentials credentials =
1045+
ServiceAccountCredentials.newBuilder()
1046+
.setClientId(CLIENT_ID)
1047+
.setClientEmail(CLIENT_EMAIL)
1048+
.setPrivateKey(privateKey)
1049+
.setPrivateKeyId(PRIVATE_KEY_ID)
1050+
.setScopes(SCOPES)
1051+
.setServiceAccountUser(USER)
1052+
.setProjectId(PROJECT_ID)
1053+
.setQuotaProjectId("my-quota-project-id")
1054+
.setHttpTransportFactory(transportFactory)
1055+
.build();
1056+
1057+
Map<String, List<String>> metadata = credentials.getRequestMetadata();
1058+
assertTrue(metadata.containsKey("x-goog-user-project"));
1059+
List<String> headerValues = metadata.get("x-goog-user-project");
1060+
assertEquals(1, headerValues.size());
1061+
assertEquals("my-quota-project-id", headerValues.get(0));
1062+
}
1063+
1064+
@Test
1065+
public void getRequestMetadataNoQuotaProjectId() throws IOException {
1066+
MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
1067+
transportFactory.transport.addClient(CLIENT_ID, "unused-client-secret");
1068+
transportFactory.transport.addServiceAccount(CLIENT_EMAIL, ACCESS_TOKEN);
1069+
1070+
PrivateKey privateKey = ServiceAccountCredentials.privateKeyFromPkcs8(PRIVATE_KEY_PKCS8);
1071+
GoogleCredentials credentials =
1072+
ServiceAccountCredentials.newBuilder()
1073+
.setClientId(CLIENT_ID)
1074+
.setClientEmail(CLIENT_EMAIL)
1075+
.setPrivateKey(privateKey)
1076+
.setPrivateKeyId(PRIVATE_KEY_ID)
1077+
.setScopes(SCOPES)
1078+
.setServiceAccountUser(USER)
1079+
.setProjectId(PROJECT_ID)
1080+
.setHttpTransportFactory(transportFactory)
1081+
.build();
1082+
1083+
Map<String, List<String>> metadata = credentials.getRequestMetadata();
1084+
assertFalse(metadata.containsKey("x-goog-user-project"));
1085+
}
1086+
1087+
@Test
1088+
public void getRequestMetadataWithCallback() throws IOException {
1089+
MockTokenServerTransportFactory transportFactory = new MockTokenServerTransportFactory();
1090+
transportFactory.transport.addClient(CLIENT_ID, "unused-client-secret");
1091+
transportFactory.transport.addServiceAccount(CLIENT_EMAIL, ACCESS_TOKEN);
1092+
1093+
PrivateKey privateKey = ServiceAccountCredentials.privateKeyFromPkcs8(PRIVATE_KEY_PKCS8);
1094+
GoogleCredentials credentials =
1095+
ServiceAccountCredentials.newBuilder()
1096+
.setClientId(CLIENT_ID)
1097+
.setClientEmail(CLIENT_EMAIL)
1098+
.setPrivateKey(privateKey)
1099+
.setPrivateKeyId(PRIVATE_KEY_ID)
1100+
.setScopes(SCOPES)
1101+
.setServiceAccountUser(USER)
1102+
.setProjectId(PROJECT_ID)
1103+
.setQuotaProjectId("my-quota-project-id")
1104+
.setHttpTransportFactory(transportFactory)
1105+
.build();
1106+
1107+
final Map<String, List<String>> plainMetadata = credentials.getRequestMetadata();
1108+
final AtomicBoolean success = new AtomicBoolean(false);
1109+
credentials.getRequestMetadata(
1110+
null,
1111+
null,
1112+
new RequestMetadataCallback() {
1113+
@Override
1114+
public void onSuccess(Map<String, List<String>> metadata) {
1115+
assertEquals(plainMetadata, metadata);
1116+
success.set(true);
1117+
}
1118+
1119+
@Override
1120+
public void onFailure(Throwable exception) {
1121+
fail("Should not throw a failure.");
1122+
}
1123+
});
1124+
1125+
assertTrue("Should have run onSuccess() callback", success.get());
1126+
}
1127+
10351128
static GenericJson writeServiceAccountJson(
10361129
String clientId,
10371130
String clientEmail,

‎oauth2_http/javatests/com/google/auth/oauth2/ServiceAccountJwtAccessCredentialsTest.java‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import com.google.api.client.json.webtoken.JsonWebSignature;
4747
import com.google.api.client.util.Clock;
4848
import com.google.auth.Credentials;
49+
import com.google.auth.RequestMetadataCallback;
4950
import com.google.auth.TestClock;
5051
import com.google.auth.http.AuthHttpConstants;
5152
import com.google.auth.oauth2.GoogleCredentialsTest.MockHttpTransportFactory;
@@ -61,6 +62,7 @@
6162
import java.util.List;
6263
import java.util.Map;
6364
import java.util.concurrent.TimeUnit;
65+
import java.util.concurrent.atomic.AtomicBoolean;
6466
import org.junit.Test;
6567
import org.junit.runner.RunWith;
6668
import org.junit.runners.JUnit4;
@@ -758,6 +760,76 @@ public void jwtWithClaims_defaultAudience() throws IOException {
758760
verifyJwtAccess(metadata, SA_CLIENT_EMAIL, URI.create("default-audience"), SA_PRIVATE_KEY_ID);
759761
}
760762

763+
@Test
764+
public void getRequestMetadataSetsQuotaProjectId() throws IOException {
765+
PrivateKey privateKey = ServiceAccountCredentials.privateKeyFromPkcs8(SA_PRIVATE_KEY_PKCS8);
766+
ServiceAccountJwtAccessCredentials credentials =
767+
ServiceAccountJwtAccessCredentials.newBuilder()
768+
.setClientId(SA_CLIENT_ID)
769+
.setClientEmail(SA_CLIENT_EMAIL)
770+
.setPrivateKey(privateKey)
771+
.setPrivateKeyId(SA_PRIVATE_KEY_ID)
772+
.setQuotaProjectId("my-quota-project-id")
773+
.setDefaultAudience(URI.create("default-audience"))
774+
.build();
775+
776+
Map<String, List<String>> metadata = credentials.getRequestMetadata();
777+
assertTrue(metadata.containsKey("x-goog-user-project"));
778+
List<String> headerValues = metadata.get("x-goog-user-project");
779+
assertEquals(1, headerValues.size());
780+
assertEquals("my-quota-project-id", headerValues.get(0));
781+
}
782+
783+
@Test
784+
public void getRequestMetadataNoQuotaProjectId() throws IOException {
785+
PrivateKey privateKey = ServiceAccountCredentials.privateKeyFromPkcs8(SA_PRIVATE_KEY_PKCS8);
786+
ServiceAccountJwtAccessCredentials credentials =
787+
ServiceAccountJwtAccessCredentials.newBuilder()
788+
.setClientId(SA_CLIENT_ID)
789+
.setClientEmail(SA_CLIENT_EMAIL)
790+
.setPrivateKey(privateKey)
791+
.setPrivateKeyId(SA_PRIVATE_KEY_ID)
792+
.setDefaultAudience(URI.create("default-audience"))
793+
.build();
794+
795+
Map<String, List<String>> metadata = credentials.getRequestMetadata();
796+
assertFalse(metadata.containsKey("x-goog-user-project"));
797+
}
798+
799+
@Test
800+
public void getRequestMetadataWithCallback() throws IOException {
801+
PrivateKey privateKey = ServiceAccountCredentials.privateKeyFromPkcs8(SA_PRIVATE_KEY_PKCS8);
802+
ServiceAccountJwtAccessCredentials credentials =
803+
ServiceAccountJwtAccessCredentials.newBuilder()
804+
.setClientId(SA_CLIENT_ID)
805+
.setClientEmail(SA_CLIENT_EMAIL)
806+
.setPrivateKey(privateKey)
807+
.setPrivateKeyId(SA_PRIVATE_KEY_ID)
808+
.setQuotaProjectId("my-quota-project-id")
809+
.setDefaultAudience(URI.create("default-audience"))
810+
.build();
811+
812+
final Map<String, List<String>> plainMetadata = credentials.getRequestMetadata();
813+
final AtomicBoolean success = new AtomicBoolean(false);
814+
credentials.getRequestMetadata(
815+
null,
816+
null,
817+
new RequestMetadataCallback() {
818+
@Override
819+
public void onSuccess(Map<String, List<String>> metadata) {
820+
assertEquals(plainMetadata, metadata);
821+
success.set(true);
822+
}
823+
824+
@Override
825+
public void onFailure(Throwable exception) {
826+
fail("Should not throw a failure.");
827+
}
828+
});
829+
830+
assertTrue("Should have run onSuccess() callback", success.get());
831+
}
832+
761833
private void verifyJwtAccess(
762834
Map<String, List<String>> metadata,
763835
String expectedEmail,

0 commit comments

Comments
 (0)