Fix SQL injection in connector sample path and validate source queries

amazon-redshift-plugin/src/main/java/io/cdap/plugin/amazon/redshift/RedshiftConnector.java

@@ -82,7 +82,7 @@ protected SchemaReader getSchemaReader(String sessionID) {
 
   @Override
   protected String getTableName(String database, String schema, String table) {
-    return String.format("\"%s\".\"%s\"", schema, table);
+    return quoteIdentifier(schema, '"') + "." + quoteIdentifier(table, '"');
   }
 
   @Override

amazon-redshift-plugin/src/test/java/io/cdap/plugin/amazon/redshift/RedshiftConnectorUnitTest.java

@@ -41,6 +41,15 @@ public void getTableNameTest() {
                         CONNECTOR.getTableName("db", "schema", "table"));
   }
 
+  /**
+   * Unit test for getTableName() with identifiers containing double quotes
+   */
+  @Test
+  public void getTableNameEscapesQuotesTest() {
+    Assert.assertEquals("\"my\"\"schema\".\"t\"\"; DELETE FROM x; --\"",
+                        CONNECTOR.getTableName("db", "my\"schema", "t\"; DELETE FROM x; --"));
+  }
+
   @Test
   public void getRandomQuery() {
     Assert.assertEquals("SELECT * FROM TestData\n" +

cloudsql-mysql-plugin/src/main/java/io/cdap/plugin/cloudsql/mysql/CloudSQLMySQLConnector.java

@@ -82,7 +82,7 @@ protected SchemaReader getSchemaReader(String sessionID) {
 
   @Override
   protected String getTableName(String database, String schema, String table) {
-    return String.format("`%s`.`%s`", database, table);
+    return quoteIdentifier(database, '`') + "." + quoteIdentifier(table, '`');
   }
 
   @Override

cloudsql-postgresql-plugin/src/main/java/io/cdap/plugin/cloudsql/postgres/CloudSQLPostgreSQLConnector.java

@@ -86,7 +86,7 @@ protected SchemaReader getSchemaReader(String sessionID) {
 
   @Override
   protected String getTableName(String database, String schema, String table) {
-    return String.format("\"%s\".\"%s\"", schema, table);
+    return quoteIdentifier(schema, '"') + "." + quoteIdentifier(table, '"');
   }
 
   @Override

database-commons/pom.xml

@@ -77,6 +77,10 @@
       <groupId>dev.failsafe</groupId>
       <artifactId>failsafe</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.github.jsqlparser</groupId>
+      <artifactId>jsqlparser</artifactId>
+    </dependency>
 
     <!-- test dependencies -->
     <dependency>

database-commons/src/main/java/io/cdap/plugin/db/connector/AbstractDBSpecificConnector.java

@@ -150,8 +150,17 @@ protected String getConnectionString(String database) {
   }
 
   protected String getTableName(String database, String schema, String table) {
-    return schema == null ? String.format("\"%s\".\"%s\"", database, table)
-      : String.format("\"%s\".\"%s\".\"%s\"", database, schema, table);
+    return schema == null ? quoteIdentifier(database, '"') + "." + quoteIdentifier(table, '"')
+      : quoteIdentifier(database, '"') + "." + quoteIdentifier(schema, '"') + "." + quoteIdentifier(table, '"');
+  }
+
+  /**
+   * Wraps the given identifier in the given quote character, doubling any embedded quote characters so that a
+   * user supplied database, schema, table or column name can never break out of the quoted identifier.
+   */
+  public static String quoteIdentifier(String identifier, char quote) {
+    String quoteString = String.valueOf(quote);
+    return quoteString + identifier.replace(quoteString, quoteString + quoteString) + quoteString;
   }
 
   protected String getTableQuery(String database, String schema, String table) {

database-commons/src/main/java/io/cdap/plugin/db/source/AbstractDBSource.java

@@ -52,6 +52,7 @@
 import io.cdap.plugin.util.DriverCleanup;
 import io.cdap.plugin.util.RetryPolicyUtil;
 import io.cdap.plugin.util.RetryUtils;
+import io.cdap.plugin.util.SelectQueryValidator;
 import org.apache.hadoop.io.LongWritable;
 import org.apache.hadoop.mapreduce.MRJobConfig;
 import org.apache.hadoop.mapreduce.lib.db.DBConfiguration;
@@ -106,6 +107,7 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
     StageConfigurer stageConfigurer = pipelineConfigurer.getStageConfigurer();
     FailureCollector collector = stageConfigurer.getFailureCollector();
     sourceConfig.validate(collector);
+    boolean queriesValid = validateQueries(collector);
     if (sourceConfig.getSchema() != null) {
       stageConfigurer.setOutputSchema(sourceConfig.getSchema());
       return;
@@ -119,7 +121,8 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
     Class<? extends Driver> driverClass = DBUtils.getDriverClass(
       pipelineConfigurer, sourceConfig, ConnectionConfig.JDBC_PLUGIN_TYPE);
 
-    if (sourceConfig.canConnect()) {
+    // never run a query that failed the read-only validation, not even to fetch its schema
+    if (sourceConfig.canConnect() && queriesValid) {
       try {
         stageConfigurer.setOutputSchema(getSchema(driverClass));
       } catch (IllegalAccessException | InstantiationException e) {
@@ -228,6 +231,43 @@ protected String getErrorDetailsProviderClassName() {
     return null;
   }
 
+  /**
+   * Returns whether square brackets quote identifiers in the database's SQL dialect, as in SQL Server.
+   * Override this method to allow such identifiers in the import and bounding queries.
+   */
+  protected boolean useSquareBracketQuotes() {
+    return false;
+  }
+
+  /**
+   * Validates that the import and bounding queries are single read-only SELECT statements, since they are executed
+   * with the credential of the connection. Queries that are macros are validated at runtime, once they are evaluated.
+   *
+   * @return true if the queries are valid or cannot be validated yet
+   */
+  private boolean validateQueries(FailureCollector collector) {
+    boolean importQueryValid = validateQuery(DBSourceConfig.IMPORT_QUERY, "Import Query",
+                                             sourceConfig.getImportQuery(), collector);
+    boolean boundingQueryValid = validateQuery(DBSourceConfig.BOUNDING_QUERY, "Bounding Query",
+                                               sourceConfig.getBoundingQuery(), collector);
+    return importQueryValid && boundingQueryValid;
+  }
+
+  private boolean validateQuery(String property, String label, @Nullable String query, FailureCollector collector) {
+    if (sourceConfig.containsMacro(property) || Strings.isNullOrEmpty(query)) {
+      return true;
+    }
+    try {
+      SelectQueryValidator.validate(query, useSquareBracketQuotes());
+      return true;
+    } catch (IllegalArgumentException e) {
+      collector.addFailure(String.format("Invalid %s: %s", label, e.getMessage()),
+                           "Provide a single SELECT statement that only reads data.")
+        .withConfigProperty(property);
+      return false;
+    }
+  }
+
   /**
    * Returns the DBErrorDetailsProvider instance.
    * Override this method to provide a custom DBErrorDetailsProvider instance.
@@ -272,6 +312,7 @@ private Connection getConnection() {
   public void prepareRun(BatchSourceContext context) throws Exception {
     FailureCollector collector = context.getFailureCollector();
     sourceConfig.validate(collector);
+    validateQueries(collector);
     collector.getOrThrowException();
 
     LOG.debug("pluginType = {}; pluginName = {}; connectionString = {}; importQuery = {}; " +

databricks-plugin/src/main/java/io/cdap/plugin/databricks/DatabricksConnector.java

@@ -116,15 +116,15 @@ protected SchemaReader getSchemaReader(String sessionID) {
   @Override
   protected String getTableName(String database, String schema, String table) {
     if (database == null && schema == null) {
-      return String.format("`%s`", table);
+      return quoteIdentifier(table, '`');
     }
     if (database == null) {
-      return String.format("`%s`.`%s`", schema, table);
+      return quoteIdentifier(schema, '`') + "." + quoteIdentifier(table, '`');
     }
     if (schema == null) {
-      return String.format("`%s`.`%s`", database, table);
+      return quoteIdentifier(database, '`') + "." + quoteIdentifier(table, '`');
     }
-    return String.format("`%s`.`%s`.`%s`", database, schema, table);
+    return quoteIdentifier(database, '`') + "." + quoteIdentifier(schema, '`') + "." + quoteIdentifier(table, '`');
   }
 
   @Override
@@ -136,6 +136,7 @@ protected String getRandomQuery(String tableName, int limit) {
 
   @Override
   protected String getStratifiedQuery(String tableName, int limit, String strata, String sessionID) {
+    String strataColumn = quoteIdentifier(strata, '`');
     return String.format("WITH t_%s AS (\n" +
         "    SELECT *,\n" +
         "    ROW_NUMBER() OVER (ORDER BY %s, RAND()) AS sqn_%s,\n" +
@@ -146,8 +147,8 @@ protected String getStratifiedQuery(String tableName, int limit, String strata,
         "WHERE MOD(sqn_%s, GREATEST(1, CAST(c_%s / %d AS BIGINT))) = 1\n" +
         "ORDER BY %s\n" +
         "LIMIT %d",
-      sessionID, strata, sessionID, sessionID, tableName, sessionID, sessionID, sessionID,
-      limit, strata, limit);
+      sessionID, strataColumn, sessionID, sessionID, tableName, sessionID, sessionID, sessionID,
+      limit, strataColumn, limit);
   }
 
   @Override

databricks-plugin/src/test/java/io/cdap/plugin/databricks/DatabricksConnectorUnitTest.java

@@ -38,6 +38,21 @@ public void testGetTableName() {
                         CONNECTOR.getTableName(null, null, "my_table"));
   }
 
+  @Test
+  public void testGetTableNameEscapesBackticks() {
+    Assert.assertEquals("`main`.`default`.`t``; DELETE FROM x; --`",
+                        CONNECTOR.getTableName("main", "default", "t`; DELETE FROM x; --"));
+    Assert.assertEquals("`my``schema`.`my``table`",
+                        CONNECTOR.getTableName(null, "my`schema", "my`table"));
+  }
+
+  @Test
+  public void testGetStratifiedQueryEscapesStrata() {
+    String query = CONNECTOR.getStratifiedQuery("`main`.`default`.`my_table`", 10, "id`; DROP TABLE t; --", "s1");
+    Assert.assertTrue(query.contains("ORDER BY `id``; DROP TABLE t; --`, RAND()"));
+    Assert.assertTrue(query.contains("ORDER BY `id``; DROP TABLE t; --`\n"));
+  }
+
   @Test
   public void testGetRandomQuery() {
     Assert.assertEquals("SELECT * FROM `main`.`default`.`my_table`\n" +
@@ -49,13 +64,13 @@ public void testGetRandomQuery() {
   public void testGetStratifiedQuery() {
     Assert.assertEquals("WITH t_s1 AS (\n" +
                           "    SELECT *,\n" +
-                          "    ROW_NUMBER() OVER (ORDER BY id, RAND()) AS sqn_s1,\n" +
+                          "    ROW_NUMBER() OVER (ORDER BY `id`, RAND()) AS sqn_s1,\n" +
                           "    COUNT(*) OVER () AS c_s1\n" +
                           "    FROM `main`.`default`.`my_table`\n" +
                           "  )\n" +
                           "SELECT * FROM t_s1\n" +
                           "WHERE MOD(sqn_s1, GREATEST(1, CAST(c_s1 / 10 AS BIGINT))) = 1\n" +
-                          "ORDER BY id\n" +
+                          "ORDER BY `id`\n" +
                           "LIMIT 10",
                         CONNECTOR.getStratifiedQuery("`main`.`default`.`my_table`", 10, "id", "s1"));
   }

generic-database-plugin/src/main/java/io/cdap/plugin/jdbc/DatabaseSource.java

@@ -46,6 +46,14 @@ protected String createConnectionString() {
     return databaseSourceConfig.connectionString;
   }
 
+  @Override
+  protected boolean useSquareBracketQuotes() {
+    // SQL Server quotes identifiers with square brackets
+    String connectionString = databaseSourceConfig.connectionString;
+    return connectionString != null && (connectionString.toLowerCase().startsWith("jdbc:sqlserver:")
+      || connectionString.toLowerCase().startsWith("jdbc:jtds:sqlserver:"));
+  }
+
   /**
    * Generic database source configuration.
    */

generic-database-plugin/src/test/java/io/cdap/plugin/db/sink/DBSourceTestRun.java

@@ -20,16 +20,19 @@
 import io.cdap.cdap.api.common.Bytes;
 import io.cdap.cdap.api.data.format.StructuredRecord;
 import io.cdap.cdap.api.dataset.table.Table;
+import io.cdap.cdap.datapipeline.SmartWorkflow;
 import io.cdap.cdap.etl.api.batch.BatchSource;
 import io.cdap.cdap.etl.mock.batch.MockSink;
 import io.cdap.cdap.etl.proto.v2.ETLBatchConfig;
 import io.cdap.cdap.etl.proto.v2.ETLPlugin;
 import io.cdap.cdap.etl.proto.v2.ETLStage;
+import io.cdap.cdap.proto.ProgramRunStatus;
 import io.cdap.cdap.proto.artifact.AppRequest;
 import io.cdap.cdap.proto.id.ApplicationId;
 import io.cdap.cdap.proto.id.NamespaceId;
 import io.cdap.cdap.test.ApplicationManager;
 import io.cdap.cdap.test.DataSetManager;
+import io.cdap.cdap.test.WorkflowManager;
 import io.cdap.plugin.GenericDatabasePluginTestBase;
 import io.cdap.plugin.common.Constants;
 import io.cdap.plugin.db.ConnectionConfig;
@@ -41,7 +44,11 @@
 import java.math.BigDecimal;
 import java.math.MathContext;
 import java.nio.ByteBuffer;
+import java.sql.Connection;
 import java.sql.Date;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.sql.Statement;
 import java.sql.Time;
 import java.text.SimpleDateFormat;
 import java.time.LocalDate;
@@ -52,6 +59,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 /**
  * Test for ETL using databases.
@@ -329,6 +337,61 @@ public void testUserNamePasswordCombinations() throws Exception {
     deployApplication(appId, appRequest);
   }
 
+  @Test
+  public void testNonSelectImportQueryFailsDeployment() throws Exception {
+    int rowCount = countRows("my_table");
+    String importQuery = "SELECT * FROM \"my_table\" WHERE $CONDITIONS; DELETE FROM \"my_table\"";
+    ETLPlugin sourceConfig = new ETLPlugin(
+      DatabaseConstants.PLUGIN_NAME,
+      BatchSource.PLUGIN_TYPE,
+      ImmutableMap.<String, String>builder()
+        .put(ConnectionConfig.CONNECTION_STRING, getConnectionURL())
+        .put(ConnectionConfig.JDBC_PLUGIN_NAME, JDBC_DRIVER_NAME)
+        .put(AbstractDBSource.DBSourceConfig.IMPORT_QUERY, importQuery)
+        .put(AbstractDBSource.DBSourceConfig.BOUNDING_QUERY, "SELECT MIN(ID),MAX(ID) from \"my_table\"")
+        .put(AbstractDBSource.DBSourceConfig.SPLIT_BY, "ID")
+        .put(Constants.Reference.REFERENCE_NAME, "DBNonSelectQueryTest")
+        .build(),
+      null);
+    ETLBatchConfig etlConfig = getETLBatchConfig(sourceConfig, MockSink.getPlugin("nonSelectQueryOutput"));
+    assertDeploymentFailure(NamespaceId.DEFAULT.app("dbSourceNonSelectQueryTest"), etlConfig, DATAPIPELINE_ARTIFACT,
+                            "Deploying DB Source with a non SELECT import query should have failed.");
+    Assert.assertEquals(rowCount, countRows("my_table"));
+  }
+
+  @Test
+  public void testNonSelectImportQueryMacroFailsRun() throws Exception {
+    int rowCount = countRows("my_table");
+    ETLPlugin sourceConfig = new ETLPlugin(
+      DatabaseConstants.PLUGIN_NAME,
+      BatchSource.PLUGIN_TYPE,
+      ImmutableMap.<String, String>builder()
+        .put(ConnectionConfig.CONNECTION_STRING, getConnectionURL())
+        .put(ConnectionConfig.JDBC_PLUGIN_NAME, JDBC_DRIVER_NAME)
+        .put(AbstractDBSource.DBSourceConfig.IMPORT_QUERY, "${importQuery}")
+        .put(AbstractDBSource.DBSourceConfig.BOUNDING_QUERY, "SELECT MIN(ID),MAX(ID) from \"my_table\"")
+        .put(AbstractDBSource.DBSourceConfig.SPLIT_BY, "ID")
+        .put(Constants.Reference.REFERENCE_NAME, "DBNonSelectQueryMacroTest")
+        .build(),
+      null);
+    ApplicationManager appManager = deployETL(sourceConfig, MockSink.getPlugin("nonSelectQueryMacroOutput"),
+                                              DATAPIPELINE_ARTIFACT, "dbSourceNonSelectQueryMacroTest");
+    WorkflowManager workflowManager = appManager.getWorkflowManager(SmartWorkflow.NAME);
+    workflowManager.start(ImmutableMap.of(
+      "importQuery", "SELECT * FROM \"my_table\" WHERE $CONDITIONS; DELETE FROM \"my_table\""));
+    workflowManager.waitForRun(ProgramRunStatus.FAILED, 5, TimeUnit.MINUTES);
+    Assert.assertEquals(rowCount, countRows("my_table"));
+  }
+
+  private int countRows(String table) throws SQLException {
+    try (Connection connection = getConnection();
+         Statement statement = connection.createStatement();
+         ResultSet resultSet = statement.executeQuery(String.format("SELECT COUNT(*) FROM \"%s\"", table))) {
+      resultSet.next();
+      return resultSet.getInt(1);
+    }
+  }
+
   @Test
   public void testNonExistentDBTable() throws Exception {
     // source

mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerConnector.java

@@ -117,7 +117,8 @@ public StructuredRecord transform(LongWritable longWritable, SqlServerSourceDBRe
 
   @Override
   protected String getTableName(String database, String schema, String table) {
-    return String.format("\"%s\".\"%s\".\"%s\"", database, schema, table);
+    return quoteIdentifier(database, '"') + "." + quoteIdentifier(schema, '"') + "."
+      + quoteIdentifier(table, '"');
   }
 
   @Override
@@ -137,6 +138,7 @@ protected String getRandomQuery(String tableName, int limit) {
 
   @Override
   protected String getStratifiedQuery(String tableName, int limit, String strata, String sessionID) {
+    String strataColumn = quoteIdentifier(strata, '"');
     return String.format("WITH t_%s AS (\n" +
                            "    SELECT *,\n" +
                            "    ROW_NUMBER() OVER (ORDER BY %s, RAND()) AS sqn_%s,\n" +
@@ -146,8 +148,8 @@ protected String getStratifiedQuery(String tableName, int limit, String strata,
                            "SELECT TOP %d * FROM t_%s\n" +
                            "WHERE sqn_%s %% CAST(0.5 * ((c_%s / %d + 1) + ABS(c_%s / %d - 1)) AS bigint) = 1\n" +
                            "ORDER BY %s",
-                         sessionID, strata, sessionID, sessionID, tableName, limit, sessionID, sessionID, sessionID,
-                         limit, sessionID, limit, strata);
+                         sessionID, strataColumn, sessionID, sessionID, tableName, limit, sessionID, sessionID,
+                         sessionID, limit, sessionID, limit, strataColumn);
   }
 
   @Override

mssql-plugin/src/main/java/io/cdap/plugin/mssql/SqlServerSource.java

@@ -81,6 +81,11 @@ protected String getErrorDetailsProviderClassName() {
     return SqlServerErrorDetailsProvider.class.getName();
   }
 
+  @Override
+  protected boolean useSquareBracketQuotes() {
+    return true;
+  }
+
   @Override
   protected DBErrorDetailsProvider getErrorDetailsProvider() {
     if (dbErrorDetailsProvider == null) {