[🍒 ] Add USER_AGENT, HTTP Basic Auth option to OAuthHandler and fix content type header

cdap-app-templates/cdap-etl/cdap-data-pipeline-base/src/main/java/io/cdap/cdap/datapipeline/oauth/OAuthProvider.java

@@ -26,15 +26,24 @@ public class OAuthProvider {
   private final String tokenRefreshURL;
   @Nullable
   private final OAuthClientCredentials clientCreds;
+  private final CredentialEncodingStrategy strategy;
+
+  // Optional string to send as a USER_AGENT header
+  @Nullable
+  private final String userAgent;
 
   public OAuthProvider(String name,
                        String loginURL,
                        String tokenRefreshURL,
-                       @Nullable OAuthClientCredentials clientCreds) {
+                       @Nullable OAuthClientCredentials clientCreds,
+                       @Nullable CredentialEncodingStrategy strategy,
+                       @Nullable String userAgent) {
     this.name = name;
     this.loginURL = loginURL;
     this.tokenRefreshURL = tokenRefreshURL;
     this.clientCreds = clientCreds;
+    this.strategy = strategy;
+    this.userAgent = userAgent;
   }
 
   public String getName() {
@@ -54,6 +63,22 @@ public OAuthClientCredentials getClientCredentials() {
     return clientCreds;
   }
 
+  public CredentialEncodingStrategy getCredentialEncodingStrategy() {
+    return strategy;
+  }
+
+  @Nullable
+  public String getUserAgent() {
+    return userAgent;
+  }
+
+  public enum CredentialEncodingStrategy {
+    // (default) Sends client ID & secret as part of the POST request body
+    FORM_BODY,
+    // Sends client ID & secret as part of a HTTP Basic Auth header
+    BASIC_AUTH,
+  }
+
   public static Builder newBuilder() {
     return new Builder();
   }
@@ -66,6 +91,8 @@ public static class Builder {
     private String loginURL;
     private String tokenRefreshURL;
     private OAuthClientCredentials clientCreds;
+    private CredentialEncodingStrategy strategy;
+    private String userAgent;
 
     public Builder() {}
 
@@ -89,11 +116,25 @@ public Builder withClientCredentials(@Nullable OAuthClientCredentials clientCred
       return this;
     }
 
+    public Builder withCredentialEncodingStrategy(@Nullable CredentialEncodingStrategy strategy) {
+      this.strategy = strategy;
+      return this;
+    }
+
+    public Builder withUserAgent(@Nullable String userAgent) {
+      this.userAgent = userAgent;
+      return this;
+    }
+
     public OAuthProvider build() {
       Preconditions.checkNotNull(name, "OAuth provider name missing");
       Preconditions.checkNotNull(loginURL, "Login URL missing");
       Preconditions.checkNotNull(tokenRefreshURL, "Token refresh URL missing");
-      return new OAuthProvider(name, loginURL, tokenRefreshURL, clientCreds);
+      // Default to FORM_BODY strategy
+      if (strategy == null) {
+        this.strategy = CredentialEncodingStrategy.FORM_BODY;
+      }
+      return new OAuthProvider(name, loginURL, tokenRefreshURL, clientCreds, strategy, userAgent);
     }
   }
 }

cdap-app-templates/cdap-etl/cdap-data-pipeline-base/src/main/java/io/cdap/cdap/datapipeline/oauth/PutOAuthProviderRequest.java

@@ -25,12 +25,22 @@ public class PutOAuthProviderRequest {
   private final String tokenRefreshURL;
   private final String clientId;
   private final String clientSecret;
+  private final OAuthProvider.CredentialEncodingStrategy strategy;
+  private final String userAgent;
 
-  public PutOAuthProviderRequest(String loginURL, String tokenRefreshURL, String clientId, String clientSecret) {
+  public PutOAuthProviderRequest(
+          String loginURL,
+          String tokenRefreshURL,
+          String clientId,
+          String clientSecret,
+          OAuthProvider.CredentialEncodingStrategy strategy,
+          String userAgent) {
     this.loginURL = loginURL;
     this.tokenRefreshURL = tokenRefreshURL;
     this.clientId = clientId;
     this.clientSecret = clientSecret;
+    this.strategy = strategy;
+    this.userAgent = userAgent;
   }
 
   public String getLoginURL() {
@@ -48,4 +58,12 @@ public String getClientId() {
   public String getClientSecret() {
     return clientSecret;
   }
+
+  public OAuthProvider.CredentialEncodingStrategy getCredentialEncodingStrategy() {
+    return strategy;
+  }
+
+  public String getUserAgent() {
+    return userAgent;
+  }
 }

cdap-app-templates/cdap-etl/cdap-data-pipeline-base/src/main/java/io/cdap/cdap/datapipeline/service/OAuthHandler.java

@@ -28,6 +28,7 @@
 import io.cdap.cdap.datapipeline.oauth.GetAccessTokenResponse;
 import io.cdap.cdap.datapipeline.oauth.OAuthClientCredentials;
 import io.cdap.cdap.datapipeline.oauth.OAuthProvider;
+import io.cdap.cdap.datapipeline.oauth.OAuthProvider.CredentialEncodingStrategy;
 import io.cdap.cdap.datapipeline.oauth.OAuthRefreshToken;
 import io.cdap.cdap.datapipeline.oauth.OAuthStore;
 import io.cdap.cdap.datapipeline.oauth.OAuthStoreException;
@@ -43,6 +44,7 @@
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
+import java.util.Base64;
 import java.util.Optional;
 import javax.ws.rs.DefaultValue;
 import javax.ws.rs.GET;
@@ -115,6 +117,8 @@ public void putOAuthProvider(HttpServiceRequest request, HttpServiceResponder re
         PutOAuthProviderRequest putOAuthProviderRequest = GSON.fromJson(
             StandardCharsets.UTF_8.decode(request.getContent()).toString(),
             PutOAuthProviderRequest.class);
+        CredentialEncodingStrategy strategy = putOAuthProviderRequest.getCredentialEncodingStrategy();
+        String userAgent = putOAuthProviderRequest.getUserAgent();
         // Validate URLs
         URL loginURL = new URL(putOAuthProviderRequest.getLoginURL());
         URL tokenRefreshURL = new URL(putOAuthProviderRequest.getTokenRefreshURL());
@@ -132,6 +136,8 @@ public void putOAuthProvider(HttpServiceRequest request, HttpServiceResponder re
                                               .withLoginURL(loginURL.toString())
                                               .withTokenRefreshURL(tokenRefreshURL.toString())
                                               .withClientCredentials(clientCredentials)
+                                              .withCredentialEncodingStrategy(strategy)
+                                              .withUserAgent(userAgent)
                                               .build();
         oauthStore.writeProvider(provider, reuseClientCredentials);
         responder.sendStatus(HttpURLConnection.HTTP_OK);
@@ -190,7 +196,7 @@ public void putOAuthCredential(HttpServiceRequest request, HttpServiceResponder
                 + response.getResponseCode()
                 + " , response message: "
                 + response.getResponseMessage()
-                + " , respone body: "
+                + " , response body: "
                 + response.getResponseBodyAsString());
       }
 
@@ -248,7 +254,7 @@ public void getOAuthCredential(HttpServiceRequest request, HttpServiceResponder
                 + response.getResponseCode()
                 + " , response message: "
                 + response.getResponseMessage()
-                + " , respone body: "
+                + " , response body: "
                 + response.getResponseBodyAsString());
       }
 
@@ -307,37 +313,109 @@ private boolean checkCredIsValid(HttpResponse response) throws OAuthServiceExcep
     return !(refreshTokenResponse.getAccessToken() == null || refreshTokenResponse.getAccessToken().isEmpty());
   }
 
+  /**
+   * Create the request body for refresh token & access token requests
+   * @param strategy which encoding strategy is used to send client ID + secret
+   * @param grantType whether an authorization code used to fetch a refresh token or a refresh token used to fetch an
+   *                  access token is used
+   * @param code used when building a request to get a refresh token
+   * @param redirectURI used when building a request to get an access token
+   * @param refreshToken used when building a request to get an access token
+   * @param clientCreds the client ID + secret
+   * @return request body
+   */
+  private String buildRequestBody(CredentialEncodingStrategy strategy,
+                                  String grantType,
+                                  String code,
+                                  String redirectURI,
+                                  String refreshToken,
+                                  OAuthClientCredentials clientCreds) {
+    switch (strategy) {
+      case BASIC_AUTH:
+        return grantType.equals("authorization_code")
+                ? String.format("code=%s&redirect_uri=%s&grant_type=%s", code, redirectURI, grantType)
+                : String.format("grant_type=%s&refresh_token=%s", grantType, refreshToken);
+      case FORM_BODY: // fall-through
+      default:
+        return grantType.equals("authorization_code")
+                ? String.format("code=%s&redirect_uri=%s&client_id=%s&client_secret=%s&grant_type=%s",
+                code, redirectURI, clientCreds.getClientId(), clientCreds.getClientSecret(), grantType)
+                : String.format("grant_type=%s&client_id=%s&client_secret=%s&refresh_token=%s",
+                grantType, clientCreds.getClientId(), clientCreds.getClientSecret(), refreshToken);
+    }
+  }
+
+  /** Build HTTP request for getting tokens */
+  private HttpRequest.Builder buildHttpRequest(String body,
+                                               CredentialEncodingStrategy strategy,
+                                               OAuthClientCredentials clientCreds,
+                                               String refreshTokenURL,
+                                               boolean addContentType,
+                                               String userAgent) throws MalformedURLException {
+    HttpRequest.Builder requestBuilder = HttpRequest.post(new URL(refreshTokenURL))
+            .withBody(body);
+
+    if (addContentType) {
+      requestBuilder.addHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_FORM_URLENCODED);
+    }
+
+    if (strategy == CredentialEncodingStrategy.BASIC_AUTH) {
+      requestBuilder.addHeader(HttpHeaders.AUTHORIZATION, getBasicAuthHeader(clientCreds));
+    }
+
+    if (userAgent != null) {
+      requestBuilder.addHeader(HttpHeaders.USER_AGENT, userAgent);
+    }
+
+    return requestBuilder;
+  }
+
+  /**
+   * Build the HttpRequest to request a refresh token from the OAuth provider
+   * @param provider
+   * @param code the authorization code given after the user accepts OAuth from the provider
+   * @param redirectURI
+   */
   private HttpRequest createGetRefreshTokenRequest(OAuthProvider provider, String code, String redirectURI)
       throws OAuthServiceException {
     OAuthClientCredentials clientCreds = provider.getClientCredentials();
+    CredentialEncodingStrategy strategy = provider.getCredentialEncodingStrategy();
+    String tokenRefreshURL = provider.getTokenRefreshURL();
+    String body = buildRequestBody(strategy, "authorization_code", code, redirectURI, null, clientCreds);
+    String userAgent = provider.getUserAgent();
+
     try {
-      return HttpRequest.post(new URL(provider.getTokenRefreshURL()))
-          .withBody(String.format(
-              "code=%s&redirect_uri=%s&client_id=%s&client_secret=%s&grant_type=authorization_code",
-              code, redirectURI, clientCreds.getClientId(), clientCreds.getClientSecret()))
-          .addHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_FORM_URLENCODED)
-          .build();
+      return buildHttpRequest(body, strategy, clientCreds, tokenRefreshURL, true, userAgent).build();
     } catch (MalformedURLException e) {
       throw new OAuthServiceException(HttpURLConnection.HTTP_INTERNAL_ERROR, "Malformed URL", e);
     }
   }
 
+  /**
+   * Build the HttpRequest to request an access token for making data requests from the OAuth provider
+   * @param provider
+   * @param refreshToken the refresh token requested previously from the provider
+   */
   private HttpRequest createGetAccessTokenRequest(OAuthProvider provider, String refreshToken)
       throws OAuthServiceException {
     OAuthClientCredentials clientCreds = provider.getClientCredentials();
+    CredentialEncodingStrategy strategy = provider.getCredentialEncodingStrategy();
+    String tokenRefreshURL = provider.getTokenRefreshURL();
+    String body = buildRequestBody(strategy, "refresh_token", null, null, refreshToken, clientCreds);
+    String userAgent = provider.getUserAgent();
+
     try {
-      return HttpRequest.post(new URL(provider.getTokenRefreshURL()))
-          .withBody(
-              String.format("grant_type=refresh_token&client_id=%s&client_secret=%s&refresh_token=%s",
-                clientCreds.getClientId(),
-                clientCreds.getClientSecret(),
-                refreshToken))
-          .build();
+      return buildHttpRequest(body, strategy, clientCreds, tokenRefreshURL, false, userAgent).build();
     } catch (MalformedURLException e) {
       throw new OAuthServiceException(HttpURLConnection.HTTP_INTERNAL_ERROR, "Malformed URL", e);
     }
   }
 
+  private String getBasicAuthHeader(OAuthClientCredentials clientCreds) {
+    String authInfo = String.format("%s:%s", clientCreds.getClientId(), clientCreds.getClientSecret());
+    return String.format("Basic %s", Base64.getEncoder().encode(authInfo.getBytes()));
+  }
+
   private OAuthProvider getProvider(String provider) throws OAuthServiceException {
     try {
       Optional<OAuthProvider> providerOptional = oauthStore.getProvider(provider);