diff --git a/ruoyi-common/src/main/java/com/ruoyi/common/utils/sql/SqlUtil.java b/ruoyi-common/src/main/java/com/ruoyi/common/utils/sql/SqlUtil.java index 1c4c1713e..e19577a8b 100644 --- a/ruoyi-common/src/main/java/com/ruoyi/common/utils/sql/SqlUtil.java +++ b/ruoyi-common/src/main/java/com/ruoyi/common/utils/sql/SqlUtil.java @@ -10,6 +10,11 @@ import com.ruoyi.common.utils.StringUtils; */ public class SqlUtil { + /** + * 定义常用的 sql关键字 + */ + public static String SQL_REGEX = "select |insert |delete |update |drop |count |exec |chr |mid |master |truncate |char |and |declare "; + /** * 仅支持字母、数字、下划线、空格、逗号、小数点(支持多个字段排序) */ @@ -34,4 +39,23 @@ public class SqlUtil { return value.matches(SQL_PATTERN); } + + /** + * SQL关键字检查 + */ + public static void filterKeyword(String value) + { + if (StringUtils.isEmpty(value)) + { + return; + } + String[] sqlKeywords = StringUtils.split(SQL_REGEX, "\\|"); + for (int i = 0; i < sqlKeywords.length; i++) + { + if (StringUtils.indexOfIgnoreCase(value, sqlKeywords[i]) > -1) + { + throw new BaseException("参数存在SQL注入风险"); + } + } + } } diff --git a/ruoyi-generator/src/main/java/com/ruoyi/generator/controller/GenController.java b/ruoyi-generator/src/main/java/com/ruoyi/generator/controller/GenController.java index c1b87076a..7b2cf659a 100644 --- a/ruoyi-generator/src/main/java/com/ruoyi/generator/controller/GenController.java +++ b/ruoyi-generator/src/main/java/com/ruoyi/generator/controller/GenController.java @@ -31,6 +31,7 @@ import com.ruoyi.common.core.text.Convert; import com.ruoyi.common.enums.BusinessType; import com.ruoyi.common.utils.StringUtils; import com.ruoyi.common.utils.security.PermissionUtils; +import com.ruoyi.common.utils.sql.SqlUtil; import com.ruoyi.generator.domain.GenTable; import com.ruoyi.generator.domain.GenTableColumn; import com.ruoyi.generator.service.IGenTableColumnService; @@ -196,31 +197,33 @@ public class GenController extends BaseController @ResponseBody public AjaxResult create(String sql) { - List sqlStatements = SQLUtils.parseStatements(sql, DbType.mysql); - List tableNames = new ArrayList<>(); - for (SQLStatement sqlStatement : sqlStatements) + try { - if (sqlStatement instanceof MySqlCreateTableStatement) + SqlUtil.filterKeyword(sql); + List sqlStatements = SQLUtils.parseStatements(sql, DbType.mysql); + List tableNames = new ArrayList<>(); + for (SQLStatement sqlStatement : sqlStatements) { - MySqlCreateTableStatement createTableStatement = (MySqlCreateTableStatement) sqlStatement; - String tableName = createTableStatement.getTableName(); - tableName = tableName.replaceAll("`", ""); - - int msg = genTableService.createTable(createTableStatement.toString()); - if (msg == 0) + if (sqlStatement instanceof MySqlCreateTableStatement) { - tableNames.add(tableName); + MySqlCreateTableStatement createTableStatement = (MySqlCreateTableStatement) sqlStatement; + if (genTableService.createTable(createTableStatement.toString())) + { + String tableName = createTableStatement.getTableName().replaceAll("`", ""); + tableNames.add(tableName); + } } } - else - { - return AjaxResult.error("请输入建表语句"); - } + List tableList = genTableService.selectDbTableListByNames(tableNames.toArray(new String[tableNames.size()])); + String operName = Convert.toStr(PermissionUtils.getPrincipalProperty("loginName")); + genTableService.importGenTable(tableList, operName); + return AjaxResult.success(); + } + catch (Exception e) + { + logger.error(e.getMessage(), e); + return AjaxResult.error("创建表结构异常" + e.getMessage()); } - List tableList = genTableService.selectDbTableListByNames((tableNames.toArray(new String[tableNames.size()]))); - String operName = Convert.toStr(PermissionUtils.getPrincipalProperty("loginName")); - genTableService.importGenTable(tableList, operName); - return AjaxResult.success(); } /** diff --git a/ruoyi-generator/src/main/java/com/ruoyi/generator/service/IGenTableService.java b/ruoyi-generator/src/main/java/com/ruoyi/generator/service/IGenTableService.java index 40dc6759f..4af8fea26 100644 --- a/ruoyi-generator/src/main/java/com/ruoyi/generator/service/IGenTableService.java +++ b/ruoyi-generator/src/main/java/com/ruoyi/generator/service/IGenTableService.java @@ -72,7 +72,7 @@ public interface IGenTableService * @param sql 创建表语句 * @return 结果 */ - public int createTable(String sql); + public boolean createTable(String sql); /** * 导入表结构 diff --git a/ruoyi-generator/src/main/java/com/ruoyi/generator/service/impl/GenTableServiceImpl.java b/ruoyi-generator/src/main/java/com/ruoyi/generator/service/impl/GenTableServiceImpl.java index 5ace7be87..5d0c095a8 100644 --- a/ruoyi-generator/src/main/java/com/ruoyi/generator/service/impl/GenTableServiceImpl.java +++ b/ruoyi-generator/src/main/java/com/ruoyi/generator/service/impl/GenTableServiceImpl.java @@ -157,9 +157,9 @@ public class GenTableServiceImpl implements IGenTableService * @return 结果 */ @Override - public int createTable(String sql) + public boolean createTable(String sql) { - return genTableMapper.createTable(sql); + return genTableMapper.createTable(sql) == 0; } /**