Skip to content

Commit

Permalink
feat(gax): append cred-type header for auth metrics (#3186)
Browse files Browse the repository at this point in the history 
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`googleapis/google-auth-library-java#1503
included in 1.28.0.
  • Loading branch information
@zhumin8
zhumin8 authored Oct 4, 2024
1 parent 3c00e8c commit ca3ec24
Show file tree
Hide file tree
Showing 2 changed files with 103 additions and 2 deletions.
23 changes: 21 additions & 2 deletions gax-java/gax/src/main/java/com/google/api/gax/rpc/ClientContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -210,7 +211,8 @@ public static ClientContext create(StubSettings settings) throws IOException {
if (transportChannelProvider.needsExecutor() && settings.getExecutorProvider() != null) {
transportChannelProvider = transportChannelProvider.withExecutor(backgroundExecutor);
}
Map<String, String> headers = getHeadersFromSettings(settings);

Map<String, String> headers = getHeaders(settings, credentials);
if (transportChannelProvider.needsHeaders()) {
transportChannelProvider = transportChannelProvider.withHeaders(headers);
}
Expand Down Expand Up @@ -318,8 +320,11 @@ static GdchCredentials getGdchCredentials(
/**
* Getting a header map from HeaderProvider and InternalHeaderProvider from settings with Quota
* Project Id.
*
* <p>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<String, String> getHeadersFromSettings(StubSettings settings) {
private static Map<String, String> getHeaders(StubSettings settings, Credentials credentials) {
// Resolve conflicts when merging headers from multiple sources
Map<String, String> userHeaders = settings.getHeaderProvider().getHeaders();
Map<String, String> internalHeaders = settings.getInternalHeaderProvider().getHeaders();
Expand All @@ -346,6 +351,20 @@ private static Map<String, String> getHeadersFromSettings(StubSettings settings)
effectiveHeaders.putAll(userHeaders);
effectiveHeaders.putAll(conflictResolution);

return appendCredentialTypeToHeaderIfPresent(effectiveHeaders, credentials);
}

private static Map<String, String> appendCredentialTypeToHeaderIfPresent(
Map<String, String> 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);
}

Expand Down
82 changes: 82 additions & 0 deletions gax-java/gax/src/test/java/com/google/api/gax/rpc/ClientContextTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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<String, String> 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<String, String> 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<String, String> 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<String, String> 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<String, String> 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";

Expand Down

0 comments on commit ca3ec24

Please sign in to comment.