代码生成创建表检查关键字，防止注入风险

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注入风险");
+            }
+        }
+    }
 }

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,6 +197,9 @@ public class GenController extends BaseController
     @ResponseBody
     public AjaxResult create(String sql)
     {
+        try
+        {
+            SqlUtil.filterKeyword(sql);
             List<SQLStatement> sqlStatements = SQLUtils.parseStatements(sql, DbType.mysql);
             List<String> tableNames = new ArrayList<>();
             for (SQLStatement sqlStatement : sqlStatements)
@@ -203,25 +207,24 @@ public class GenController extends BaseController
                 if (sqlStatement instanceof MySqlCreateTableStatement)
                 {
                     MySqlCreateTableStatement createTableStatement = (MySqlCreateTableStatement) sqlStatement;
-                String tableName = createTableStatement.getTableName();
+                    if (genTableService.createTable(createTableStatement.toString()))
-                tableName = tableName.replaceAll("`", "");
-
-                int msg = genTableService.createTable(createTableStatement.toString());
-                if (msg == 0)
                     {
+                        String tableName = createTableStatement.getTableName().replaceAll("`", "");
                         tableNames.add(tableName);
                     }
                 }
-            else
-            {
-                return AjaxResult.error("请输入建表语句");
             }
-        }
+            List<GenTable> tableList = genTableService.selectDbTableListByNames(tableNames.toArray(new String[tableNames.size()]));
-        List<GenTable> 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());
+        }
+    }
 
     /**
      * 预览代码

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);
 
     /**
      * 导入表结构

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;
     }
 
     /**