From 20bb200c8019ad1df8acbfe210cea7d5e9a9a57c Mon Sep 17 00:00:00 2001 From: Summer Ji Date: Thu, 23 Jul 2020 22:31:04 -0700 Subject: [PATCH] feat: enable setting quota_project_id (#1128) --- .../com/google/api/gax/rpc/ClientContext.java | 53 +++- .../google/api/gax/rpc/ClientSettings.java | 16 + .../com/google/api/gax/rpc/StubSettings.java | 59 ++++ .../QuotaProjectIdHidingCredentials.java | 92 ++++++ .../google/api/gax/rpc/ClientContextTest.java | 190 +++++++++++ .../api/gax/rpc/ClientSettingsTest.java | 294 ++++++++++++++++++ .../QuotaProjectIdHidingCredentialsTest.java | 105 +++++++ 7 files changed, 803 insertions(+), 6 deletions(-) create mode 100644 gax/src/main/java/com/google/api/gax/rpc/internal/QuotaProjectIdHidingCredentials.java create mode 100644 gax/src/test/java/com/google/api/gax/rpc/internal/QuotaProjectIdHidingCredentialsTest.java diff --git a/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java b/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java index 377ec730a..1dca243d9 100644 --- a/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java +++ b/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java @@ -35,6 +35,7 @@ import com.google.api.gax.core.BackgroundResource; import com.google.api.gax.core.ExecutorAsBackgroundResource; import com.google.api.gax.core.ExecutorProvider; +import com.google.api.gax.rpc.internal.QuotaProjectIdHidingCredentials; import com.google.api.gax.tracing.ApiTracerFactory; import com.google.api.gax.tracing.NoopApiTracerFactory; import com.google.auth.Credentials; @@ -60,6 +61,7 @@ */ @AutoValue public abstract class ClientContext { + private static final String QUOTA_PROJECT_ID_HEADER_KEY = "x-goog-user-project"; /** * The objects that need to be closed in order to clean up the resources created in the process of @@ -96,6 +98,9 @@ public abstract class ClientContext { @Nullable public abstract String getEndpoint(); + @Nullable + public abstract String getQuotaProjectId(); + /** Gets the {@link ApiTracerFactory} that will be used to generate traces for operations. */ @BetaApi("The surface for tracing is not stable yet and may change in the future.") @Nonnull @@ -110,7 +115,8 @@ public static Builder newBuilder() { .setClock(NanoClock.getDefaultClock()) .setStreamWatchdog(null) .setStreamWatchdogCheckInterval(Duration.ZERO) - .setTracerFactory(NoopApiTracerFactory.getInstance()); + .setTracerFactory(NoopApiTracerFactory.getInstance()) + .setQuotaProjectId(null); } public abstract Builder toBuilder(); @@ -135,15 +141,19 @@ public static ClientContext create(StubSettings settings) throws IOException { Credentials credentials = settings.getCredentialsProvider().getCredentials(); + if (settings.getQuotaProjectId() != null) { + // If the quotaProjectId is set, wrap original credentials with correct quotaProjectId as + // QuotaProjectIdHidingCredentials. + // Ensure that a custom set quota project id takes priority over one detected by credentials. + // Avoid the backend receiving possibly conflict values of quotaProjectId + credentials = new QuotaProjectIdHidingCredentials(credentials); + } + TransportChannelProvider transportChannelProvider = settings.getTransportChannelProvider(); if (transportChannelProvider.needsExecutor()) { transportChannelProvider = transportChannelProvider.withExecutor((Executor) executor); } - Map headers = - ImmutableMap.builder() - .putAll(settings.getHeaderProvider().getHeaders()) - .putAll(settings.getInternalHeaderProvider().getHeaders()) - .build(); + Map headers = getHeadersFromSettings(settings); if (transportChannelProvider.needsHeaders()) { transportChannelProvider = transportChannelProvider.withHeaders(headers); } @@ -200,12 +210,41 @@ public static ClientContext create(StubSettings settings) throws IOException { .setClock(clock) .setDefaultCallContext(defaultCallContext) .setEndpoint(settings.getEndpoint()) + .setQuotaProjectId(settings.getQuotaProjectId()) .setStreamWatchdog(watchdog) .setStreamWatchdogCheckInterval(settings.getStreamWatchdogCheckInterval()) .setTracerFactory(settings.getTracerFactory()) .build(); } + /** + * Getting a header map from HeaderProvider and InternalHeaderProvider from settings with Quota + * Project Id. + */ + private static Map getHeadersFromSettings(StubSettings settings) { + ImmutableMap.Builder headersBuilder = ImmutableMap.builder(); + if (settings.getQuotaProjectId() != null) { + headersBuilder.put(QUOTA_PROJECT_ID_HEADER_KEY, settings.getQuotaProjectId()); + for (Map.Entry entry : settings.getHeaderProvider().getHeaders().entrySet()) { + if (entry.getKey().equals(QUOTA_PROJECT_ID_HEADER_KEY)) { + continue; + } + headersBuilder.put(entry); + } + for (Map.Entry entry : + settings.getInternalHeaderProvider().getHeaders().entrySet()) { + if (entry.getKey().equals(QUOTA_PROJECT_ID_HEADER_KEY)) { + continue; + } + headersBuilder.put(entry); + } + } else { + headersBuilder.putAll(settings.getHeaderProvider().getHeaders()); + headersBuilder.putAll(settings.getInternalHeaderProvider().getHeaders()); + } + return headersBuilder.build(); + } + @AutoValue.Builder public abstract static class Builder { @@ -229,6 +268,8 @@ public abstract static class Builder { public abstract Builder setEndpoint(String endpoint); + public abstract Builder setQuotaProjectId(String QuotaProjectId); + @BetaApi("The surface for streaming is not stable yet and may change in the future.") public abstract Builder setStreamWatchdog(Watchdog watchdog); diff --git a/gax/src/main/java/com/google/api/gax/rpc/ClientSettings.java b/gax/src/main/java/com/google/api/gax/rpc/ClientSettings.java index 13b733641..7575fea2b 100644 --- a/gax/src/main/java/com/google/api/gax/rpc/ClientSettings.java +++ b/gax/src/main/java/com/google/api/gax/rpc/ClientSettings.java @@ -93,6 +93,10 @@ public final String getEndpoint() { return stubSettings.getEndpoint(); } + public final String getQuotaProjectId() { + return stubSettings.getQuotaProjectId(); + } + @BetaApi("The surface for streaming is not stable yet and may change in the future.") @Nullable public final WatchdogProvider getWatchdogProvider() { @@ -114,6 +118,7 @@ public String toString() { .add("internalHeaderProvider", getInternalHeaderProvider()) .add("clock", getClock()) .add("endpoint", getEndpoint()) + .add("quotaProjectId", getQuotaProjectId()) .add("watchdogProvider", getWatchdogProvider()) .add("watchdogCheckInterval", getWatchdogCheckInterval()) .toString(); @@ -216,6 +221,11 @@ public B setEndpoint(String endpoint) { return self(); } + public B setQuotaProjectId(String quotaProjectId) { + stubSettings.setQuotaProjectId(quotaProjectId); + return self(); + } + @BetaApi("The surface for streaming is not stable yet and may change in the future.") public B setWatchdogProvider(@Nullable WatchdogProvider watchdogProvider) { stubSettings.setStreamWatchdogProvider(watchdogProvider); @@ -264,6 +274,11 @@ public String getEndpoint() { return stubSettings.getEndpoint(); } + /** Gets the QuotaProjectId that was previously set on this Builder. */ + public String getQuotaProjectId() { + return stubSettings.getQuotaProjectId(); + } + @BetaApi("The surface for streaming is not stable yet and may change in the future.") @Nullable public WatchdogProvider getWatchdogProvider() { @@ -294,6 +309,7 @@ public String toString() { .add("internalHeaderProvider", getInternalHeaderProvider()) .add("clock", getClock()) .add("endpoint", getEndpoint()) + .add("quotaProjectId", getQuotaProjectId()) .add("watchdogProvider", getWatchdogProvider()) .add("watchdogCheckInterval", getWatchdogCheckInterval()) .toString(); diff --git a/gax/src/main/java/com/google/api/gax/rpc/StubSettings.java b/gax/src/main/java/com/google/api/gax/rpc/StubSettings.java index 2cac53602..b22f6a451 100644 --- a/gax/src/main/java/com/google/api/gax/rpc/StubSettings.java +++ b/gax/src/main/java/com/google/api/gax/rpc/StubSettings.java @@ -41,6 +41,8 @@ import com.google.api.gax.core.NoCredentialsProvider; import com.google.api.gax.tracing.ApiTracerFactory; import com.google.api.gax.tracing.NoopApiTracerFactory; +import com.google.auth.Credentials; +import com.google.auth.oauth2.QuotaProjectIdProvider; import com.google.common.base.MoreObjects; import com.google.common.base.Preconditions; import java.io.IOException; @@ -60,6 +62,8 @@ */ public abstract class StubSettings> { + static final String QUOTA_PROJECT_ID_HEADER_KEY = "x-goog-user-project"; + private final ExecutorProvider executorProvider; private final CredentialsProvider credentialsProvider; private final HeaderProvider headerProvider; @@ -67,6 +71,7 @@ public abstract class StubSettings> { private final TransportChannelProvider transportChannelProvider; private final ApiClock clock; private final String endpoint; + private final String quotaProjectId; @Nullable private final WatchdogProvider streamWatchdogProvider; @Nonnull private final Duration streamWatchdogCheckInterval; @Nonnull private final ApiTracerFactory tracerFactory; @@ -80,6 +85,7 @@ protected StubSettings(Builder builder) { this.internalHeaderProvider = builder.internalHeaderProvider; this.clock = builder.clock; this.endpoint = builder.endpoint; + this.quotaProjectId = builder.quotaProjectId; this.streamWatchdogProvider = builder.streamWatchdogProvider; this.streamWatchdogCheckInterval = builder.streamWatchdogCheckInterval; this.tracerFactory = builder.tracerFactory; @@ -115,6 +121,10 @@ public final String getEndpoint() { return endpoint; } + public final String getQuotaProjectId() { + return quotaProjectId; + } + @BetaApi("The surface for streaming is not stable yet and may change in the future.") @Nullable public final WatchdogProvider getStreamWatchdogProvider() { @@ -146,6 +156,7 @@ public String toString() { .add("internalHeaderProvider", internalHeaderProvider) .add("clock", clock) .add("endpoint", endpoint) + .add("quotaProjectId", quotaProjectId) .add("streamWatchdogProvider", streamWatchdogProvider) .add("streamWatchdogCheckInterval", streamWatchdogCheckInterval) .add("tracerFactory", tracerFactory) @@ -164,6 +175,7 @@ public abstract static class Builder< private TransportChannelProvider transportChannelProvider; private ApiClock clock; private String endpoint; + private String quotaProjectId; @Nullable private WatchdogProvider streamWatchdogProvider; @Nonnull private Duration streamWatchdogCheckInterval; @Nonnull private ApiTracerFactory tracerFactory; @@ -177,11 +189,29 @@ protected Builder(StubSettings settings) { this.internalHeaderProvider = settings.internalHeaderProvider; this.clock = settings.clock; this.endpoint = settings.endpoint; + this.quotaProjectId = settings.quotaProjectId; this.streamWatchdogProvider = settings.streamWatchdogProvider; this.streamWatchdogCheckInterval = settings.streamWatchdogCheckInterval; this.tracerFactory = settings.tracerFactory; } + /** Get Quota Project ID from Client Context * */ + private static String getQuotaProjectIdFromClientContext(ClientContext clientContext) { + if (clientContext.getQuotaProjectId() != null) { + return clientContext.getQuotaProjectId(); + } + if (clientContext.getCredentials() instanceof QuotaProjectIdProvider) { + return ((QuotaProjectIdProvider) clientContext.getCredentials()).getQuotaProjectId(); + } + if (clientContext.getHeaders().containsKey(QUOTA_PROJECT_ID_HEADER_KEY)) { + return clientContext.getHeaders().get(QUOTA_PROJECT_ID_HEADER_KEY); + } + if (clientContext.getInternalHeaders().containsKey(QUOTA_PROJECT_ID_HEADER_KEY)) { + return clientContext.getInternalHeaders().get(QUOTA_PROJECT_ID_HEADER_KEY); + } + return null; + } + protected Builder(ClientContext clientContext) { if (clientContext == null) { this.executorProvider = InstantiatingExecutorProvider.newBuilder().build(); @@ -191,6 +221,7 @@ protected Builder(ClientContext clientContext) { this.internalHeaderProvider = new NoHeaderProvider(); this.clock = NanoClock.getDefaultClock(); this.endpoint = null; + this.quotaProjectId = null; this.streamWatchdogProvider = InstantiatingWatchdogProvider.create(); this.streamWatchdogCheckInterval = Duration.ofSeconds(10); this.tracerFactory = NoopApiTracerFactory.getInstance(); @@ -208,6 +239,7 @@ protected Builder(ClientContext clientContext) { FixedWatchdogProvider.create(clientContext.getStreamWatchdog()); this.streamWatchdogCheckInterval = clientContext.getStreamWatchdogCheckInterval(); this.tracerFactory = clientContext.getTracerFactory(); + this.quotaProjectId = getQuotaProjectIdFromClientContext(clientContext); } } @@ -234,6 +266,14 @@ public B setExecutorProvider(ExecutorProvider executorProvider) { /** Sets the CredentialsProvider to use for getting the credentials to make calls with. */ public B setCredentialsProvider(CredentialsProvider credentialsProvider) { this.credentialsProvider = Preconditions.checkNotNull(credentialsProvider); + try { + Credentials credentials = credentialsProvider.getCredentials(); + if (this.quotaProjectId == null && credentials instanceof QuotaProjectIdProvider) { + this.quotaProjectId = ((QuotaProjectIdProvider) credentials).getQuotaProjectId(); + } + } catch (IOException e) { + System.out.println("fail to fetch credentials"); + } return self(); } @@ -247,6 +287,10 @@ public B setCredentialsProvider(CredentialsProvider credentialsProvider) { @BetaApi("The surface for customizing headers is not stable yet and may change in the future.") public B setHeaderProvider(HeaderProvider headerProvider) { this.headerProvider = headerProvider; + if (this.quotaProjectId == null + && headerProvider.getHeaders().containsKey(QUOTA_PROJECT_ID_HEADER_KEY)) { + this.quotaProjectId = headerProvider.getHeaders().get(QUOTA_PROJECT_ID_HEADER_KEY); + } return self(); } @@ -260,6 +304,10 @@ public B setHeaderProvider(HeaderProvider headerProvider) { @BetaApi("The surface for customizing headers is not stable yet and may change in the future.") protected B setInternalHeaderProvider(HeaderProvider internalHeaderProvider) { this.internalHeaderProvider = internalHeaderProvider; + if (this.quotaProjectId == null + && internalHeaderProvider.getHeaders().containsKey(QUOTA_PROJECT_ID_HEADER_KEY)) { + this.quotaProjectId = internalHeaderProvider.getHeaders().get(QUOTA_PROJECT_ID_HEADER_KEY); + } return self(); } @@ -298,6 +346,11 @@ public B setEndpoint(String endpoint) { return self(); } + public B setQuotaProjectId(String quotaProjectId) { + this.quotaProjectId = quotaProjectId; + return self(); + } + /** * Sets how often the {@link Watchdog} will check ongoing streaming RPCs. Defaults to 10 secs. * Use {@link Duration#ZERO} to disable. @@ -364,6 +417,11 @@ public String getEndpoint() { return endpoint; } + /** Gets the QuotaProjectId that was previously set on this Builder. */ + public String getQuotaProjectId() { + return quotaProjectId; + } + @BetaApi("The surface for streaming is not stable yet and may change in the future.") @Nonnull public Duration getStreamWatchdogCheckInterval() { @@ -396,6 +454,7 @@ public String toString() { .add("internalHeaderProvider", internalHeaderProvider) .add("clock", clock) .add("endpoint", endpoint) + .add("quotaProjectId", quotaProjectId) .add("streamWatchdogProvider", streamWatchdogProvider) .add("streamWatchdogCheckInterval", streamWatchdogCheckInterval) .add("tracerFactory", tracerFactory) diff --git a/gax/src/main/java/com/google/api/gax/rpc/internal/QuotaProjectIdHidingCredentials.java b/gax/src/main/java/com/google/api/gax/rpc/internal/QuotaProjectIdHidingCredentials.java new file mode 100644 index 000000000..51bbade9a --- /dev/null +++ b/gax/src/main/java/com/google/api/gax/rpc/internal/QuotaProjectIdHidingCredentials.java @@ -0,0 +1,92 @@ +/* + * Copyright 2020 Google LLC + * + * 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 LLC 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.api.gax.rpc.internal; + +import com.google.api.client.util.Beta; +import com.google.auth.Credentials; +import com.google.common.collect.ImmutableMap; +import java.io.IOException; +import java.net.URI; +import java.util.List; +import java.util.Map; + +/** + * {@code QuotaProjectIdHidingCredentials} is intended to be used to wrap a {@link Credentials} to + * hide quota project id + * + *

Ensure that a custom set quota project id takes priority over one detected by credentials. + */ +@Beta +public class QuotaProjectIdHidingCredentials extends Credentials { + private static final String QUOTA_PROJECT_ID_HEADER_KEY = "x-goog-user-project"; + private final Credentials wrappedCredentials; + + public QuotaProjectIdHidingCredentials(Credentials credentials) { + this.wrappedCredentials = credentials; + } + + @Override + public String getAuthenticationType() { + return this.wrappedCredentials.getAuthenticationType(); + } + + @Override + public Map> getRequestMetadata(URI uri) throws IOException { + ImmutableMap.Builder> metaBuilder = ImmutableMap.builder(); + for (Map.Entry> entry : + this.wrappedCredentials.getRequestMetadata(uri).entrySet()) { + if (entry.getKey().equals(QUOTA_PROJECT_ID_HEADER_KEY)) { + continue; + } + metaBuilder.put(entry); + } + return metaBuilder.build(); + } + + @Override + public Map> getRequestMetadata() throws IOException { + return this.getRequestMetadata(null); + } + + @Override + public boolean hasRequestMetadata() { + return this.wrappedCredentials.hasRequestMetadata(); + } + + @Override + public boolean hasRequestMetadataOnly() { + return this.wrappedCredentials.hasRequestMetadataOnly(); + } + + @Override + public void refresh() throws IOException { + this.wrappedCredentials.refresh(); + } +} diff --git a/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java b/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java index 45be3da78..0f7305ca2 100644 --- a/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java +++ b/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java @@ -33,15 +33,18 @@ import com.google.api.core.ApiClock; import com.google.api.gax.core.BackgroundResource; +import com.google.api.gax.core.CredentialsProvider; import com.google.api.gax.core.ExecutorProvider; import com.google.api.gax.core.FixedCredentialsProvider; import com.google.api.gax.rpc.testing.FakeChannel; import com.google.api.gax.rpc.testing.FakeClientSettings; import com.google.api.gax.rpc.testing.FakeTransportChannel; import com.google.auth.Credentials; +import com.google.auth.oauth2.GoogleCredentials; import com.google.common.collect.ImmutableMap; import com.google.common.truth.Truth; import java.io.IOException; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.concurrent.Executor; @@ -336,4 +339,191 @@ public void testWatchdogProvider() throws IOException { List resources = context.getBackgroundResources(); assertThat(resources.get(2)).isInstanceOf(Watchdog.class); } + + @Test + public void testMergeHeaders_getQuotaProjectIdFromHeadersProvider() throws IOException { + final String QUOTA_PROJECT_ID_KEY = "x-goog-user-project"; + final String QUOTA_PROJECT_ID_FROM_SETTINGS = "quota_project_id_from_settings"; + FakeClientSettings.Builder builder = new FakeClientSettings.Builder(); + + InterceptingExecutor executor = new InterceptingExecutor(1); + FakeTransportChannel transportChannel = FakeTransportChannel.create(new FakeChannel()); + FakeTransportProvider transportProvider = + new FakeTransportProvider(transportChannel, executor, true, null, null); + + HeaderProvider headerProvider = Mockito.mock(HeaderProvider.class); + Mockito.when(headerProvider.getHeaders()).thenReturn(ImmutableMap.of("header_k1", "v1")); + HeaderProvider internalHeaderProvider = Mockito.mock(HeaderProvider.class); + Mockito.when(internalHeaderProvider.getHeaders()) + .thenReturn(ImmutableMap.of("internal_header_k1", "v1")); + + builder.setTransportChannelProvider(transportProvider); + builder.setCredentialsProvider( + FixedCredentialsProvider.create(Mockito.mock(Credentials.class))); + builder.setHeaderProvider(headerProvider); + builder.setInternalHeaderProvider(internalHeaderProvider); + builder.setQuotaProjectId(QUOTA_PROJECT_ID_FROM_SETTINGS); + + ClientContext context = ClientContext.create(builder.build()); + List resources = context.getBackgroundResources(); + FakeTransportChannel fakeTransportChannel = (FakeTransportChannel) resources.get(0); + assertThat(fakeTransportChannel.getHeaders().size()) + .isEqualTo( + headerProvider.getHeaders().size() + internalHeaderProvider.getHeaders().size() + 1); + assertThat(fakeTransportChannel.getHeaders().get(QUOTA_PROJECT_ID_KEY)) + .isEqualTo(QUOTA_PROJECT_ID_FROM_SETTINGS); + } + + @Test + public void testMergeHeaders_getQuotaProjectIdFromSettings() throws IOException { + final String QUOTA_PROJECT_ID_KEY = "x-goog-user-project"; + final String QUOTA_PROJECT_ID_FROM_HEADERS = "quota_project_id_from_headers"; + final String QUOTA_PROJECT_ID_FROM_INTERNAL_HEADERS = "quota_project_id_from_internal_headers"; + final String QUOTA_PROJECT_ID_FROM_SETTINGS = "quota_project_id_from_settings"; + FakeClientSettings.Builder builder = new FakeClientSettings.Builder(); + + InterceptingExecutor executor = new InterceptingExecutor(1); + FakeTransportChannel transportChannel = FakeTransportChannel.create(new FakeChannel()); + FakeTransportProvider transportProvider = + new FakeTransportProvider(transportChannel, executor, true, null, null); + + HeaderProvider headerProvider = + new HeaderProvider() { + @Override + public Map getHeaders() { + return ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_HEADERS, "k2", "v2"); + } + }; + HeaderProvider internalHeaderProvider = + new HeaderProvider() { + @Override + public Map getHeaders() { + return ImmutableMap.of( + QUOTA_PROJECT_ID_KEY, + QUOTA_PROJECT_ID_FROM_INTERNAL_HEADERS, + "internal_header_k1", + "v1"); + } + }; + + builder.setTransportChannelProvider(transportProvider); + builder.setCredentialsProvider( + FixedCredentialsProvider.create(Mockito.mock(Credentials.class))); + builder.setHeaderProvider(headerProvider); + builder.setInternalHeaderProvider(internalHeaderProvider); + builder.setQuotaProjectId(QUOTA_PROJECT_ID_FROM_SETTINGS); + + ClientContext context = ClientContext.create(builder.build()); + List resources = context.getBackgroundResources(); + FakeTransportChannel fakeTransportChannel = (FakeTransportChannel) resources.get(0); + assertThat(fakeTransportChannel.getHeaders().size()) + .isEqualTo( + headerProvider.getHeaders().size() + internalHeaderProvider.getHeaders().size() - 1); + assertThat(fakeTransportChannel.getHeaders().get(QUOTA_PROJECT_ID_KEY)) + .isEqualTo(QUOTA_PROJECT_ID_FROM_SETTINGS); + } + + @Test + public void testMergeHeaders_noQuotaProjectIdSet() throws IOException { + final String QUOTA_PROJECT_ID_KEY = "x-goog-user-project"; + FakeClientSettings.Builder builder = new FakeClientSettings.Builder(); + + InterceptingExecutor executor = new InterceptingExecutor(1); + FakeTransportChannel transportChannel = FakeTransportChannel.create(new FakeChannel()); + FakeTransportProvider transportProvider = + new FakeTransportProvider(transportChannel, executor, true, null, null); + + HeaderProvider headerProvider = Mockito.mock(HeaderProvider.class); + Mockito.when(headerProvider.getHeaders()).thenReturn(ImmutableMap.of("header_k1", "v1")); + HeaderProvider internalHeaderProvider = Mockito.mock(HeaderProvider.class); + Mockito.when(internalHeaderProvider.getHeaders()) + .thenReturn(ImmutableMap.of("internal_header_k1", "v1")); + + builder.setTransportChannelProvider(transportProvider); + builder.setCredentialsProvider( + FixedCredentialsProvider.create(Mockito.mock(Credentials.class))); + builder.setHeaderProvider(headerProvider); + builder.setInternalHeaderProvider(internalHeaderProvider); + + ClientContext context = ClientContext.create(builder.build()); + List resources = context.getBackgroundResources(); + FakeTransportChannel fakeTransportChannel = (FakeTransportChannel) resources.get(0); + assertThat(fakeTransportChannel.getHeaders().size()) + .isEqualTo(headerProvider.getHeaders().size() + internalHeaderProvider.getHeaders().size()); + assertThat(fakeTransportChannel.getHeaders().containsKey(QUOTA_PROJECT_ID_KEY)).isFalse(); + } + + @Test + public void testHidingQuotaProjectId_quotaSetFromSetting() throws IOException { + final String QUOTA_PROJECT_ID_KEY = "x-goog-user-project"; + final String QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE = "quota_project_id_from_credentials"; + FakeClientSettings.Builder builder = new FakeClientSettings.Builder(); + + InterceptingExecutor executor = new InterceptingExecutor(1); + FakeTransportChannel transportChannel = FakeTransportChannel.create(new FakeChannel()); + FakeTransportProvider transportProvider = + new FakeTransportProvider(transportChannel, executor, true, null, null); + Map> metaDataWithQuota = + ImmutableMap.of( + "k1", + Collections.singletonList("v1"), + QUOTA_PROJECT_ID_KEY, + Collections.singletonList(QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE)); + final Credentials credentialsWithQuotaProjectId = Mockito.mock(GoogleCredentials.class); + Mockito.when(credentialsWithQuotaProjectId.getRequestMetadata(null)) + .thenReturn(metaDataWithQuota); + HeaderProvider headerProviderWithQuota = Mockito.mock(HeaderProvider.class); + HeaderProvider internalHeaderProvider = Mockito.mock(HeaderProvider.class); + + builder.setExecutorProvider(new FakeExecutorProvider(executor, true)); + builder.setTransportChannelProvider(transportProvider); + builder.setCredentialsProvider( + new CredentialsProvider() { + @Override + public Credentials getCredentials() throws IOException { + return credentialsWithQuotaProjectId; + } + }); + builder.setHeaderProvider(headerProviderWithQuota); + builder.setInternalHeaderProvider(internalHeaderProvider); + builder.setQuotaProjectId(QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE); + + ClientContext clientContext = ClientContext.create(builder.build()); + assertThat(clientContext.getCredentials().getRequestMetadata().size()) + .isEqualTo(metaDataWithQuota.size() - 1); + assertThat( + clientContext.getCredentials().getRequestMetadata().containsKey(QUOTA_PROJECT_ID_KEY)) + .isFalse(); + } + + @Test + public void testHidingQuotaProjectId_noQuotaSetFromSetting() throws IOException { + final String QUOTA_PROJECT_ID_KEY = "x-goog-user-project"; + FakeClientSettings.Builder builder = new FakeClientSettings.Builder(); + + InterceptingExecutor executor = new InterceptingExecutor(1); + FakeTransportChannel transportChannel = FakeTransportChannel.create(new FakeChannel()); + FakeTransportProvider transportProvider = + new FakeTransportProvider(transportChannel, executor, true, null, null); + Map> metaData = ImmutableMap.of("k1", Collections.singletonList("v1")); + final Credentials credentialsWithoutQuotaProjectId = Mockito.mock(GoogleCredentials.class); + Mockito.when(credentialsWithoutQuotaProjectId.getRequestMetadata(null)).thenReturn(metaData); + HeaderProvider headerProviderWithQuota = Mockito.mock(HeaderProvider.class); + HeaderProvider internalHeaderProvider = Mockito.mock(HeaderProvider.class); + + builder.setExecutorProvider(new FakeExecutorProvider(executor, true)); + builder.setTransportChannelProvider(transportProvider); + builder.setCredentialsProvider( + new CredentialsProvider() { + @Override + public Credentials getCredentials() throws IOException { + return credentialsWithoutQuotaProjectId; + } + }); + builder.setHeaderProvider(headerProviderWithQuota); + builder.setInternalHeaderProvider(internalHeaderProvider); + + ClientContext clientContext = ClientContext.create(builder.build()); + assertThat(clientContext.getCredentials().getRequestMetadata(null)).isEqualTo(metaData); + } } diff --git a/gax/src/test/java/com/google/api/gax/rpc/ClientSettingsTest.java b/gax/src/test/java/com/google/api/gax/rpc/ClientSettingsTest.java index ee490107d..cf842a923 100644 --- a/gax/src/test/java/com/google/api/gax/rpc/ClientSettingsTest.java +++ b/gax/src/test/java/com/google/api/gax/rpc/ClientSettingsTest.java @@ -29,6 +29,8 @@ */ package com.google.api.gax.rpc; +import static org.junit.Assert.fail; + import com.google.api.core.ApiClock; import com.google.api.core.ApiFunction; import com.google.api.core.NanoClock; @@ -41,7 +43,12 @@ import com.google.api.gax.rpc.testing.FakeCallContext; import com.google.api.gax.rpc.testing.FakeClientSettings; import com.google.auth.Credentials; +import com.google.auth.oauth2.GoogleCredentials; +import com.google.common.collect.ImmutableMap; import com.google.common.truth.Truth; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; import java.util.ArrayList; import java.util.Collections; import java.util.List; @@ -55,6 +62,49 @@ @RunWith(JUnit4.class) public class ClientSettingsTest { + private static final String QUOTA_PROJECT_ID_KEY = "x-goog-user-project"; + private static final String QUOTA_PROJECT_ID_FROM_HEADER_VALUE = "quota_project_id_from_headers"; + private static final String QUOTA_PROJECT_ID_FROM_BUILDERS = "quota_project_id_from_builders"; + private static final String QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE = + "quota_project_id_from_internal_headers"; + private static final String QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE = + "quota_project_id_from_credentials"; + private static final String QUOTA_PROJECT_ID_FROM_CONTEXT = + "quota_project_id_from_client_context"; + private static final String JSON_KEY_QUOTA_PROJECT_ID = + "{\n" + + " \"private_key_id\": \"somekeyid\",\n" + + " \"private_key\": \"-----BEGIN PRIVATE KEY-----\\nMIIEvgIBADANBgkqhkiG9w0BAQEFAASCBKgwggS" + + "kAgEAAoIBAQC+K2hSuFpAdrJI\\nnCgcDz2M7t7bjdlsadsasad+fvRSW6TjNQZ3p5LLQY1kSZRqBqylRkzteMOyHg" + + "aR\\n0Pmxh3ILCND5men43j3h4eDbrhQBuxfEMalkG92sL+PNQSETY2tnvXryOvmBRwa/\\nQP/9dJfIkIDJ9Fw9N4" + + "Bhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nknddadwkwewcVxHFhcZJO+XWf6ofLUXpRwiTZakGMn8EE1uVa2" + + "LgczOjwWHGi99MFjxSer5m9\\n1tCa3/KEGKiS/YL71JvjwX3mb+cewlkcmweBKZHM2JPTk0ZednFSpVZMtycjkbLa" + + "\\ndYOS8V85AgMBewECggEBAKksaldajfDZDV6nGqbFjMiizAKJolr/M3OQw16K6o3/\\n0S31xIe3sSlgW0+UbYlF" + + "4U8KifhManD1apVSC3csafaspP4RZUHFhtBywLO9pR5c\\nr6S5aLp+gPWFyIp1pfXbWGvc5VY/v9x7ya1VEa6rXvL" + + "sKupSeWAW4tMj3eo/64ge\\nsdaceaLYw52KeBYiT6+vpsnYrEkAHO1fF/LavbLLOFJmFTMxmsNaG0tuiJHgjshB\\" + + "n82DpMCbXG9YcCgI/DbzuIjsdj2JC1cascSP//3PmefWysucBQe7Jryb6NQtASmnv\\nCdDw/0jmZTEjpe4S1lxfHp" + + "lAhHFtdgYTvyYtaLZiVVkCgYEA8eVpof2rceecw/I6\\n5ng1q3Hl2usdWV/4mZMvR0fOemacLLfocX6IYxT1zA1FF" + + "JlbXSRsJMf/Qq39mOR2\\nSpW+hr4jCoHeRVYLgsbggtrevGmILAlNoqCMpGZ6vDmJpq6ECV9olliDvpPgWOP+\\nm" + + "YPDreFBGxWvQrADNbRt2dmGsrsCgYEAyUHqB2wvJHFqdmeBsaacewzV8x9WgmeX\\ngUIi9REwXlGDW0Mz50dxpxcK" + + "CAYn65+7TCnY5O/jmL0VRxU1J2mSWyWTo1C+17L0\\n3fUqjxL1pkefwecxwecvC+gFFYdJ4CQ/MHHXU81Lwl1iWdF" + + "Cd2UoGddYaOF+KNeM\\nHC7cmqra+JsCgYEAlUNywzq8nUg7282E+uICfCB0LfwejuymR93CtsFgb7cRd6ak\\nECR" + + "8FGfCpH8ruWJINllbQfcHVCX47ndLZwqv3oVFKh6pAS/vVI4dpOepP8++7y1u\\ncoOvtreXCX6XqfrWDtKIvv0vjl" + + "HBhhhp6mCcRpdQjV38H7JsyJ7lih/oNjECgYAt\\nkndj5uNl5SiuVxHFhcZJO+XWf6ofLUregtevZakGMn8EE1uVa" + + "2AY7eafmoU/nZPT\\n00YB0TBATdCbn/nBSuKDESkhSg9s2GEKQZG5hBmL5uCMfo09z3SfxZIhJdlerreP\\nJ7gSi" + + "dI12N+EZxYd4xIJh/HFDgp7RRO87f+WJkofMQKBgGTnClK1VMaCRbJZPriw\\nEfeFCoOX75MxKwXs6xgrw4W//AYG" + + "GUjDt83lD6AZP6tws7gJ2IwY/qP7+lyhjEqN\\nHtfPZRGFkGZsdaksdlaksd323423d+15/UvrlRSFPNj1tWQmNKk" + + "XyRDW4IG1Oa2p\\nrALStNBx5Y9t0/LQnFI4w3aG\\n-----END PRIVATE KEY-----\\n\",\n" + + " \"project_id\": \"someprojectid\",\n" + + " \"client_email\": \"someclientid@developer.gserviceaccount.com\",\n" + + " \"client_id\": \"someclientid.apps.googleusercontent.com\",\n" + + " \"type\": \"service_account\",\n" + + " \"quota_project_id\": \"" + + QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE + + "\"\n" + + "}"; + + private static final GoogleCredentials credentialsWithQuotaProject = + loadCredentials(JSON_KEY_QUOTA_PROJECT_ID); @Test public void testEmptyBuilder() throws Exception { @@ -69,6 +119,7 @@ public void testEmptyBuilder() throws Exception { Truth.assertThat(builder.getWatchdogProvider()) .isInstanceOf(InstantiatingWatchdogProvider.class); Truth.assertThat(builder.getWatchdogCheckInterval()).isGreaterThan(Duration.ZERO); + Truth.assertThat(builder.getQuotaProjectId()).isNull(); FakeClientSettings settings = builder.build(); Truth.assertThat(settings.getExecutorProvider()) @@ -84,6 +135,7 @@ public void testEmptyBuilder() throws Exception { Truth.assertThat(settings.getWatchdogProvider()) .isInstanceOf(InstantiatingWatchdogProvider.class); Truth.assertThat(settings.getWatchdogCheckInterval()).isGreaterThan(Duration.ZERO); + Truth.assertThat((settings.getQuotaProjectId())).isSameInstanceAs(builder.getQuotaProjectId()); String settingsString = settings.toString(); Truth.assertThat(settingsString).contains("executorProvider"); @@ -93,6 +145,7 @@ public void testEmptyBuilder() throws Exception { Truth.assertThat(settingsString).contains("headerProvider"); Truth.assertThat(settingsString).contains("watchdogProvider"); Truth.assertThat(settingsString).contains("watchdogCheckInterval"); + Truth.assertThat(settingsString).contains(("quotaProjectId")); } @Test @@ -107,6 +160,7 @@ public void testBuilder() throws Exception { HeaderProvider internalHeaderProvider = Mockito.mock(HeaderProvider.class); WatchdogProvider watchdogProvider = Mockito.mock(WatchdogProvider.class); Duration watchdogCheckInterval = Duration.ofSeconds(13); + String quotaProjectId = "test_quota_project_id"; builder.setExecutorProvider(executorProvider); builder.setTransportChannelProvider(transportProvider); @@ -116,6 +170,7 @@ public void testBuilder() throws Exception { builder.setClock(clock); builder.setWatchdogProvider(watchdogProvider); builder.setWatchdogCheckInterval(watchdogCheckInterval); + builder.setQuotaProjectId(quotaProjectId); Truth.assertThat(builder.getExecutorProvider()).isSameInstanceAs(executorProvider); Truth.assertThat(builder.getTransportChannelProvider()).isSameInstanceAs(transportProvider); @@ -125,6 +180,7 @@ public void testBuilder() throws Exception { Truth.assertThat(builder.getInternalHeaderProvider()).isSameInstanceAs(internalHeaderProvider); Truth.assertThat(builder.getWatchdogProvider()).isSameInstanceAs(watchdogProvider); Truth.assertThat(builder.getWatchdogCheckInterval()).isSameInstanceAs(watchdogCheckInterval); + Truth.assertThat(builder.getQuotaProjectId()).isEqualTo(quotaProjectId); String builderString = builder.toString(); Truth.assertThat(builderString).contains("executorProvider"); @@ -135,10 +191,12 @@ public void testBuilder() throws Exception { Truth.assertThat(builderString).contains("internalHeaderProvider"); Truth.assertThat(builderString).contains("watchdogProvider"); Truth.assertThat(builderString).contains("watchdogCheckInterval"); + Truth.assertThat(builderString).contains("quotaProjectId"); } @Test public void testBuilderFromClientContext() throws Exception { + final String QUOTA_PROJECT_ID_FROM_CONTEXT = "some_quota_project_id_from_context"; ApiClock clock = Mockito.mock(ApiClock.class); ApiCallContext callContext = FakeCallContext.createDefault(); Map headers = Collections.singletonMap("spiffykey", "spiffyvalue"); @@ -159,6 +217,7 @@ public void testBuilderFromClientContext() throws Exception { .setHeaders(headers) .setStreamWatchdog(watchdog) .setStreamWatchdogCheckInterval(watchdogCheckInterval) + .setQuotaProjectId(QUOTA_PROJECT_ID_FROM_CONTEXT) .build(); FakeClientSettings.Builder builder = new FakeClientSettings.Builder(clientContext); @@ -173,6 +232,7 @@ public void testBuilderFromClientContext() throws Exception { Truth.assertThat(builder.getWatchdogProvider()).isInstanceOf(FixedWatchdogProvider.class); Truth.assertThat(builder.getWatchdogProvider().getWatchdog()).isSameInstanceAs(watchdog); Truth.assertThat(builder.getWatchdogCheckInterval()).isEqualTo(watchdogCheckInterval); + Truth.assertThat(builder.getQuotaProjectId()).isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT); } @Test @@ -187,6 +247,7 @@ public void testBuilderFromSettings() throws Exception { HeaderProvider internalHeaderProvider = Mockito.mock(HeaderProvider.class); WatchdogProvider watchdogProvider = Mockito.mock(WatchdogProvider.class); Duration watchdogCheckInterval = Duration.ofSeconds(14); + String quotaProjectId = "test_builder_from_settings_quotaProjectId"; builder.setExecutorProvider(executorProvider); builder.setTransportChannelProvider(transportProvider); @@ -196,6 +257,7 @@ public void testBuilderFromSettings() throws Exception { builder.setInternalHeaderProvider(internalHeaderProvider); builder.setWatchdogProvider(watchdogProvider); builder.setWatchdogCheckInterval(watchdogCheckInterval); + builder.setQuotaProjectId(quotaProjectId); FakeClientSettings settings = builder.build(); FakeClientSettings.Builder newBuilder = new FakeClientSettings.Builder(settings); @@ -209,6 +271,7 @@ public void testBuilderFromSettings() throws Exception { .isSameInstanceAs(internalHeaderProvider); Truth.assertThat(newBuilder.getWatchdogProvider()).isSameInstanceAs(watchdogProvider); Truth.assertThat(newBuilder.getWatchdogCheckInterval()).isEqualTo(watchdogCheckInterval); + Truth.assertThat(newBuilder.getQuotaProjectId()).isEqualTo(quotaProjectId); } @Test @@ -238,4 +301,235 @@ public Void apply(UnaryCallSettings.Builder input) { Truth.assertThat(builders.get(1).getRetryableCodes()) .containsExactly(StatusCode.Code.DEADLINE_EXCEEDED); } + + static GoogleCredentials loadCredentials(String credentialFile) { + try { + InputStream keyStream = new ByteArrayInputStream(credentialFile.getBytes()); + return GoogleCredentials.fromStream(keyStream); + } catch (IOException e) { + fail("Couldn't create fake JSON credentials."); + } + return null; + } + + @Test + public void testBuilderFromSettings_QuotaProjectId() { + + CredentialsProvider credentialsProvider_no_quota = Mockito.mock(CredentialsProvider.class); + HeaderProvider headerProvider_no_quota = Mockito.mock(HeaderProvider.class); + HeaderProvider internalHeaderProvider_no_quota = Mockito.mock(HeaderProvider.class); + HeaderProvider headerProvider_with_quota = + new HeaderProvider() { + @Override + public Map getHeaders() { + return Collections.singletonMap( + QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_HEADER_VALUE); + } + }; + HeaderProvider internalHeaderProvider_with_quota = + new HeaderProvider() { + @Override + public Map getHeaders() { + return Collections.singletonMap( + QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE); + } + }; + CredentialsProvider credentialsProvider_with_quota = + new CredentialsProvider() { + @Override + public Credentials getCredentials() throws IOException { + return credentialsWithQuotaProject; + } + }; + + // Case for setting quota_project_id from builder only + // expect value is from builders + FakeClientSettings.Builder builder_setQuotaOnly = new FakeClientSettings.Builder(); + builder_setQuotaOnly.setCredentialsProvider(credentialsProvider_no_quota); + builder_setQuotaOnly.setHeaderProvider(headerProvider_no_quota); + builder_setQuotaOnly.setInternalHeaderProvider(internalHeaderProvider_no_quota); + builder_setQuotaOnly.setQuotaProjectId(QUOTA_PROJECT_ID_FROM_BUILDERS); + + // Case for setting quota_project_id from HeaderProvider Only + // expect value is from HeaderProvider + FakeClientSettings.Builder builder_setQuotaFromHeadersOnly = new FakeClientSettings.Builder(); + builder_setQuotaFromHeadersOnly.setHeaderProvider(headerProvider_with_quota); + + // Case for setting quota_project_id from HeaderProvider and set from builders + // expect value is from builders + FakeClientSettings.Builder builder_setQuotaFromHeadersAndBuilders = + new FakeClientSettings.Builder(); + builder_setQuotaFromHeadersOnly.setHeaderProvider(headerProvider_with_quota); + builder_setQuotaFromHeadersAndBuilders.setQuotaProjectId(QUOTA_PROJECT_ID_FROM_BUILDERS); + + // Case for setting quota_project_id from InternalHeaderProvider and set from builders + // expect value is from InternalHeaderProvider + FakeClientSettings.Builder builder_setQuotaFromInternalHeadersOnly = + new FakeClientSettings.Builder(); + builder_setQuotaFromInternalHeadersOnly.setInternalHeaderProvider( + internalHeaderProvider_with_quota); + + // Case for setting quota_project_id from InternalHeaderProvider and set from builders + // expect value is from builders + FakeClientSettings.Builder builder_setQuotaFromInternalHeadersAndBuilders = + new FakeClientSettings.Builder(); + builder_setQuotaFromInternalHeadersAndBuilders.setInternalHeaderProvider( + internalHeaderProvider_with_quota); + builder_setQuotaFromInternalHeadersAndBuilders.setQuotaProjectId( + QUOTA_PROJECT_ID_FROM_BUILDERS); + + // Case for setting quota_project_id from CredentialProvider Only + // expect value is from CredentialProvider + FakeClientSettings.Builder builder_setQuotaFromCredentialsProvider = + new FakeClientSettings.Builder(); + builder_setQuotaFromCredentialsProvider.setCredentialsProvider(credentialsProvider_with_quota); + + // Case for setting quota_project_id from CredentialProvider and Builders + // expect value is from builders + FakeClientSettings.Builder builder_setQuotaFromCredentialsProviderAndBuilder = + new FakeClientSettings.Builder(); + builder_setQuotaFromCredentialsProviderAndBuilder.setCredentialsProvider( + credentialsProvider_with_quota); + builder_setQuotaFromCredentialsProviderAndBuilder.setQuotaProjectId( + QUOTA_PROJECT_ID_FROM_BUILDERS); + + // Case for setting quota_project_id from All three sources + // expect value is from builders + FakeClientSettings.Builder builder_setQuotaFromAllSources = new FakeClientSettings.Builder(); + builder_setQuotaFromAllSources.setHeaderProvider(headerProvider_with_quota); + builder_setQuotaFromAllSources.setInternalHeaderProvider(headerProvider_with_quota); + builder_setQuotaFromAllSources.setCredentialsProvider(credentialsProvider_with_quota); + builder_setQuotaFromAllSources.setQuotaProjectId(QUOTA_PROJECT_ID_FROM_BUILDERS); + + // Case for setting quota_project_id from All three sources but set from builders first + // expect value is from builders + FakeClientSettings.Builder builder_setQuotaFromAllSourcesOrder = + new FakeClientSettings.Builder(); + builder_setQuotaFromAllSourcesOrder.setQuotaProjectId(QUOTA_PROJECT_ID_FROM_BUILDERS); + builder_setQuotaFromAllSourcesOrder.setHeaderProvider(headerProvider_with_quota); + builder_setQuotaFromAllSourcesOrder.setInternalHeaderProvider(headerProvider_with_quota); + builder_setQuotaFromAllSourcesOrder.setCredentialsProvider(credentialsProvider_with_quota); + + Truth.assertThat(builder_setQuotaOnly.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_BUILDERS); + Truth.assertThat(builder_setQuotaFromHeadersOnly.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_HEADER_VALUE); + Truth.assertThat(builder_setQuotaFromHeadersAndBuilders.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_BUILDERS); + Truth.assertThat((builder_setQuotaFromInternalHeadersOnly).getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE); + Truth.assertThat(builder_setQuotaFromInternalHeadersAndBuilders.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_BUILDERS); + Truth.assertThat(builder_setQuotaFromCredentialsProvider.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE); + Truth.assertThat(builder_setQuotaFromCredentialsProviderAndBuilder.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_BUILDERS); + Truth.assertThat(builder_setQuotaFromAllSources.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_BUILDERS); + Truth.assertThat(builder_setQuotaFromAllSourcesOrder.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_BUILDERS); + } + + @Test + public void testBuilderFromClientContext_QuotaProjectId() { + ApiCallContext callContext = FakeCallContext.createDefault(); + + ClientContext clientContextQuotaOnly = + ClientContext.newBuilder() + .setTransportChannel(Mockito.mock(TransportChannel.class)) + .setDefaultCallContext(callContext) + .setQuotaProjectId(QUOTA_PROJECT_ID_FROM_CONTEXT) + .build(); + FakeClientSettings.Builder builderQuotaOnly = + new FakeClientSettings.Builder(clientContextQuotaOnly); + + ClientContext clientContextCredentialOnly = + ClientContext.newBuilder() + .setTransportChannel(Mockito.mock(TransportChannel.class)) + .setDefaultCallContext(callContext) + .setCredentials(credentialsWithQuotaProject) + .build(); + FakeClientSettings.Builder builderCredentialOnly = + new FakeClientSettings.Builder(clientContextCredentialOnly); + + ClientContext clientContextCredentialAndQuota = + ClientContext.newBuilder() + .setTransportChannel(Mockito.mock(TransportChannel.class)) + .setDefaultCallContext(callContext) + .setCredentials(credentialsWithQuotaProject) + .setQuotaProjectId(QUOTA_PROJECT_ID_FROM_CONTEXT) + .build(); + FakeClientSettings.Builder builderCredentialAndQuota = + new FakeClientSettings.Builder(clientContextCredentialAndQuota); + + ClientContext clientContextHeadersOnly = + ClientContext.newBuilder() + .setTransportChannel(Mockito.mock(TransportChannel.class)) + .setDefaultCallContext(callContext) + .setHeaders(ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_HEADER_VALUE)) + .build(); + FakeClientSettings.Builder builderHeadersOnly = + new FakeClientSettings.Builder(clientContextHeadersOnly); + + ClientContext clientContextHeadersAndQuota = + ClientContext.newBuilder() + .setTransportChannel(Mockito.mock(TransportChannel.class)) + .setDefaultCallContext(callContext) + .setQuotaProjectId(QUOTA_PROJECT_ID_FROM_CONTEXT) + .setHeaders(ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_HEADER_VALUE)) + .build(); + FakeClientSettings.Builder builderHeadersAndQuota = + new FakeClientSettings.Builder(clientContextHeadersAndQuota); + + ClientContext clientContextInternalHeadersOnly = + ClientContext.newBuilder() + .setTransportChannel(Mockito.mock(TransportChannel.class)) + .setDefaultCallContext(callContext) + .setInternalHeaders( + ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE)) + .build(); + FakeClientSettings.Builder builderInternalHeadersOnly = + new FakeClientSettings.Builder(clientContextInternalHeadersOnly); + + ClientContext clientContextInternalHeadersAndQuota = + ClientContext.newBuilder() + .setTransportChannel(Mockito.mock(TransportChannel.class)) + .setDefaultCallContext(callContext) + .setInternalHeaders( + ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE)) + .setQuotaProjectId(QUOTA_PROJECT_ID_FROM_CONTEXT) + .build(); + FakeClientSettings.Builder builderInternalHeadersAndQuota = + new FakeClientSettings.Builder(clientContextInternalHeadersAndQuota); + + ClientContext clientContextQuotaFromAllSources = + ClientContext.newBuilder() + .setTransportChannel(Mockito.mock(TransportChannel.class)) + .setDefaultCallContext(callContext) + .setHeaders( + ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE)) + .setCredentials(credentialsWithQuotaProject) + .setQuotaProjectId(QUOTA_PROJECT_ID_FROM_CONTEXT) + .setInternalHeaders( + ImmutableMap.of(QUOTA_PROJECT_ID_KEY, QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE)) + .build(); + FakeClientSettings.Builder builderQuotaFromAllSources = + new FakeClientSettings.Builder(clientContextQuotaFromAllSources); + + Truth.assertThat(builderQuotaOnly.getQuotaProjectId()).isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT); + Truth.assertThat(builderCredentialOnly.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE); + Truth.assertThat(builderCredentialAndQuota.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT); + Truth.assertThat(builderHeadersOnly.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_HEADER_VALUE); + Truth.assertThat(builderHeadersAndQuota.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT); + Truth.assertThat(builderInternalHeadersOnly.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_INTERNAL_HEADER_VALUE); + Truth.assertThat(builderInternalHeadersAndQuota.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT); + Truth.assertThat(builderQuotaFromAllSources.getQuotaProjectId()) + .isEqualTo(QUOTA_PROJECT_ID_FROM_CONTEXT); + } } diff --git a/gax/src/test/java/com/google/api/gax/rpc/internal/QuotaProjectIdHidingCredentialsTest.java b/gax/src/test/java/com/google/api/gax/rpc/internal/QuotaProjectIdHidingCredentialsTest.java new file mode 100644 index 000000000..af36e0a5b --- /dev/null +++ b/gax/src/test/java/com/google/api/gax/rpc/internal/QuotaProjectIdHidingCredentialsTest.java @@ -0,0 +1,105 @@ +/* + * Copyright 2020 Google LLC + * + * 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 LLC 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.api.gax.rpc.internal; + +import com.google.auth.Credentials; +import com.google.auth.oauth2.GoogleCredentials; +import com.google.common.collect.ImmutableMap; +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import org.junit.Assert; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.Mockito; + +@RunWith(JUnit4.class) +public class QuotaProjectIdHidingCredentialsTest { + private static final String QUOTA_PROJECT_ID_KEY = "x-goog-user-project"; + private static final String QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE = + "quota_project_id_from_credentials"; + + @Test + public void quotaProjectIdHidingCredentials_getRequestMetadata() throws IOException { + // Credentials with quota project id + Map> metaDataWithQuota = + ImmutableMap.of( + "k1", + Collections.singletonList("v1"), + QUOTA_PROJECT_ID_KEY, + Collections.singletonList(QUOTA_PROJECT_ID_FROM_CREDENTIALS_VALUE)); + Credentials credentialsWithQuotaProjectId = Mockito.mock(GoogleCredentials.class); + Mockito.when(credentialsWithQuotaProjectId.getRequestMetadata(null)) + .thenReturn(metaDataWithQuota); + QuotaProjectIdHidingCredentials quotaProjectIdHidingCredentials = + new QuotaProjectIdHidingCredentials(credentialsWithQuotaProjectId); + Map> metaDataHidingQuota = + quotaProjectIdHidingCredentials.getRequestMetadata(); + + Assert.assertTrue(metaDataWithQuota.containsKey(QUOTA_PROJECT_ID_KEY)); + Assert.assertFalse(metaDataHidingQuota.containsKey(QUOTA_PROJECT_ID_KEY)); + Assert.assertEquals(metaDataWithQuota.size() - 1, metaDataHidingQuota.size()); + + // Credentials without quota project id + Map> metaDataWithoutQuota = + ImmutableMap.of("k1", Collections.singletonList("v1")); + Credentials credentialsWithoutQuotaProjectId = Mockito.mock(GoogleCredentials.class); + Mockito.when(credentialsWithoutQuotaProjectId.getRequestMetadata(null)) + .thenReturn(metaDataWithoutQuota); + QuotaProjectIdHidingCredentials quotaProjectIdHidingCredentialsWithout = + new QuotaProjectIdHidingCredentials(credentialsWithoutQuotaProjectId); + Map> metaDataHidingQuotaWithout = + quotaProjectIdHidingCredentials.getRequestMetadata(); + + Assert.assertEquals( + quotaProjectIdHidingCredentialsWithout.getRequestMetadata(), metaDataWithoutQuota); + } + + @Test + public void quotaProjectIdHidingCredentials_getAuthenticationType() throws IOException { + final String mockType = "mock_type"; + Credentials credentials = Mockito.mock(GoogleCredentials.class); + Mockito.when(credentials.getAuthenticationType()).thenReturn(mockType); + Mockito.when(credentials.hasRequestMetadata()).thenReturn(true); + Mockito.when(credentials.hasRequestMetadataOnly()).thenReturn(false); + + QuotaProjectIdHidingCredentials quotaProjectIdHidingCredentials = + new QuotaProjectIdHidingCredentials(credentials); + quotaProjectIdHidingCredentials.refresh(); + + Assert.assertEquals(quotaProjectIdHidingCredentials.getAuthenticationType(), mockType); + Assert.assertTrue(quotaProjectIdHidingCredentials.hasRequestMetadata()); + Assert.assertFalse(quotaProjectIdHidingCredentials.hasRequestMetadataOnly()); + + Mockito.verify(credentials, Mockito.atLeastOnce()).refresh(); + } +}