diff --git a/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java b/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java index 6c0be7351c3..7f647383d3a 100644 --- a/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java +++ b/catalogs/catalog-jdbc-common/src/main/java/com/datastrato/gravitino/catalog/jdbc/operation/JdbcDatabaseOperations.java @@ -52,6 +52,7 @@ public void create(String databaseName, String comment, Map 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/catalogs/catalog-jdbc-common/src/test/java/com/datastrato/gravitino/catalog/jdbc/operation/SqliteDatabaseOperations.java b/catalogs/catalog-jdbc-common/src/test/java/com/datastrato/gravitino/catalog/jdbc/operation/SqliteDatabaseOperations.java index 1054e14d094..33b7fe71bee 100644 --- a/catalogs/catalog-jdbc-common/src/test/java/com/datastrato/gravitino/catalog/jdbc/operation/SqliteDatabaseOperations.java +++ b/catalogs/catalog-jdbc-common/src/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/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlDatabaseOperations.java b/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlDatabaseOperations.java index 6aa5f8e627f..abb8bb22fa8 100644 --- a/catalogs/catalog-jdbc-mysql/src/main/java/com/datastrato/gravitino/catalog/mysql/operation/MysqlDatabaseOperations.java +++ b/catalogs/catalog-jdbc-mysql/src/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/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java b/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java index 91de3648da6..b2b910e7f17 100644 --- a/catalogs/catalog-jdbc-postgresql/src/main/java/com/datastrato/gravitino/catalog/postgresql/operation/PostgreSqlSchemaOperations.java +++ b/catalogs/catalog-jdbc-postgresql/src/main/java/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/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/mysql/TestMysqlDatabaseOperations.java b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/mysql/TestMysqlDatabaseOperations.java index d6313b676ad..bcbec68a16a 100644 --- a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/mysql/TestMysqlDatabaseOperations.java +++ b/integration-test/src/test/java/com/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/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlSchemaOperations.java b/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlSchemaOperations.java index 6b2462ad61e..dfc30308d46 100644 --- a/integration-test/src/test/java/com/datastrato/gravitino/integration/test/catalog/jdbc/postgresql/TestPostgreSqlSchemaOperations.java +++ b/integration-test/src/test/java/com/datastrato/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); + }); + } }