Repository: googleapis/java-storage
Pull request #1309 (2022-03-25 16:11:14)

fix: add gccl-invocation-id interceptor
Joint effort with @BenWhitehead who is the author of the original design, thank you Ben!
- [x] Ensure the tests and linter pass
- [x] Add option to disable the invocation id
- [x] Make sure this header isn't passed in for signed URLs
- [x] Appropriate docs were updated (if necessary)




Patch:
diff --git a/google-cloud-storage/pom.xml b/google-cloud-storage/pom.xml
index db88133b..b82f9f0f 100644
--- a/google-cloud-storage/pom.xml
+++ b/google-cloud-storage/pom.xml
@@ -92,6 +92,10 @@
       <groupId>com.fasterxml.jackson.core</groupId>
       <artifactId>jackson-core</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.google.code.findbugs</groupId>
+      <artifactId>jsr305</artifactId>
+    </dependency>
 
     <!-- Test dependencies -->
     <dependency>
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/Retrying.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/Retrying.java
index 3daf2c17..51ce17bf 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/Retrying.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/Retrying.java
@@ -22,6 +22,7 @@ import com.google.api.core.ApiClock;
 import com.google.api.gax.retrying.ResultRetryAlgorithm;
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.cloud.RetryHelper.RetryHelperException;
+import com.google.cloud.storage.spi.v1.HttpRpcContext;
 import java.util.concurrent.Callable;
 import java.util.function.Function;
 
