fix: media upload to have applicationName as User-Agent (#2227)

* Upload to have applicationName as User-Agent Test that fails * request initializer to set user-agent header
googleapis · Jan 20, 2023 · 2595de0 · 2595de0
1 parent 0d1e5db
commit 2595de0
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 2 deletions.
diff --git a/.../src/main/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequest.java b/.../src/main/java/com/google/api/client/googleapis/services/AbstractGoogleClientRequest.java
@@ -351,15 +351,42 @@ public final MediaHttpUploader getMediaHttpUploader() {
    */
   protected final void initializeMediaUpload(AbstractInputStreamContent mediaContent) {
     HttpRequestFactory requestFactory = abstractGoogleClient.getRequestFactory();
+    String applicationName = abstractGoogleClient.getApplicationName();
+    HttpRequestInitializer requestInitializer =
+        mediaUploadRequestUserAgentInitializer(applicationName, requestFactory.getInitializer());
     this.uploader =
-        new MediaHttpUploader(
-            mediaContent, requestFactory.getTransport(), requestFactory.getInitializer());
+        new MediaHttpUploader(mediaContent, requestFactory.getTransport(), requestInitializer);
     this.uploader.setInitiationRequestMethod(requestMethod);
     if (httpContent != null) {
       this.uploader.setMetadata(httpContent);
     }
   }
 
+  private static HttpRequestInitializer mediaUploadRequestUserAgentInitializer(
+      final String applicationName, final HttpRequestInitializer originalInitializer) {
+    if (applicationName == null) {
+      return originalInitializer;
+    }
+    if (originalInitializer == null) {
+      return new HttpRequestInitializer() {
+        @Override
+        public void initialize(HttpRequest request) {
+          HttpHeaders headers = request.getHeaders();
+          headers.setUserAgent(applicationName);
+        }
+      };
+    } else {
+      return new HttpRequestInitializer() {
+        @Override
+        public void initialize(HttpRequest request) throws IOException {
+          originalInitializer.initialize(request);
+          HttpHeaders headers = request.getHeaders();
+          headers.setUserAgent(applicationName);
+        }
+      };
+    }
+  }
+
   /** Returns the media HTTP downloader or {@code null} for none. */
   public final MediaHttpDownloader getMediaHttpDownloader() {
     return downloader;

diff --git a/...ent/src/test/java/com/google/api/client/googleapis/services/AbstractGoogleClientTest.java b/...ent/src/test/java/com/google/api/client/googleapis/services/AbstractGoogleClientTest.java
@@ -16,6 +16,7 @@
 import com.google.api.client.googleapis.testing.services.MockGoogleClient;
 import com.google.api.client.googleapis.testing.services.MockGoogleClientRequest;
 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.HttpRequest;
 import com.google.api.client.http.HttpRequestInitializer;
@@ -33,6 +34,8 @@
 import com.google.api.client.util.Key;
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
+import java.util.ArrayList;
+import java.util.List;
 import junit.framework.TestCase;
 
 /**
@@ -132,6 +135,7 @@ private static class MediaTransport extends MockHttpTransport {
     int bytesUploaded;
     int contentLength = MediaHttpUploader.DEFAULT_CHUNK_SIZE;
     boolean contentLengthNotSpecified;
+    List<String> userAgentsRecorded = new ArrayList<>();
 
     protected MediaTransport() {}
 
@@ -169,6 +173,7 @@ public LowLevelHttpResponse execute() {
           String expectedContentRange = "bytes " + bytesRange + "/" + contentLength;
           assertEquals(expectedContentRange, getFirstHeaderValue("Content-Range"));
           bytesUploaded += MediaHttpUploader.DEFAULT_CHUNK_SIZE;
+          userAgentsRecorded.add(getFirstHeaderValue("User-Agent"));
 
           if (bytesUploaded == contentLength) {
             // Return 200 since the upload is complete.
@@ -207,6 +212,32 @@ public void testMediaUpload() throws Exception {
     assertEquals("somevalue", result.foo);
   }
 
+  public void testMediaUpload_applicationNameAsUserAgent() throws Exception {
+    MediaTransport fakeTransport = new MediaTransport();
+    String applicationName = "Foo/1.0 (BAR:Baz/1.0) XYZ/1.0";
+    AbstractGoogleClient client =
+        new MockGoogleClient.Builder(
+                fakeTransport, TEST_RESUMABLE_REQUEST_URL, "", JSON_OBJECT_PARSER, null)
+            .setApplicationName(applicationName)
+            .build();
+    InputStream is = new ByteArrayInputStream(new byte[MediaHttpUploader.DEFAULT_CHUNK_SIZE]);
+    InputStreamContent mediaContent = new InputStreamContent(TEST_CONTENT_TYPE, is);
+    mediaContent.setLength(MediaHttpUploader.DEFAULT_CHUNK_SIZE);
+    MockGoogleClientRequest<A> rq =
+        new MockGoogleClientRequest<A>(client, "POST", "", null, A.class);
+
+    rq.initializeMediaUpload(mediaContent);
+    MediaHttpUploader mediaHttpUploader = rq.getMediaHttpUploader();
+    mediaHttpUploader.upload(new GenericUrl(TEST_RESUMABLE_REQUEST_URL));
+
+    assertEquals(1, fakeTransport.userAgentsRecorded.size());
+    for (String userAgent : fakeTransport.userAgentsRecorded) {
+      assertTrue(
+          "UserAgent header does not have expected value in requests",
+          userAgent.contains(applicationName));
+    }
+  }
+
   private static class GZipCheckerInitializer implements HttpRequestInitializer {
 
     private boolean gzipDisabled;