feat: snowflake key auth server

app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/SnowflakePlugin.java

@@ -7,6 +7,7 @@
 import com.appsmith.external.models.ActionConfiguration;
 import com.appsmith.external.models.ActionExecutionRequest;
 import com.appsmith.external.models.ActionExecutionResult;
+import com.appsmith.external.models.AuthenticationDTO;
 import com.appsmith.external.models.DBAuth;
 import com.appsmith.external.models.DatasourceConfiguration;
 import com.appsmith.external.models.DatasourceStructure;
@@ -15,32 +16,29 @@
 import com.appsmith.external.plugins.BasePlugin;
 import com.appsmith.external.plugins.PluginExecutor;
 import com.external.plugins.exceptions.SnowflakeErrorMessages;
+import com.external.utils.SnowflakeKeyUtils;
 import com.external.utils.SqlUtils;
 import com.zaxxer.hikari.HikariConfig;
 import com.zaxxer.hikari.HikariDataSource;
 import com.zaxxer.hikari.HikariPoolMXBean;
 import com.zaxxer.hikari.pool.HikariPool;
 import lombok.extern.slf4j.Slf4j;
+import net.snowflake.client.jdbc.SnowflakeBasicDataSource;
+import org.bouncycastle.pkcs.PKCSException;
 import org.pf4j.Extension;
 import org.pf4j.PluginWrapper;
 import org.springframework.util.StringUtils;
 import reactor.core.publisher.Mono;
 import reactor.core.scheduler.Scheduler;
 import reactor.core.scheduler.Schedulers;
 
+import java.security.*;
+import java.sql.*;
 import java.sql.Connection;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.sql.Statement;
-import java.util.ArrayList;
-import java.util.Comparator;
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.LinkedHashMap;
-import java.util.Map;
-import java.util.Properties;
-import java.util.Set;
+import java.util.*;
 
+import static com.appsmith.external.constants.Authentication.BASIC;
+import static com.appsmith.external.constants.Authentication.SNOWFLAKE_KEY_PAIR_AUTH;
 import static com.appsmith.external.constants.PluginConstants.PluginName.SNOWFLAKE_PLUGIN_NAME;
 import static com.external.utils.ExecutionUtils.getRowsFromQueryResult;
 import static com.external.utils.SnowflakeDatasourceUtils.getConnectionFromHikariConnectionPool;
