From f42a7078522461018da96196a67657809d28b15b Mon Sep 17 00:00:00 2001
From: Aristide Niyungeko <2230766+aristide-n@users.noreply.github.com>
Date: Mon, 16 Aug 2021 16:10:23 -0700
Subject: [PATCH] feat: support downloading an emulator from an access
 controlled URL (#513)

* feat: support downloading an emulator from an access controlled URL

fixes googleapis/java-datastore#376
used in googleapis/java-datastore#492

* fix format
---
 .../google/cloud/testing/BaseEmulatorHelper.java   | 14 +++++++++++++-
 .../cloud/testing/BaseEmulatorHelperTest.java      |  4 +++-
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/google-cloud-core/src/main/java/com/google/cloud/testing/BaseEmulatorHelper.java b/google-cloud-core/src/main/java/com/google/cloud/testing/BaseEmulatorHelper.java
index 9a0898960a..9679c6299c 100644
--- a/google-cloud-core/src/main/java/com/google/cloud/testing/BaseEmulatorHelper.java
+++ b/google-cloud-core/src/main/java/com/google/cloud/testing/BaseEmulatorHelper.java
@@ -38,6 +38,7 @@
 import java.net.HttpURLConnection;
 import java.net.ServerSocket;
 import java.net.URL;
+import java.net.URLConnection;
 import java.nio.channels.Channels;
 import java.nio.channels.ReadableByteChannel;
 import java.nio.file.Files;
@@ -316,6 +317,7 @@ protected static class DownloadableEmulatorRunner implements EmulatorRunner {
     private final String md5CheckSum;
     private final URL downloadUrl;
     private final String fileName;
+    private String accessToken;
     private Process process;
     private static final Logger log = Logger.getLogger(DownloadableEmulatorRunner.class.getName());
 
@@ -328,6 +330,12 @@ public DownloadableEmulatorRunner(
       this.fileName = splitUrl[splitUrl.length - 1];
     }
 
+    public DownloadableEmulatorRunner(
+        List<String> commandText, URL downloadUrl, String md5CheckSum, String accessToken) {
+      this(commandText, downloadUrl, md5CheckSum);
+      this.accessToken = accessToken;
+    }
+
     @Override
     public boolean isAvailable() {
       try {
@@ -420,7 +428,11 @@ private File downloadZipFile() throws IOException {
         if (log.isLoggable(Level.FINE)) {
           log.fine("Fetching emulator");
         }
-        ReadableByteChannel rbc = Channels.newChannel(downloadUrl.openStream());
+        URLConnection urlConnection = downloadUrl.openConnection();
+        if (accessToken != null) {
+          urlConnection.setRequestProperty("Authorization", "Bearer " + accessToken);
+        }
+        ReadableByteChannel rbc = Channels.newChannel(urlConnection.getInputStream());
         try (FileOutputStream fos = new FileOutputStream(zipFile)) {
           fos.getChannel().transferFrom(rbc, 0, Long.MAX_VALUE);
         }
diff --git a/google-cloud-core/src/test/java/com/google/cloud/testing/BaseEmulatorHelperTest.java b/google-cloud-core/src/test/java/com/google/cloud/testing/BaseEmulatorHelperTest.java
index 3a799e2044..b002a6198b 100644
--- a/google-cloud-core/src/test/java/com/google/cloud/testing/BaseEmulatorHelperTest.java
+++ b/google-cloud-core/src/test/java/com/google/cloud/testing/BaseEmulatorHelperTest.java
@@ -108,6 +108,7 @@ public void testEmulatorHelperDownloadWithRetries()
     String mockInputStream = "mockInputStream";
     String mockProtocol = "mockProtocol";
     String mockFile = "mockFile";
+    String mockAccessToken = "mockAccessToken";
     String mockCommandText = "mockCommandText";
 
     MockURLStreamHandler mockURLStreamHandler = EasyMock.createMock(MockURLStreamHandler.class);
@@ -119,6 +120,7 @@ public void testEmulatorHelperDownloadWithRetries()
     EasyMock.expect(mockURLConnection.getInputStream())
         .andReturn(new ByteArrayInputStream(mockInputStream.getBytes()))
         .anyTimes();
+    mockURLConnection.setRequestProperty("Authorization", "Bearer " + mockAccessToken);
     EasyMock.expect(mockURLStreamHandler.openConnection(EasyMock.anyObject(URL.class)))
         .andThrow(new EOFException())
         .times(1);
@@ -130,7 +132,7 @@ public void testEmulatorHelperDownloadWithRetries()
     URL url = new URL(mockProtocol, null, 0, mockFile, mockURLStreamHandler);
     BaseEmulatorHelper.DownloadableEmulatorRunner runner =
         new BaseEmulatorHelper.DownloadableEmulatorRunner(
-            ImmutableList.of(mockCommandText), url, null);
+            ImmutableList.of(mockCommandText), url, null, mockAccessToken);
 
     File cachedFile = new File(System.getProperty("java.io.tmpdir"), mockExternalForm);
     cachedFile.delete(); // Clear the cached version so we're always testing the download