From ca3ec24f506f0d43623a92ba0eed908f874fe488 Mon Sep 17 00:00:00 2001 From: Min Zhu Date: Fri, 4 Oct 2024 15:31:13 -0400 Subject: [PATCH] feat(gax): append cred-type header for auth metrics (#3186) This is POC change in gax-java for auth metrics requirements on token usage. See go/googleapis-auth-metric-design for context. [Credentials](https://github.com/googleapis/google-auth-library-java/blob/main/credentials/java/com/google/auth/Credentials.java) will expose `getMetricsCredentialType()` method, this change appends it to existing `x-goog-api-client` header Note: Currently implement in gax at client level. There are 2 edge cases not covered and will create followups for: if handwritten library overrides credentials at rpc level; If handwritten library does not build on gax. (ref: b/370039645, b/370038458) related change in `google-auth-library`https://github.com/googleapis/google-auth-library-java/pull/1503 included in 1.28.0. --- .../com/google/api/gax/rpc/ClientContext.java | 23 +++++- .../google/api/gax/rpc/ClientContextTest.java | 82 +++++++++++++++++++ 2 files changed, 103 insertions(+), 2 deletions(-) diff --git a/gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java b/gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java index 0e73bb38ef..5bce1ac6bb 100644 --- a/gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java +++ b/gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java @@ -44,6 +44,7 @@ import com.google.api.gax.tracing.ApiTracerFactory; import com.google.api.gax.tracing.BaseApiTracerFactory; import com.google.auth.ApiKeyCredentials; +import com.google.auth.CredentialTypeForMetrics; import com.google.auth.Credentials; import com.google.auth.oauth2.GdchCredentials; import com.google.auto.value.AutoValue; @@ -210,7 +211,8 @@ public static ClientContext create(StubSettings settings) throws IOException { if (transportChannelProvider.needsExecutor() && settings.getExecutorProvider() != null) { transportChannelProvider = transportChannelProvider.withExecutor(backgroundExecutor); } - Map headers = getHeadersFromSettings(settings); + + Map headers = getHeaders(settings, credentials); if (transportChannelProvider.needsHeaders()) { transportChannelProvider = transportChannelProvider.withHeaders(headers); } @@ -318,8 +320,11 @@ static GdchCredentials getGdchCredentials( /** * Getting a header map from HeaderProvider and InternalHeaderProvider from settings with Quota * Project Id. + * + *

Then if credentials is present and its type for metrics is not {@code + * CredentialTypeForMetrics.DO_NOT_SEND}, append this type info to x-goog-api-client header. */ - private static Map getHeadersFromSettings(StubSettings settings) { + private static Map getHeaders(StubSettings settings, Credentials credentials) { // Resolve conflicts when merging headers from multiple sources Map userHeaders = settings.getHeaderProvider().getHeaders(); Map internalHeaders = settings.getInternalHeaderProvider().getHeaders(); @@ -346,6 +351,20 @@ private static Map getHeadersFromSettings(StubSettings settings) effectiveHeaders.putAll(userHeaders); effectiveHeaders.putAll(conflictResolution); + return appendCredentialTypeToHeaderIfPresent(effectiveHeaders, credentials); + } + + private static Map appendCredentialTypeToHeaderIfPresent( + Map effectiveHeaders, Credentials credentials) { + CredentialTypeForMetrics credentialTypeForMetrics = + credentials == null + ? CredentialTypeForMetrics.DO_NOT_SEND + : credentials.getMetricsCredentialType(); + if (credentialTypeForMetrics != CredentialTypeForMetrics.DO_NOT_SEND) { + effectiveHeaders.computeIfPresent( + ApiClientHeaderProvider.getDefaultApiClientHeaderKey(), + (key, value) -> value + " cred-type/" + credentialTypeForMetrics.getLabel()); + } return ImmutableMap.copyOf(effectiveHeaders); } diff --git a/gax-java/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java b/gax-java/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java index c4d2027632..826864a49c 100644 --- a/gax-java/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java +++ b/gax-java/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java @@ -39,6 +39,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import com.google.api.core.ApiClock; import com.google.api.gax.core.BackgroundResource; @@ -53,6 +54,7 @@ import com.google.api.gax.rpc.testing.FakeStubSettings; import com.google.api.gax.rpc.testing.FakeTransportChannel; import com.google.auth.ApiKeyCredentials; +import com.google.auth.CredentialTypeForMetrics; import com.google.auth.Credentials; import com.google.auth.oauth2.ComputeEngineCredentials; import com.google.auth.oauth2.GdchCredentials; @@ -677,6 +679,86 @@ void testUserAgentConcat() throws Exception { .containsEntry("user-agent", "user-supplied-agent internal-agent"); } + @Test + void testApiClientHeaderAppendsCredType() throws Exception { + GoogleCredentials googleCredentials = Mockito.mock(GoogleCredentials.class); + when(googleCredentials.getMetricsCredentialType()) + .thenReturn(CredentialTypeForMetrics.USER_CREDENTIALS); + + Map headers = + setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders( + FixedCredentialsProvider.create(googleCredentials), + FixedHeaderProvider.create("x-goog-api-client", "internal-agent")); + + assertThat(headers).containsEntry("x-goog-api-client", "internal-agent cred-type/u"); + } + + @Test + void testApiClientHeaderDoNotAppendsCredType_whenNoApiClientHeader() throws Exception { + GoogleCredentials googleCredentials = Mockito.mock(GoogleCredentials.class); + when(googleCredentials.getMetricsCredentialType()) + .thenReturn(CredentialTypeForMetrics.USER_CREDENTIALS); + + Map headers = + setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders( + FixedCredentialsProvider.create(googleCredentials), + FixedHeaderProvider.create("some-other-header", "internal-agent")); + + assertThat(headers).doesNotContainKey("x-goog-api-client"); + assertThat(headers).containsEntry("some-other-header", "internal-agent"); + } + + @Test + void testApiClientHeaderDoNotAppendsCredType_whenNullCredentials() throws IOException { + Map headers = + setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders( + NoCredentialsProvider.create(), + FixedHeaderProvider.create("x-goog-api-client", "internal-agent")); + + assertThat(headers).containsKey("x-goog-api-client"); + assertThat(headers.get("x-goog-api-client")).doesNotContain("cred-type/"); + } + + @Test + void testApiClientHeaderDoNotAppendsCredType_whenCredTypeDoNotSend() throws Exception { + GoogleCredentials googleCredentials = Mockito.mock(GoogleCredentials.class); + when(googleCredentials.getMetricsCredentialType()) + .thenReturn(CredentialTypeForMetrics.DO_NOT_SEND); + + Map headers = + setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders( + FixedCredentialsProvider.create(googleCredentials), + FixedHeaderProvider.create("x-goog-api-client", "internal-agent")); + + assertThat(headers).containsKey("x-goog-api-client"); + assertThat(headers.get("x-goog-api-client")).doesNotContain("cred-type/"); + } + + private Map setupTestForCredentialTokenUsageMetricsAndGetTransportChannelHeaders( + CredentialsProvider credentialsProvider, HeaderProvider headerProvider) throws IOException { + TransportChannelProvider transportChannelProvider = + new FakeTransportProvider( + FakeTransportChannel.create(new FakeChannel()), + null, + true, + null, + null, + DEFAULT_ENDPOINT); + + ClientSettings.Builder builder = + new FakeClientSettings.Builder() + .setExecutorProvider( + FixedExecutorProvider.create(Mockito.mock(ScheduledExecutorService.class))) + .setTransportChannelProvider(transportChannelProvider) + .setCredentialsProvider(credentialsProvider) + .setInternalHeaderProvider(headerProvider); + + ClientContext clientContext = ClientContext.create(builder.build()); + FakeTransportChannel transportChannel = + (FakeTransportChannel) clientContext.getTransportChannel(); + return transportChannel.getHeaders(); + } + private static String endpoint = "https://foo.googleapis.com"; private static String mtlsEndpoint = "https://foo.mtls.googleapis.com";