CDAP-21265 : Add support for PKCE

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

@@ -0,0 +1,31 @@
+/*
+ * Copyright © 2024 Cask Data, Inc.
+ * 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 io.cdap.cdap.datapipeline.oauth;
+
+/**
+ * The type of OAuth 2.0 authorization flow to use.
+ */
+public enum AuthType {
+  /**
+   * Standard OAuth 2.0 Authorization Code flow.
+   */
+  STANDARD,
+
+  /**
+   * OAuth 2.0 Authorization Code flow with Proof Key for Code Exchange (PKCE).
+   * Provides enhanced security by dynamically generating a code challenge and verifier.
+   */
+  PKCE
+}

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

@@ -33,18 +33,23 @@ public class OAuthProvider {
   @Nullable
   private final String userAgent;
 
+  @Nullable
+  private final AuthType authType;
+
   public OAuthProvider(String name,
                        String loginURL,
                        String tokenRefreshURL,
                        @Nullable OAuthClientCredentials clientCreds,
                        @Nullable CredentialEncodingStrategy strategy,
-                       @Nullable String userAgent) {
+                       @Nullable String userAgent,
+                       @Nullable AuthType authType) {
     this.name = name;
     this.loginURL = loginURL;
     this.tokenRefreshURL = tokenRefreshURL;
     this.clientCreds = clientCreds;
     this.strategy = strategy;
     this.userAgent = userAgent;
+    this.authType = authType != null ? authType : AuthType.STANDARD;
   }
 
   public String getName() {
@@ -74,6 +79,10 @@ public String getUserAgent() {
     return userAgent;
   }
 
+  public AuthType getAuthType() {
+    return authType;
+  }
+
   public enum CredentialEncodingStrategy {
     // (default) Sends client ID & secret as part of the POST request body
     FORM_BODY,
@@ -95,6 +104,7 @@ public static class Builder {
     private OAuthClientCredentials clientCreds;
     private CredentialEncodingStrategy strategy;
     private String userAgent;
+    private AuthType authType;
 
     public Builder() {}
 
@@ -128,6 +138,11 @@ public Builder withUserAgent(@Nullable String userAgent) {
       return this;
     }
 
+    public Builder withAuthType(@Nullable AuthType authType) {
+      this.authType = authType;
+      return this;
+    }
+
     public OAuthProvider build() {
       Preconditions.checkNotNull(name, "OAuth provider name missing");
       Preconditions.checkNotNull(loginURL, "Login URL missing");
@@ -136,7 +151,7 @@ public OAuthProvider build() {
       if (strategy == null) {
         this.strategy = CredentialEncodingStrategy.FORM_BODY;
       }
-      return new OAuthProvider(name, loginURL, tokenRefreshURL, clientCreds, strategy, userAgent);
+      return new OAuthProvider(name, loginURL, tokenRefreshURL, clientCreds, strategy, userAgent, authType);
     }
   }
 }

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

@@ -46,13 +46,16 @@ public class OAuthStore {
   private static final String TOKEN_REFRESH_URL_COL = "tokenrefreshurl";
   private static final String CREDENTIAL_ENCODING_STRATEGY_COL = "credentialencodingstrategy";
   private static final String USER_AGENT_COL = "useragent";
+  private static final String AUTH_TYPE_COL = "authtype";
   private static final String CLIENT_CREDS_KEY_PREFIX = "oauthclientcreds";
   private static final String ACCESS_TOKEN_KEY_PREFIX = "oauthaccesstoken";
   private static final String REFRESH_TOKEN_KEY_PREFIX = "oauthrefreshtoken";
+  private static final String PKCE_KEY_PREFIX = "oauth-pkce";
   private static final Gson GSON = new Gson();
   private final TransactionRunner transactionRunner;
   private final SecureStore secureStore;
   private final SecureStoreManager secureStoreManager;
+  private final long codeVerifierTTL;
 
   public static final StructuredTableId TABLE_ID = new StructuredTableId("oauth");
   public static final StructuredTableSpecification TABLE_SPEC = new StructuredTableSpecification.Builder()
@@ -61,17 +64,20 @@ public class OAuthStore {
                   Fields.stringType(LOGIN_URL_COL),
                   Fields.stringType(TOKEN_REFRESH_URL_COL),
                   Fields.stringType(CREDENTIAL_ENCODING_STRATEGY_COL),
-                  Fields.stringType(USER_AGENT_COL))
+                  Fields.stringType(USER_AGENT_COL),
+                  Fields.stringType(AUTH_TYPE_COL))
       .withPrimaryKeys(OAUTH_PROVIDER_COL)
       .build();
 
   public OAuthStore(
       TransactionRunner transactionRunner,
       SecureStore secureStore,
-      SecureStoreManager secureStoreManager) {
+      SecureStoreManager secureStoreManager,
+      long codeVerifierTTL) {
     this.transactionRunner = transactionRunner;
     this.secureStore = secureStore;
     this.secureStoreManager = secureStoreManager;
+    this.codeVerifierTTL = codeVerifierTTL;
   }
 
   /**
@@ -299,6 +305,10 @@ private static String getAccessTokenKey(String oauthProvider, String credentialI
     return String.format("%s-%s-%s", ACCESS_TOKEN_KEY_PREFIX, oauthProvider.toLowerCase(), credentialId.toLowerCase());
   }
 
+  private static String getPKCEKey(String oauthProvider, String state) {
+    return String.format("%s-%s-%s", PKCE_KEY_PREFIX, oauthProvider.toLowerCase(), state.toLowerCase());
+  }
+
   private static List<Field<?>> getKey(String name) {
     List<Field<?>> keyFields = new ArrayList<>(1);
     keyFields.add(Fields.stringField(OAUTH_PROVIDER_COL, name));
@@ -311,6 +321,7 @@ private static OAuthProvider fromRow(StructuredRow row, OAuthClientCredentials c
     String tokenRefreshURL = row.getString(TOKEN_REFRESH_URL_COL);
     String credentialEncodingStrategy = row.getString(CREDENTIAL_ENCODING_STRATEGY_COL);
     String userAgent = row.getString(USER_AGENT_COL);
+    String authTypeStr = row.getString(AUTH_TYPE_COL);
 
     return OAuthProvider.newBuilder()
         .withName(name)
@@ -321,6 +332,7 @@ private static OAuthProvider fromRow(StructuredRow row, OAuthClientCredentials c
             Optional.ofNullable(credentialEncodingStrategy)
                 .map(OAuthProvider.CredentialEncodingStrategy::valueOf).orElse(null))
         .withUserAgent(userAgent)
+        .withAuthType(authTypeStr != null ? AuthType.valueOf(authTypeStr) : AuthType.STANDARD)
         .build();
   }
 
@@ -333,6 +345,36 @@ private static List<Field<?>> getRow(OAuthProvider oauthProvider) {
             CREDENTIAL_ENCODING_STRATEGY_COL,
             oauthProvider.getCredentialEncodingStrategy().toString()));
     fields.add(Fields.stringField(USER_AGENT_COL, oauthProvider.getUserAgent()));
+    fields.add(Fields.stringField(AUTH_TYPE_COL, oauthProvider.getAuthType().toString()));
     return fields;
   }
+
+  public void writePKCECodeVerifier(String provider, String state, String codeVerifier) throws Exception {
+    String key = getPKCEKey(provider, state);
+    secureStoreManager.put(NamespaceId.SYSTEM.getNamespace(), key, codeVerifier,
+        "PKCE Code Verifier", Collections.emptyMap(), codeVerifierTTL);
+  }
+
+  public String getPKCECodeVerifier(String provider, String state) throws OAuthStoreException {
+    String key = getPKCEKey(provider, state);
+    try {
+      return new String(secureStore.getData(NamespaceId.SYSTEM.getNamespace(), key), StandardCharsets.UTF_8);
+    } catch (IOException e) {
+      throw new OAuthStoreException("Failed to read PKCE code verifier from secure storage", e);
+    } catch (Exception e) {
+      throw new OAuthStoreException("PKCE code verifier not found or expired for state or an "
+          + "unexpected error reading from secure store: ", e);
+    }
+  }
+
+  public void deletePKCECodeVerifier(String provider, String state) throws OAuthStoreException {
+    String key = getPKCEKey(provider, state);
+    try {
+      secureStoreManager.delete(NamespaceId.SYSTEM.getNamespace(), key);
+    } catch (IOException e) {
+      throw new OAuthStoreException("Failed to delete PKCE code verifier from secure storage", e);
+    } catch (Exception e) {
+      // Ignore if not found
+    }
+  }
 }

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

@@ -17,16 +17,21 @@
 
 package io.cdap.cdap.datapipeline.oauth;
 
+import javax.annotation.Nullable;
+
 /**
  * OAuth credential REST PUT request body.
  */
 public class PutOAuthCredentialRequest {
   private final String oneTimeCode;
   private final String redirectURI;
+  @Nullable
+  private final String state;
 
-  public PutOAuthCredentialRequest(String oneTimeCode, String redirectURI) {
+  public PutOAuthCredentialRequest(String oneTimeCode, String redirectURI, @Nullable String state) {
     this.oneTimeCode = oneTimeCode;
     this.redirectURI = redirectURI;
+    this.state = state;
   }
 
   public String getOneTimeCode() {
@@ -36,4 +41,9 @@ public String getOneTimeCode() {
   public String getRedirectURI() {
     return redirectURI;
   }
+
+  @Nullable
+  public String getState() {
+    return state;
+  }
 }

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

@@ -27,20 +27,23 @@ public class PutOAuthProviderRequest {
   private final String clientSecret;
   private final OAuthProvider.CredentialEncodingStrategy strategy;
   private final String userAgent;
+  private final AuthType authType;
 
   public PutOAuthProviderRequest(
           String loginURL,
           String tokenRefreshURL,
           String clientId,
           String clientSecret,
           OAuthProvider.CredentialEncodingStrategy strategy,
-          String userAgent) {
+          String userAgent,
+          AuthType authType) {
     this.loginURL = loginURL;
     this.tokenRefreshURL = tokenRefreshURL;
     this.clientId = clientId;
     this.clientSecret = clientSecret;
     this.strategy = strategy;
     this.userAgent = userAgent;
+    this.authType = authType;
   }
 
   public String getLoginURL() {
@@ -66,4 +69,8 @@ public OAuthProvider.CredentialEncodingStrategy getCredentialEncodingStrategy()
   public String getUserAgent() {
     return userAgent;
   }
+
+  public AuthType getAuthType() {
+    return authType;
+  }
 }