[apache#2179] Improvement: Potential SQL injection point in generateD…

…ropDatabaseSql
zivali · Feb 24, 2024 · a16eb09 · a16eb09
1 parent 5131d17
commit a16eb09
Show file tree

Hide file tree

Showing 6 changed files with 84 additions and 0 deletions.
diff --git a/...src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java b/...src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java
@@ -52,6 +52,7 @@ public void create(String databaseName, String comment, Map<String, String> prop
   public void delete(String databaseName, boolean cascade) throws NoSuchSchemaException {
     LOG.info("Beginning to drop database {}", databaseName);
     try (final Connection connection = getConnection()) {
+      validateDatabaseName(databaseName);
       JdbcConnectorUtils.executeUpdate(connection, generateDropDatabaseSql(databaseName, cascade));
       LOG.info("Finished dropping database {}", databaseName);
     } catch (final SQLException se) {
@@ -93,6 +94,8 @@ protected abstract String generateCreateDatabaseSql(
    */
   protected abstract String generateDropDatabaseSql(String databaseName, boolean cascade);
 
+  protected abstract void validateDatabaseName(String databaseName);
+
   protected Connection getConnection() throws SQLException {
     return dataSource.getConnection();
   }

diff --git a/...c/test/java/com/datastrato/gravitino/catalog/jdbc/operation/SqliteDatabaseOperations.java b/...c/test/java/com/datastrato/gravitino/catalog/jdbc/operation/SqliteDatabaseOperations.java
@@ -91,4 +91,9 @@ public String generateCreateDatabaseSql(
   public String generateDropDatabaseSql(String databaseName, boolean cascade) {
     return null;
   }
+
+  @Override
+  protected void validateDatabaseName(String databaseName) {
+    // do nothing
+  }
 }
diff --git a/...c/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlDatabaseOperations.java b/...c/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlDatabaseOperations.java
@@ -90,6 +90,19 @@ public String generateDropDatabaseSql(String databaseName, boolean cascade) {
     return dropDatabaseSql;
   }
 
+  @Override
+  protected void validateDatabaseName(String databaseName) {
+    if (StringUtils.isEmpty(databaseName)) {
+      throw new IllegalArgumentException("Database name cannot be empty.");
+    }
+    if (databaseName.length() > 64) {
+      throw new IllegalArgumentException("Database name cannot be longer than 64 characters.");
+    }
+    if (!databaseName.matches("^[\\w\\p{L}$]*$")) {
+      throw new IllegalArgumentException("Invalid database name.");
+    }
+  }
+
   @Override
   public JdbcSchema load(String databaseName) throws NoSuchSchemaException {
     try (final Connection connection = this.dataSource.getConnection()) {

diff --git a/...ava/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java b/...ava/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java
@@ -129,6 +129,19 @@ public String generateDropDatabaseSql(String schema, boolean cascade) {
     return sqlBuilder.toString();
   }
 
+  @Override
+  protected void validateDatabaseName(String schema) {
+    if (StringUtils.isEmpty(schema)) {
+      throw new IllegalArgumentException("Schema name cannot be empty.");
+    }
+    if (schema.length() > 63) {
+      throw new IllegalArgumentException("Schema name cannot be longer than 63 characters.");
+    }
+    if (!schema.matches("^[_a-zA-Z\\p{L}][\\w\\p{L}$]*$")) {
+      throw new IllegalArgumentException("Invalid schema name.");
+    }
+  }
+
   @Override
   protected Connection getConnection() throws SQLException {
     Connection connection = dataSource.getConnection();

diff --git a/...datastrato/gravitino/integration/test/catalog/jdbc/mysql/TestMysqlDatabaseOperations.java b/...datastrato/gravitino/integration/test/catalog/jdbc/mysql/TestMysqlDatabaseOperations.java
@@ -10,6 +10,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import org.apache.commons.lang3.StringUtils;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
@@ -29,4 +30,24 @@ public void testBaseOperationDatabase() {
     testBaseOperation(databaseName, properties, comment);
     testDropDatabase(databaseName);
   }
+
+  @Test
+  public void testDropDatabaseWithSqlInjection() {
+    String databaseName = GravitinoITUtils.genRandomName("ct_db");
+    // testDropDatabase should throw an exception with string that might contain SQL injection
+    String sqlInjection = databaseName + " UNION SELECT NOW()";
+    Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          testDropDatabase(sqlInjection);
+        });
+
+    // testDropDatabase should throw an exception with input that has more than 65 characters
+    String invalidInput = StringUtils.repeat("a", 65);
+    Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          testDropDatabase(invalidInput);
+        });
+  }
 }
diff --git a/...to/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlSchemaOperations.java b/...to/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlSchemaOperations.java
@@ -102,4 +102,33 @@ public void testCreateMultipleSchema() throws SQLException {
     JdbcSchema load = postgreSqlSchemaOperations.load(schema_1);
     Assertions.assertEquals(schema_1, load.name());
   }
+
+  @Test
+  public void testDropDatabaseWithSqlInjection() {
+    String databaseName = GravitinoITUtils.genRandomName("ct_db");
+    // testDropDatabase should throw an exception with string that might contain SQL injection
+    String sqlInjection = databaseName + " UNION SELECT NOW()";
+    Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          testDropDatabase(sqlInjection);
+        });
+
+    // testDropDatabase should throw an exception with input that has more than 64 characters
+    String invalidInput = StringUtils.repeat("a", 64);
+    Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          testDropDatabase(invalidInput);
+        });
+
+    // testDropDatabase should throw an exception with schema name that starts with special
+    // character
+    String invalidInput2 = GravitinoITUtils.genRandomName("$test_db");
+    Assertions.assertThrows(
+        IllegalArgumentException.class,
+        () -> {
+          testDropDatabase(invalidInput2);
+        });
+  }
 }