@@ -159,58 +157,31 @@ public Mono<ActionExecutionResult> execute(
         @Override
         public Mono<HikariDataSource> createConnectionClient(
                 DatasourceConfiguration datasourceConfiguration, Properties properties) {
-            return Mono.fromCallable(() -> {
-                        HikariConfig config = new HikariConfig();
-
-                        config.setDriverClassName(properties.getProperty("driver_name"));
-
-                        config.setMinimumIdle(
-                                Integer.parseInt(properties.get("minimumIdle").toString()));
-                        config.setMaximumPoolSize(Integer.parseInt(
-                                properties.get("maximunPoolSize").toString()));
-
-                        config.setInitializationFailTimeout(Long.parseLong(
-                                properties.get("initializationFailTimeout").toString()));
-                        config.setConnectionTimeout(Long.parseLong(
-                                properties.get("connectionTimeoutMillis").toString()));
-
-                        if (Authentication.SNOWFLAKE_KEY_PAIR_AUTH.equals(
-                                datasourceConfiguration.getAuthentication().getAuthenticationType())) {
-                            KeyPairAuth authentication = (KeyPairAuth) datasourceConfiguration.getAuthentication();
-                            // Set authentication properties
-                            if (authentication.getUsername() != null) {
-                                config.setUsername(authentication.getUsername());
-                            }
-                        } else {
-                            // Set authentication properties
-                            DBAuth authentication = (DBAuth) datasourceConfiguration.getAuthentication();
-                            if (authentication.getUsername() != null) {
-                                config.setUsername(authentication.getUsername());
-                            }
-                            if (authentication.getPassword() != null) {
-                                config.setPassword(authentication.getPassword());
-                            }
-                        }
-
-                        // Set up the connection URL
-                        StringBuilder urlBuilder = new StringBuilder(
-                                "jdbc:snowflake://" + datasourceConfiguration.getUrl() + ".snowflakecomputing.com?");
-                        config.setJdbcUrl(urlBuilder.toString());
+            return getHikariConfig(datasourceConfiguration, properties)
+                    .flatMap(config -> Mono.fromCallable(() -> {
+                                // Set up the connection URL
+                                StringBuilder urlBuilder = new StringBuilder("jdbc:snowflake://"
+                                        + datasourceConfiguration.getUrl() + ".snowflakecomputing.com?");
+                                config.setJdbcUrl(urlBuilder.toString());
 
-                        config.setDataSourceProperties(properties);
+                                config.setDataSourceProperties(properties);
 
-                        // Now create the connection pool from the configuration
-                        HikariDataSource datasource = null;
-                        try {
-                            datasource = new HikariDataSource(config);
-                        } catch (HikariPool.PoolInitializationException e) {
-                            throw new AppsmithPluginException(
-                                    AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, e.getMessage());
-                        }
+                                // Now create the connection pool from the configuration
+                                HikariDataSource datasource = null;
+                                try {
+                                    datasource = new HikariDataSource(config);
+                                } catch (HikariPool.PoolInitializationException e) {
+                                    throw new AppsmithPluginException(
+                                            AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, e.getMessage());
+                                }
 
-                        return datasource;
-                    })
-                    .subscribeOn(scheduler);
+                                return datasource;
+                            })
+                            .subscribeOn(scheduler))
+                    .onErrorMap(
+                            AppsmithPluginException.class,
+                            error -> new AppsmithPluginException(
+                                    AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, error.getMessage()));
         }
 
         @Override
@@ -234,17 +205,14 @@ public Properties addPluginSpecificProperties(
         @Override
         public Properties addAuthParamsToConnectionConfig(
                 DatasourceConfiguration datasourceConfiguration, Properties properties) {
-
-            if (Authentication.SNOWFLAKE_KEY_PAIR_AUTH.equals(
-                    datasourceConfiguration.getAuthentication().getAuthenticationType())) {
-                KeyPairAuth authentication = (KeyPairAuth) datasourceConfiguration.getAuthentication();
-                properties.setProperty("user", authentication.getUsername());
-            } else {
-                DBAuth authentication = (DBAuth) datasourceConfiguration.getAuthentication();
-                properties.setProperty("user", authentication.getUsername());
-                properties.setProperty("password", authentication.getPassword());
+            // Only for username password auth, we need to set these properties, for others
+            // like key-pair auth, authentication specific properties need to be set on config itself
+            AuthenticationDTO authentication = datasourceConfiguration.getAuthentication();
+            if (authentication instanceof DBAuth) {
+                DBAuth dbAuth = (DBAuth) authentication;
+                properties.setProperty("user", dbAuth.getUsername());
+                properties.setProperty("password", dbAuth.getPassword());
             }
-
             properties.setProperty(
                     "warehouse",
                     String.valueOf(
@@ -491,5 +459,108 @@ public Mono<DatasourceStructure> getStructure(
                     })
                     .subscribeOn(scheduler);
         }
+
+        private Mono<HikariConfig> getHikariConfig(
+                DatasourceConfiguration datasourceConfiguration, Properties properties) {
+            HikariConfig commonConfig = getCommonHikariConfig(properties);
+            Mono<HikariConfig> configMono = Mono.just(commonConfig);
+
+            String authenticationType =
+                    datasourceConfiguration.getAuthentication().getAuthenticationType();
+            if (authenticationType != null) {
+                switch (authenticationType) {
+                    case BASIC:
+                        configMono = configMono.flatMap(config -> getBasicAuthConfig(config, datasourceConfiguration));
+                        break;
+                    case SNOWFLAKE_KEY_PAIR_AUTH:
+                        configMono =
+                                configMono.flatMap(config -> getKeyPairAuthConfig(config, datasourceConfiguration));
+                        break;
+                    default:
+                        break;
+                }
+            }
+            return configMono;
+        }
+
+        private Mono<HikariConfig> getKeyPairAuthConfig(
+                HikariConfig config, DatasourceConfiguration datasourceConfiguration) {
+            KeyPairAuth keyPairAuthConfig = (KeyPairAuth) datasourceConfiguration.getAuthentication();
+            byte[] keyBytes = keyPairAuthConfig.getPrivateKey().getDecodedContent();
+            String passphrase = keyPairAuthConfig.getPassphrase();
+            return getPrivateKeyFromBase64(keyBytes, passphrase)
+                    .flatMap(privateKey -> {
+                        StringBuilder urlBuilder = new StringBuilder(
+                                "jdbc:snowflake://" + datasourceConfiguration.getUrl() + ".snowflakecomputing.com?");
+                        config.setJdbcUrl(urlBuilder.toString());
+
+                        SnowflakeBasicDataSource ds = new SnowflakeBasicDataSource();
+                        ds.setPrivateKey(privateKey);
+                        ds.setUser(keyPairAuthConfig.getUsername());
+                        ds.setUrl(urlBuilder.toString());
+                        ds.setWarehouse(String.valueOf(
+                                datasourceConfiguration.getProperties().get(0).getValue()));
+                        ds.setDatabaseName(String.valueOf(
+                                datasourceConfiguration.getProperties().get(1).getValue()));
+                        ds.setRole(String.valueOf(
+                                datasourceConfiguration.getProperties().get(3).getValue()));
+                        ds.setSchema(String.valueOf(
+                                datasourceConfiguration.getProperties().get(2).getValue()));
+                        config.setDataSource(ds);
+
+                        return Mono.just(config);
+                    })
+                    .onErrorMap(
+                            AppsmithPluginException.class,
+                            error -> new AppsmithPluginException(
+                                    AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR, error.getMessage()));
+        }
+
+        private Mono<HikariConfig> getBasicAuthConfig(
+                HikariConfig config, DatasourceConfiguration datasourceConfiguration) {
+            return Mono.fromCallable(() -> {
+                DBAuth authentication = (DBAuth) datasourceConfiguration.getAuthentication();
+                if (authentication.getUsername() != null) {
+                    config.setUsername(authentication.getUsername());
+                }
+                if (authentication.getPassword() != null) {
+                    config.setPassword(authentication.getPassword());
+                }
+                return config;
+            });
+        }
+
+        private HikariConfig getCommonHikariConfig(Properties properties) {
+            HikariConfig config = new HikariConfig();
+            config.setDriverClassName(properties.getProperty("driver_name"));
+
+            config.setMinimumIdle(Integer.parseInt(properties.get("minimumIdle").toString()));
+            config.setMaximumPoolSize(
+                    Integer.parseInt(properties.get("maximunPoolSize").toString()));
+
+            config.setInitializationFailTimeout(
+                    Long.parseLong(properties.get("initializationFailTimeout").toString()));
+            config.setConnectionTimeout(
+                    Long.parseLong(properties.get("connectionTimeoutMillis").toString()));
+            return config;
+        }
+
+        private Mono<PrivateKey> getPrivateKeyFromBase64(byte[] keyBytes, String passphrase) {
+            try {
+                return Mono.just(SnowflakeKeyUtils.readEncryptedPrivateKey(keyBytes, passphrase));
+            } catch (IllegalArgumentException e) {
+                return Mono.error(new AppsmithPluginException(
+                        AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
+                        SnowflakeErrorMessages.PASSPHRASE_IS_REQUIRED_FOR_ENCRYPTED_PRIVATE_KEY));
+            } catch (PKCSException e) {
+                return Mono.error(new AppsmithPluginException(
+                        AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
+                        SnowflakeErrorMessages.PASSPHRASE_OR_PRIVATE_KEY_INCORRECT));
+            } catch (Exception e) {
+                return Mono.error(new AppsmithPluginException(
+                        AppsmithPluginError.PLUGIN_DATASOURCE_ARGUMENT_ERROR,
+                        SnowflakeErrorMessages.UNABLE_TO_CREATE_CONNECTION_ERROR_MSG));
+            }
+        }
     }
 }

app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/plugins/exceptions/SnowflakeErrorMessages.java

@@ -36,4 +36,8 @@ public class SnowflakeErrorMessages extends BasePluginErrorMessages {
     public static final String DS_MISSING_PASSWORD_ERROR_MSG = "Missing password for authentication.";
 
     public static final String DS_MISSING_PRIVATE_KEY_ERROR_MSG = "Missing private key for authentication.";
+
+    public static final String PASSPHRASE_IS_REQUIRED_FOR_ENCRYPTED_PRIVATE_KEY =
+            "Passphrase is required as private key is encrypted";
+    public static final String PASSPHRASE_OR_PRIVATE_KEY_INCORRECT = "Passphrase or private key is incorrect";
 }

app/server/appsmith-plugins/snowflakePlugin/src/main/java/com/external/utils/SnowflakeKeyUtils.java

@@ -0,0 +1,46 @@
+package com.external.utils;
+
+import com.external.plugins.exceptions.SnowflakeErrorMessages;
+import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
+import org.bouncycastle.jce.provider.BouncyCastleProvider;
+import org.bouncycastle.openssl.PEMParser;
+import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
+import org.bouncycastle.openssl.jcajce.JceOpenSSLPKCS8DecryptorProviderBuilder;
+import org.bouncycastle.operator.InputDecryptorProvider;
+import org.bouncycastle.pkcs.PKCS8EncryptedPrivateKeyInfo;
+
+import java.io.StringReader;
+import java.security.PrivateKey;
+import java.security.Security;
+
+import static org.apache.commons.lang.StringUtils.isEmpty;
+
+public class SnowflakeKeyUtils {
+    static {
+        Security.addProvider(new org.bouncycastle.jce.provider.BouncyCastleProvider());
+    }
+
+    public static PrivateKey readEncryptedPrivateKey(byte[] keyBytes, String passphrase) throws Exception {
+        PrivateKeyInfo privateKeyInfo = null;
+        String privateKeyPEM = new String(keyBytes);
+        PEMParser pemParser = new PEMParser(new StringReader(privateKeyPEM));
+        Object pemObject = pemParser.readObject();
+        if (pemObject instanceof PKCS8EncryptedPrivateKeyInfo) {
+            // Handle the case where the private key is encrypted.
+            if (isEmpty(passphrase)) {
+                throw new IllegalArgumentException(
+                        SnowflakeErrorMessages.PASSPHRASE_IS_REQUIRED_FOR_ENCRYPTED_PRIVATE_KEY);
+            }
+            PKCS8EncryptedPrivateKeyInfo encryptedPrivateKeyInfo = (PKCS8EncryptedPrivateKeyInfo) pemObject;
+            InputDecryptorProvider pkcs8Prov =
+                    new JceOpenSSLPKCS8DecryptorProviderBuilder().build(passphrase.toCharArray());
+            privateKeyInfo = encryptedPrivateKeyInfo.decryptPrivateKeyInfo(pkcs8Prov);
+        } else if (pemObject instanceof PrivateKeyInfo) {
+            // Handle the case where the private key is unencrypted.
+            privateKeyInfo = (PrivateKeyInfo) pemObject;
+        }
+        pemParser.close();
+        JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider(BouncyCastleProvider.PROVIDER_NAME);
+        return converter.getPrivateKey(privateKeyInfo);
+    }
+}