Fix oauth test failures causing build to fail (#6030)

airbytehq · Sep 13, 2021 · 129a8fb · 129a8fb
1 parent 2ee79f2
commit 129a8fb
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 28 deletions.
diff --git a/airbyte-oauth/src/main/java/io/airbyte/oauth/google/GoogleOAuthFlow.java b/airbyte-oauth/src/main/java/io/airbyte/oauth/google/GoogleOAuthFlow.java
@@ -39,16 +39,17 @@
 import io.airbyte.validation.json.JsonValidationException;
 import java.io.IOException;
 import java.net.URI;
+import java.net.URISyntaxException;
 import java.net.URLEncoder;
 import java.net.http.HttpClient;
 import java.net.http.HttpClient.Version;
 import java.net.http.HttpRequest;
 import java.net.http.HttpResponse;
 import java.nio.charset.StandardCharsets;
-import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
+import org.apache.http.client.utils.URIBuilder;
 
 /**
  * Following docs from https://developers.google.com/identity/protocols/oauth2/web-server
@@ -61,7 +62,7 @@ public class GoogleOAuthFlow implements OAuthFlowImplementation {
   private final static String ACCESS_TOKEN_URL = "https://oauth2.googleapis.com/token";
 
   private final String scope;
-  private final List<String> googleQueryParameters;
+  private final Map<String, String> defaultQueryParams;
 
   private final ConfigRepository configRepository;
 
@@ -73,13 +74,15 @@ public GoogleOAuthFlow(ConfigRepository configRepository, String scope) {
   GoogleOAuthFlow(ConfigRepository configRepository, String scope, HttpClient httpClient) {
     this.configRepository = configRepository;
     this.httpClient = httpClient;
-    this.scope = UrlEncode(scope);
-    this.googleQueryParameters = List.of(
-        String.format("scope=%s", this.scope),
-        "access_type=offline",
-        "include_granted_scopes=true",
-        "response_type=code",
-        "prompt=consent");
+    this.scope = scope;
+    this.defaultQueryParams = ImmutableMap.<String, String>builder()
+        .put("scope", this.scope)
+        .put("access_type", "offline")
+        .put("include_granted_scopes", "true")
+        .put("response_type", "code")
+        .put("prompt", "consent")
+        .build();
+
   }
 
   @Override
@@ -96,18 +99,18 @@ public String getDestinationConsentUrl(UUID workspaceId, UUID destinationDefinit
   }
 
   private String getConsentUrl(UUID definitionId, String clientId, String redirectUrl) {
-    final StringBuilder result = new StringBuilder(CONSENT_URL)
-        .append("?");
-    for (String queryParameter : googleQueryParameters) {
-      result.append(queryParameter).append("&");
+    try {
+      URIBuilder uriBuilder = new URIBuilder(CONSENT_URL)
+          .addParameter("state", definitionId.toString())
+          .addParameter("client_id", clientId)
+          .addParameter("redirect_uri", redirectUrl);
+      for (Map.Entry<String, String> queryParameter : defaultQueryParams.entrySet()) {
+        uriBuilder.addParameter(queryParameter.getKey(), queryParameter.getValue());
+      }
+      return uriBuilder.toString();
+    } catch (URISyntaxException e) {
+      throw new IllegalArgumentException(e);
     }
-    return result
-        // TODO state should be randomly generated, and the 2nd step of oauth should verify its value
-        // matches the initially generated state value
-        .append("state=").append(definitionId.toString()).append("&")
-        .append("client_id=").append(clientId).append("&")
-        .append("redirect_uri=").append(redirectUrl)
-        .toString();
   }
 
   @Override

diff --git a/airbyte-oauth/src/test/java/io/airbyte/oauth/google/GoogleOAuthFlowTest.java b/airbyte-oauth/src/test/java/io/airbyte/oauth/google/GoogleOAuthFlowTest.java
@@ -39,8 +39,10 @@
 import io.airbyte.config.persistence.ConfigRepository;
 import io.airbyte.validation.json.JsonValidationException;
 import java.io.IOException;
+import java.net.URLEncoder;
 import java.net.http.HttpClient;
 import java.net.http.HttpResponse;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.List;
@@ -55,8 +57,8 @@ public class GoogleOAuthFlowTest {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(GoogleOAuthFlowTest.class);
   private static final Path CREDENTIALS_PATH = Path.of("secrets/credentials.json");
-  private static final String REDIRECT_URL = "https%3A//airbyte.io";
-  private static final String SCOPE = "https%3A//www.googleapis.com/auth/analytics.readonly";
+  private static final String REDIRECT_URL = "https://airbyte.io";
+  private static final String SCOPE = "https://www.googleapis.com/auth/analytics.readonly";
 
   private HttpClient httpClient;
   private ConfigRepository configRepository;
@@ -108,11 +110,11 @@ public void testGetSourceConsentUrl() throws IOException, ConfigNotFoundExceptio
             .build()))));
     final String actualSourceUrl = googleOAuthFlow.getSourceConsentUrl(workspaceId, definitionId, REDIRECT_URL);
     final String expectedSourceUrl = String.format(
-        "https://accounts.google.com/o/oauth2/v2/auth?scope=%s&access_type=offline&include_granted_scopes=true&response_type=code&prompt=consent&state=%s&client_id=%s&redirect_uri=%s",
-        SCOPE,
+        "https://accounts.google.com/o/oauth2/v2/auth?state=%s&client_id=%s&redirect_uri=%s&scope=%s&access_type=offline&include_granted_scopes=true&response_type=code&prompt=consent",
         definitionId,
         getClientId(),
-        REDIRECT_URL);
+        urlEncode(REDIRECT_URL),
+        urlEncode(SCOPE));
     LOGGER.info(expectedSourceUrl);
     assertEquals(expectedSourceUrl, actualSourceUrl);
   }
@@ -126,13 +128,16 @@ public void testGetDestinationConsentUrl() throws IOException, ConfigNotFoundExc
         .withConfiguration(Jsons.jsonNode(ImmutableMap.builder()
             .put("client_id", getClientId())
             .build()))));
+    // It would be better to make this comparison agnostic of the order of query params but the URI
+    // class' equals() method
+    // considers URLs with different qparam orders different URIs..
     final String actualDestinationUrl = googleOAuthFlow.getDestinationConsentUrl(workspaceId, definitionId, REDIRECT_URL);
     final String expectedDestinationUrl = String.format(
-        "https://accounts.google.com/o/oauth2/v2/auth?scope=%s&access_type=offline&include_granted_scopes=true&response_type=code&prompt=consent&state=%s&client_id=%s&redirect_uri=%s",
-        SCOPE,
+        "https://accounts.google.com/o/oauth2/v2/auth?state=%s&client_id=%s&redirect_uri=%s&scope=%s&access_type=offline&include_granted_scopes=true&response_type=code&prompt=consent",
         definitionId,
         getClientId(),
-        REDIRECT_URL);
+        urlEncode(REDIRECT_URL),
+        urlEncode(SCOPE));
     LOGGER.info(expectedDestinationUrl);
     assertEquals(expectedDestinationUrl, actualDestinationUrl);
   }
@@ -201,4 +206,8 @@ private String getClientId() throws IOException {
     }
   }
 
+  private String urlEncode(String s) {
+    return URLEncoder.encode(s, StandardCharsets.UTF_8);
+  }
+
 }