fix: update User-Agent handling for resumable uploads (#2168)

BenWhitehead · web-flow · commit 665b714f421d · 2023-08-23T13:33:17.000-04:00
The apiary library manually injects ApplicationName into each request it constructs. For resumable uploads (start and PUTs) we are using the apiary clients http client directly. Update our registered interceptor to add the user-agent if it is null.
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
@@ -123,14 +123,15 @@ public HttpStorageRpc(StorageOptions options, JsonFactory jsonFactory) {
     this.options = options;
 
     // Open Census initialization
+    String applicationName = options.getApplicationName();
     CensusHttpModule censusHttpModule = new CensusHttpModule(tracer, IS_RECORD_EVENTS);
     initializer = censusHttpModule.getHttpRequestInitializer(initializer);
-    initializer = new InvocationIdInitializer(initializer);
+    initializer = new InvocationIdInitializer(initializer, applicationName);
     batchRequestInitializer = censusHttpModule.getHttpRequestInitializer(null);
     storage =
         new Storage.Builder(transport, jsonFactory, initializer)
             .setRootUrl(options.getHost())
-            .setApplicationName(options.getApplicationName())
+            .setApplicationName(applicationName)
             .build();
   }
 
@@ -140,9 +141,12 @@ public Storage getStorage() {
 
   private static final class InvocationIdInitializer implements HttpRequestInitializer {
     @Nullable HttpRequestInitializer initializer;
+    @Nullable private final String applicationName;
 
-    private InvocationIdInitializer(@Nullable HttpRequestInitializer initializer) {
+    private InvocationIdInitializer(
+        @Nullable HttpRequestInitializer initializer, @Nullable String applicationName) {
       this.initializer = initializer;
+      this.applicationName = applicationName;
     }
 
     @Override
@@ -151,15 +155,19 @@ public void initialize(HttpRequest request) throws IOException {
       if (this.initializer != null) {
         this.initializer.initialize(request);
       }
-      request.setInterceptor(new InvocationIdInterceptor(request.getInterceptor()));
+      request.setInterceptor(
+          new InvocationIdInterceptor(request.getInterceptor(), applicationName));
     }
   }
 
   private static final class InvocationIdInterceptor implements HttpExecuteInterceptor {
-    @Nullable HttpExecuteInterceptor interceptor;
+    @Nullable private final HttpExecuteInterceptor interceptor;
+    @Nullable private final String applicationName;
 
-    private InvocationIdInterceptor(@Nullable HttpExecuteInterceptor interceptor) {
+    private InvocationIdInterceptor(
+        @Nullable HttpExecuteInterceptor interceptor, @Nullable String applicationName) {
       this.interceptor = interceptor;
+      this.applicationName = applicationName;
     }
 
     @Override
@@ -183,6 +191,13 @@ public void intercept(HttpRequest request) throws IOException {
         }
         headers.set("x-goog-api-client", newValue);
         headers.set("x-goog-gcs-idempotency-token", invocationId);
+
+        String userAgent = headers.getUserAgent();
+        if ((userAgent == null
+            || userAgent.isEmpty()
+            || (applicationName != null && !userAgent.contains(applicationName)))) {
+          headers.setUserAgent(applicationName);
+        }
       }
     }
   }
diff --git a/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITUserAgentTest.java b/google-cloud-storage/src/test/java/com/google/cloud/storage/it/ITUserAgentTest.java
@@ -0,0 +1,78 @@
+/*
+ * Copyright 2023 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.it;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import com.google.api.client.http.HttpHeaders;
+import com.google.api.client.http.HttpRequest;
+import com.google.cloud.WriteChannel;
+import com.google.cloud.storage.BlobInfo;
+import com.google.cloud.storage.BucketInfo;
+import com.google.cloud.storage.DataGenerator;
+import com.google.cloud.storage.HttpStorageOptions;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageOptions;
+import com.google.cloud.storage.TransportCompatibility.Transport;
+import com.google.cloud.storage.it.runner.StorageITRunner;
+import com.google.cloud.storage.it.runner.annotations.Backend;
+import com.google.cloud.storage.it.runner.annotations.Inject;
+import com.google.cloud.storage.it.runner.annotations.SingleBackend;
+import com.google.cloud.storage.it.runner.annotations.StorageFixture;
+import com.google.cloud.storage.it.runner.registry.Generator;
+import com.google.common.collect.ImmutableList;
+import java.util.Objects;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(StorageITRunner.class)
+@SingleBackend(Backend.PROD)
+public final class ITUserAgentTest {
+
+  @Inject
+  @StorageFixture(Transport.HTTP)
+  public Storage storage;
+
+  @Inject public BucketInfo bucket;
+  @Inject public Generator generator;
+
+  @Test
+  public void userAgentIncludesGcloudJava_writer_http() throws Exception {
+    RequestAuditing requestAuditing = new RequestAuditing();
+    HttpStorageOptions options2 =
+        StorageOptions.http().setTransportOptions(requestAuditing).build();
+    try (Storage storage = options2.getService()) {
+      try (WriteChannel writer =
+          storage.writer(BlobInfo.newBuilder(bucket, generator.randomObjectName()).build())) {
+        writer.write(DataGenerator.base64Characters().genByteBuffer(13));
+      }
+    }
+
+    ImmutableList<String> userAgents =
+        requestAuditing.getRequests().stream()
+            .map(HttpRequest::getHeaders)
+            .map(HttpHeaders::getUserAgent)
+            .filter(Objects::nonNull)
+            .collect(ImmutableList.toImmutableList());
+
+    ImmutableList<String> found =
+        userAgents.stream()
+            .filter(ua -> ua.contains("gcloud-java/"))
+            .collect(ImmutableList.toImmutableList());
+    assertThat(found).hasSize(2); // one for the create session, and one for the PUT and finalize
+  }
+}