@@ -47,11 +48,15 @@ final class Retrying {
    */
   static <T, U> U run(
       StorageOptions options, ResultRetryAlgorithm<?> algorithm, Callable<T> c, Function<T, U> f) {
+    HttpRpcContext httpRpcContext = HttpRpcContext.getInstance();
     try {
+      httpRpcContext.newInvocationId();
       T result = runWithRetries(c, options.getRetrySettings(), algorithm, options.getClock());
       return result == null ? null : f.apply(result);
     } catch (RetryHelperException e) {
       throw StorageException.coalesce(e);
+    } finally {
+      httpRpcContext.clearInvocationId();
     }
   }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageOptions.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageOptions.java
index f3de842f..49f1b9e0 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageOptions.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/StorageOptions.java
@@ -38,8 +38,9 @@ public class StorageOptions extends ServiceOptions<Storage, StorageOptions> {
   private static final String GCS_SCOPE = "https://www.googleapis.com/auth/devstorage.full_control";
   private static final Set<String> SCOPES = ImmutableSet.of(GCS_SCOPE);
   private static final String DEFAULT_HOST = "https://storage.googleapis.com";
-
+  private static final boolean DEFAULT_INCLUDE_INVOCATION_ID = true;
   private final RetryAlgorithmManager retryAlgorithmManager;
+  private final boolean includeInvocationId;
 
   public static class DefaultStorageFactory implements StorageFactory {
 
@@ -64,11 +65,13 @@ public class StorageOptions extends ServiceOptions<Storage, StorageOptions> {
   public static class Builder extends ServiceOptions.Builder<Storage, StorageOptions, Builder> {
 
     private StorageRetryStrategy storageRetryStrategy;
+    private boolean includeInvocationId;
 
     private Builder() {}
 
     private Builder(StorageOptions options) {
       super(options);
+      this.includeInvocationId = options.includeInvocationId;
     }
 
     @Override
@@ -93,6 +96,17 @@ public class StorageOptions extends ServiceOptions<Storage, StorageOptions> {
       return this;
     }
 
+    /**
+     * Override default enablement of invocation id added to x-goog-api-client header.
+     *
+     * @param includeInvocationId a boolean to change enablement of invocation id
+     * @return the builder
+     */
+    public Builder setIncludeInvocationId(boolean includeInvocationId) {
+      this.includeInvocationId = includeInvocationId;
+      return this;
+    }
+
     @Override
     public StorageOptions build() {
       return new StorageOptions(this, new StorageDefaults());
@@ -105,6 +119,7 @@ public class StorageOptions extends ServiceOptions<Storage, StorageOptions> {
         new RetryAlgorithmManager(
             MoreObjects.firstNonNull(
                 builder.storageRetryStrategy, serviceDefaults.getStorageRetryStrategy()));
+    this.includeInvocationId = builder.includeInvocationId;
   }
 
   private static class StorageDefaults implements ServiceDefaults<Storage, StorageOptions> {
@@ -127,6 +142,10 @@ public class StorageOptions extends ServiceOptions<Storage, StorageOptions> {
     public StorageRetryStrategy getStorageRetryStrategy() {
       return StorageRetryStrategy.getDefaultStorageRetryStrategy();
     }
+
+    public boolean isIncludeInvocationId() {
+      return DEFAULT_INCLUDE_INVOCATION_ID;
+    }
   }
 
   public static HttpTransportOptions getDefaultHttpTransportOptions() {
@@ -153,6 +172,11 @@ public class StorageOptions extends ServiceOptions<Storage, StorageOptions> {
     return retryAlgorithmManager;
   }
 
+  /** Returns if Invocation ID is enabled and transmitted through x-goog-api-client header. */
+  public boolean isIncludeInvocationId() {
+    return includeInvocationId;
+  }
+
   /** Returns a default {@code StorageOptions} instance. */
   public static StorageOptions getDefaultInstance() {
     return newBuilder().build();
@@ -180,6 +204,8 @@ public class StorageOptions extends ServiceOptions<Storage, StorageOptions> {
   }
 
   public static Builder newBuilder() {
-    return new Builder().setHost(DEFAULT_HOST);
+    return new Builder()
+        .setHost(DEFAULT_HOST)
+        .setIncludeInvocationId(DEFAULT_INCLUDE_INVOCATION_ID);
   }
 }
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpRpcContext.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpRpcContext.java
new file mode 100644
index 00000000..a62e7f99
--- /dev/null
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpRpcContext.java
@@ -0,0 +1,72 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.storage.spi.v1;
+
+import com.google.api.core.InternalApi;
+import java.util.UUID;
+import java.util.function.Supplier;
+import javax.annotation.Nullable;
+
+@InternalApi
+public final class HttpRpcContext {
+
+  private static final Object GET_INSTANCE_LOCK = new Object();
+
+  private static volatile HttpRpcContext instance;
+
+  private final ThreadLocal<UUID> invocationId;
+  private final Supplier<UUID> supplier;
+
+  HttpRpcContext(Supplier<UUID> randomUUID) {
+    this.invocationId = new InheritableThreadLocal<>();
+    this.supplier = randomUUID;
+  }
+
+  @InternalApi
+  @Nullable
+  public UUID getInvocationId() {
+    return invocationId.get();
+  }
+
+  @InternalApi
+  public UUID newInvocationId() {
+    invocationId.set(supplier.get());
+    return getInvocationId();
+  }
+
+  @InternalApi
+  public void clearInvocationId() {
+    invocationId.remove();
+  }
+
+  @InternalApi
+  public static HttpRpcContext init() {
+    return new HttpRpcContext(UUID::randomUUID);
+  }
+
+  @InternalApi
+  public static HttpRpcContext getInstance() {
+    if (instance == null) {
+      synchronized (GET_INSTANCE_LOCK) {
+        if (instance == null) {
+          instance = init();
+        }
+      }
+    }
+    return instance;
+  }
+}
diff --git a/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java b/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java
index d0c1ff8a..e990a004 100644
--- a/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java
+++ b/google-cloud-storage/src/main/java/com/google/cloud/storage/spi/v1/HttpStorageRpc.java
@@ -18,6 +18,7 @@ package com.google.cloud.storage.spi.v1;
 
 import static com.google.common.base.MoreObjects.firstNonNull;
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
 import static java.net.HttpURLConnection.HTTP_NOT_FOUND;
 
 import com.google.api.client.googleapis.batch.BatchRequest;
@@ -26,6 +27,7 @@ import com.google.api.client.googleapis.json.GoogleJsonError;
 import com.google.api.client.http.ByteArrayContent;
 import com.google.api.client.http.EmptyContent;
 import com.google.api.client.http.GenericUrl;
+import com.google.api.client.http.HttpExecuteInterceptor;
 import com.google.api.client.http.HttpHeaders;
 import com.google.api.client.http.HttpRequest;
 import com.google.api.client.http.HttpRequestFactory;
@@ -86,6 +88,8 @@ import java.util.LinkedList;
 import java.util.List;
 import java.util.Locale;
 import java.util.Map;
+import java.util.UUID;
+import javax.annotation.Nullable;
 
 public class HttpStorageRpc implements StorageRpc {
   public static final String DEFAULT_PROJECTION = "full";
@@ -114,6 +118,9 @@ public class HttpStorageRpc implements StorageRpc {
     // Open Census initialization
     censusHttpModule = new CensusHttpModule(tracer, true);
     initializer = censusHttpModule.getHttpRequestInitializer(initializer);
+    if (options.isIncludeInvocationId()) {
+      initializer = new InvocationIdInitializer(initializer);
+    }
     batchRequestInitializer = censusHttpModule.getHttpRequestInitializer(null);
     storage =
         new Storage.Builder(transport, new JacksonFactory(), initializer)
@@ -122,6 +129,54 @@ public class HttpStorageRpc implements StorageRpc {
             .build();
   }
 
+  private static final class InvocationIdInitializer implements HttpRequestInitializer {
+    @Nullable HttpRequestInitializer initializer;
+
+    private InvocationIdInitializer(@Nullable HttpRequestInitializer initializer) {
+      this.initializer = initializer;
+    }
+
+    @Override
+    public void initialize(HttpRequest request) throws IOException {
+      checkNotNull(request);
+      if (this.initializer != null) {
+        this.initializer.initialize(request);
+      }
+      request.setInterceptor(new InvocationIdInterceptor(request.getInterceptor()));
+    }
+  }
+
+  private static final class InvocationIdInterceptor implements HttpExecuteInterceptor {
+    @Nullable HttpExecuteInterceptor interceptor;
+
+    private InvocationIdInterceptor(@Nullable HttpExecuteInterceptor interceptor) {
+      this.interceptor = interceptor;
+    }
+
+    @Override
+    public void intercept(HttpRequest request) throws IOException {
+      checkNotNull(request);
+      if (this.interceptor != null) {
+        this.interceptor.intercept(request);
+      }
+      UUID invocationId = HttpRpcContext.getInstance().getInvocationId();
+      final String signatureKey = "Signature="; // For V2 and V4 signedURLs
+      final String builtURL = request.getUrl().build();
+      if (invocationId != null && !builtURL.contains(signatureKey)) {
+        HttpHeaders headers = request.getHeaders();
+        String existing = (String) headers.get("x-goog-api-client");
+        String invocationEntry = "gccl-invocation-id/" + invocationId;
+        final String newValue;
+        if (existing != null && !existing.isEmpty()) {
+          newValue = existing + " " + invocationEntry;
+        } else {
+          newValue = invocationEntry;
+        }
+        headers.set("x-goog-api-client", newValue);
+      }
+    }
+  }
+
   private class DefaultRpcBatch implements RpcBatch {
 
     // Batch size is limited as, due to some current service implementation details, the service


Test patch:
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageOptionsTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageOptionsTest.java
index 5c9c708c..93d36b6e 100644
--- a/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageOptionsTest.java
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/StorageOptionsTest.java
@@ -17,6 +17,8 @@
 package com.google.cloud.storage;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 
 import com.google.cloud.TransportOptions;
 import org.easymock.EasyMock;
@@ -65,4 +67,18 @@ public class StorageOptionsTest {
 
     assertThat(opts1.getHost()).isEqualTo("https://storage.googleapis.com");
   }
+
+  @Test
+  public void testDefaultInvocationId() {
+    StorageOptions opts1 = StorageOptions.getDefaultInstance();
+
+    assertTrue(opts1.isIncludeInvocationId());
+  }
+
+  @Test
+  public void testDisableInvocationId() {
+    StorageOptions opts1 = StorageOptions.newBuilder().setIncludeInvocationId(false).build();
+
+    assertFalse(opts1.isIncludeInvocationId());
+  }
 }
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/spi/v1/HttpRpcContextTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/spi/v1/HttpRpcContextTest.java
new file mode 100644
index 00000000..f1e93f5c
--- /dev/null
+++ b/google-cloud-storage/src/test/java/com/google/cloud/storage/spi/v1/HttpRpcContextTest.java
@@ -0,0 +1,195 @@
+/*
+ * Copyright 2022 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *       http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.storage.spi.v1;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+
+import com.google.api.client.http.HttpTransport;
+import com.google.api.client.http.LowLevelHttpRequest;
+import com.google.api.client.http.LowLevelHttpResponse;
+import com.google.api.client.testing.http.MockLowLevelHttpResponse;
+import com.google.cloud.TransportOptions;
+import com.google.cloud.Tuple;
+import com.google.cloud.WriteChannel;
+import com.google.cloud.http.HttpTransportOptions;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageOptions;
+import com.google.common.collect.ImmutableList;
+import java.io.IOException;
+import java.net.URL;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Optional;
+import java.util.UUID;
+import org.junit.Test;
+
+public class HttpRpcContextTest {
+  @Test
+  public void testNewInvocationId() {
+    UUID uuid = UUID.fromString("28220dff-1e8b-4770-9e10-022c2a99d8f3");
+    HttpRpcContext testContext = new HttpRpcContext(() -> uuid);
+
+    assertThat(testContext.newInvocationId()).isEqualTo(uuid);
+    assertThat(testContext.getInvocationId()).isEqualTo(uuid);
+    // call again to ensure the id is consistent with our supplier
+    assertThat(testContext.newInvocationId()).isEqualTo(uuid);
+    assertThat(testContext.getInvocationId()).isEqualTo(uuid);
+  }
+
+  @Test
+  public void testInvocationIdIsPassedThrough() {
+    MockLowLevelHttpResponse response =
+        new MockLowLevelHttpResponse()
+            .setContentType("application/json")
+            .setContent(
+                "{\n"
+                    + "  \"kind\": \"storage#serviceAccount\",\n"
+                    + "  \"email_address\": \"service-234234@gs-project-accounts.iam.gserviceaccount.com\"\n"
+                    + "}\n")
+            .setStatusCode(200);
+    AuditingHttpTransport transport = new AuditingHttpTransport(response);
+    TransportOptions transportOptions =
+        HttpTransportOptions.newBuilder().setHttpTransportFactory(() -> transport).build();
+    Storage service =
+        StorageOptions.getDefaultInstance()
+            .toBuilder()
+            .setTransportOptions(transportOptions)
+            .build()
+            .getService();
+    service.getServiceAccount("test-project");
+    Optional<Tuple<String, String>> anyXGoogApiClientWithGcclInvocationId =
+        transport.getAddHeaderCalls().stream()
+            .filter(t -> "x-goog-api-client".equals(t.x()) && t.y().contains("gccl-invocation-id/"))
+            .findFirst();
+    assertTrue(anyXGoogApiClientWithGcclInvocationId.isPresent());
+    assertThat(transport.getBuildRequestCalls()).hasSize(1);
+  }
+
+  @Test
+  public void testInvocationIdIsNotPassedThroughWhenDisabled() {
+    MockLowLevelHttpResponse response =
+        new MockLowLevelHttpResponse()
+            .setContentType("application/json")
+            .setContent(
+                "{\n"
+                    + "  \"kind\": \"storage#serviceAccount\",\n"
+                    + "  \"email_address\": \"service-234234@gs-project-accounts.iam.gserviceaccount.com\"\n"
+                    + "}\n")
+            .setStatusCode(200);
+    AuditingHttpTransport transport = new AuditingHttpTransport(response);
+    TransportOptions transportOptions =
+        HttpTransportOptions.newBuilder().setHttpTransportFactory(() -> transport).build();
+    Storage service =
+        StorageOptions.getDefaultInstance()
+            .toBuilder()
+            .setTransportOptions(transportOptions)
+            .setIncludeInvocationId(false)
+            .build()
+            .getService();
+    service.getServiceAccount("test-project");
+    Optional<Tuple<String, String>> anyXGoogApiClientWithGcclInvocationId =
+        transport.getAddHeaderCalls().stream()
+            .filter(t -> "x-goog-api-client".equals(t.x()) && t.y().contains("gccl-invocation-id/"))
+            .findFirst();
+
+    assertFalse(anyXGoogApiClientWithGcclInvocationId.isPresent());
+    assertThat(transport.getBuildRequestCalls()).hasSize(1);
+  }
+
+  @Test
+  public void testInvocationIdNotInSignedURL_v2() throws IOException {
+    URL signedUrlV2 =
+        new URL(
+            "http://www.test.com/test-bucket/test1.txt?GoogleAccessId=testClient-test@test.com&Expires=1553839761&Signature=MJUBXAZ7");
+    doTestInvocationIdNotInSignedURL(signedUrlV2);
+  }
+
+  @Test
+  public void testInvocationIdNotInSignedURL_v4() throws IOException {
+    URL signedUrlV4 =
+        new URL(
+            "http://www.test.com/test-bucket/test1.txt?X-Goog-Algorithm=&X-Goog-Credential=&X-Goog-Date=&X-Goog-Expires=&X-Goog-SignedHeaders=&X-Goog-Signature=MJUBXAZ7");
+    doTestInvocationIdNotInSignedURL(signedUrlV4);
+  }
+
+  private void doTestInvocationIdNotInSignedURL(URL signedUrl) throws IOException {
+    MockLowLevelHttpResponse response =
+        new MockLowLevelHttpResponse()
+            .setContentType("text/plain")
+            .setHeaderNames(ImmutableList.of("Location"))
+            .setHeaderValues(ImmutableList.of("http://test"))
+            .setStatusCode(201);
+    AuditingHttpTransport transport = new AuditingHttpTransport(response);
+    TransportOptions transportOptions =
+        HttpTransportOptions.newBuilder().setHttpTransportFactory(() -> transport).build();
+    Storage service =
+        StorageOptions.getDefaultInstance()
+            .toBuilder()
+            .setTransportOptions(transportOptions)
+            .build()
+            .getService();
+    WriteChannel writer = service.writer(signedUrl);
+    writer.write(ByteBuffer.wrap("hello".getBytes(StandardCharsets.UTF_8)));
+    Optional<Tuple<String, String>> anyXGoogApiClientWithGcclInvocationId =
+        transport.getAddHeaderCalls().stream()
+            .filter(t -> "x-goog-api-client".equals(t.x()) && t.y().contains("gccl-invocation-id/"))
+            .findFirst();
+    assertFalse(anyXGoogApiClientWithGcclInvocationId.isPresent());
+    assertThat(transport.getBuildRequestCalls()).hasSize(1);
+  }
+
+  private static final class AuditingHttpTransport extends HttpTransport {
+    private final LowLevelHttpResponse response;
+    private final List<Tuple<String, String>> buildRequestCalls;
+    private final List<Tuple<String, String>> addHeaderCalls;
+
+    private AuditingHttpTransport(LowLevelHttpResponse response) {
+      this.response = response;
+      this.buildRequestCalls = Collections.synchronizedList(new ArrayList<>());
+      this.addHeaderCalls = Collections.synchronizedList(new ArrayList<>());
+    }
+
+    public List<Tuple<String, String>> getBuildRequestCalls() {
+      return ImmutableList.copyOf(buildRequestCalls);
+    }
+
+    public List<Tuple<String, String>> getAddHeaderCalls() {
+      return ImmutableList.copyOf(addHeaderCalls);
+    }
+
+    @Override
+    protected LowLevelHttpRequest buildRequest(String method, String url) {
+      buildRequestCalls.add(Tuple.of(method, url));
+      return new LowLevelHttpRequest() {
+        @Override
+        public void addHeader(String name, String value) {
+          addHeaderCalls.add(Tuple.of(name, value));
+        }
+
+        @Override
+        public LowLevelHttpResponse execute() {
+          return response;
+        }
+      };
+    }
+  }
+}