fix(iot): 校验数据库数据目的表名

YunaiV
1 parent 86137492
Showing 5 changed files with 153 additions and 0 deletions
yudao-module-iot/yudao-module-iot-biz/src/main/java/cn/iocoder/yudao/module/iot/controller/admin/rule/vo/data/sink/IotDataSinkSaveReqVO.java
yudao-module-iot/yudao-module-iot-biz/src/main/java/cn/iocoder/yudao/module/iot/dal/dataobject/rule/config/IotDataSinkDatabaseConfig.java
yudao-module-iot/yudao-module-iot-biz/src/main/java/cn/iocoder/yudao/module/iot/service/rule/data/action/IotDatabaseDataRuleAction.java
yudao-module-iot/yudao-module-iot-biz/src/test/java/cn/iocoder/yudao/module/iot/controller/admin/rule/vo/data/sink/IotDataSinkSaveReqVOTest.java
yudao-module-iot/yudao-module-iot-biz/src/test/java/cn/iocoder/yudao/module/iot/service/rule/data/action/IotDatabaseDataRuleActionTest.java
@@ -5,6 +5,7 @@ import cn.iocoder.yudao.framework.common.validation.InEnum;
 import cn.iocoder.yudao.module.iot.dal.dataobject.rule.config.IotAbstractDataSinkConfig;
 import cn.iocoder.yudao.module.iot.enums.rule.IotDataSinkTypeEnum;
 import io.swagger.v3.oas.annotations.media.Schema;
+import jakarta.validation.Valid;
 import jakarta.validation.constraints.NotEmpty;
 import jakarta.validation.constraints.NotNull;
 import lombok.Data;
@@ -35,6 +36,7 @@ public class IotDataSinkSaveReqVO {
  
     @Schema(description = "数据目的配置")
     @NotNull(message = "数据目的配置不能为空")
+    @Valid
     private IotAbstractDataSinkConfig config;
  
 }
 \ No newline at end of file
 package cn.iocoder.yudao.module.iot.dal.dataobject.rule.config;
  
+import jakarta.validation.constraints.NotEmpty;
+import jakarta.validation.constraints.Pattern;
 import lombok.Data;
  
 /**
@@ -15,6 +17,13 @@ import lombok.Data;
 public class IotDataSinkDatabaseConfig extends IotAbstractDataSinkConfig {
  
     /**
+     * 目标表名的格式：只允许字母、数字、下划线，且不以数字开头；支持可选的 schema.table 格式
+     *
+     * 目的：表名会拼接到 INSERT 语句中，保存和执行时都必须校验
+     */
+    public static final String TABLE_NAME_REGEX = "^[A-Za-z_][A-Za-z0-9_]{0,127}(\\.[A-Za-z_][A-Za-z0-9_]{0,127})?$";
+
+    /**
      * JDBC 连接地址
      *
      * 例如：jdbc:mysql://localhost:3306/iot_data
@@ -37,6 +46,8 @@ public class IotDataSinkDatabaseConfig extends IotAbstractDataSinkConfig {
      *
      * 设备消息将以固定结构写入该表
      */
+    @NotEmpty(message = "目标表名不能为空")
+    @Pattern(regexp = TABLE_NAME_REGEX, message = "目标表名格式不正确，只允许字母、数字、下划线，或 schema.table 格式")
     private String tableName;
  
 }
 package cn.iocoder.yudao.module.iot.service.rule.data.action;
  
+import cn.hutool.core.util.ReUtil;
 import cn.hutool.core.util.StrUtil;
 import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
 import cn.iocoder.yudao.module.iot.core.mq.message.IotDeviceMessage;
@@ -33,6 +34,10 @@ public class IotDatabaseDataRuleAction extends
     @Override
     public void execute(IotDeviceMessage message, IotDataSinkDatabaseConfig config) throws Exception {
         try {
+            // 0. 校验表名：执行时再次校验，避免历史保存的非法表名拼接到 SQL 中
+            if (!ReUtil.isMatch(IotDataSinkDatabaseConfig.TABLE_NAME_REGEX, config.getTableName())) {
+                throw new IllegalArgumentException(StrUtil.format("数据库目的的表名({}) 不合法", config.getTableName()));
+            }
             // 1. 获取或创建 JdbcTemplate
             JdbcTemplate jdbcTemplate = getProducer(config);
  
+package cn.iocoder.yudao.module.iot.controller.admin.rule.vo.data.sink;
+
+import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
+import cn.iocoder.yudao.module.iot.dal.dataobject.rule.config.IotAbstractDataSinkConfig;
+import cn.iocoder.yudao.module.iot.dal.dataobject.rule.config.IotDataSinkDatabaseConfig;
+import cn.iocoder.yudao.module.iot.dal.dataobject.rule.config.IotDataSinkHttpConfig;
+import cn.iocoder.yudao.module.iot.enums.rule.IotDataSinkTypeEnum;
+import jakarta.validation.ConstraintViolation;
+import jakarta.validation.Validation;
+import jakarta.validation.Validator;
+import org.junit.jupiter.api.Test;
+
+import java.util.Arrays;
+import java.util.Set;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+/**
+ * {@link IotDataSinkSaveReqVO} 的参数校验单元测试
+ *
+ * @author 芋道源码
+ */
+public class IotDataSinkSaveReqVOTest {
+
+    private final Validator validator = Validation.buildDefaultValidatorFactory().getValidator();
+
+    @Test
+    public void testValidate_databaseTableNameValid() {
+        for (String tableName : Arrays.asList("iot_device_message", "iot_data.device_message")) {
+            // 准备参数
+            IotDataSinkSaveReqVO reqVO = buildReqVO(IotDataSinkTypeEnum.DATABASE.getType(), buildDatabaseConfig(tableName));
+
+            // 调用，并断言
+            assertTrue(validator.validate(reqVO).isEmpty(), tableName);
+        }
+    }
+
+    @Test
+    public void testValidate_databaseTableNameInvalid() {
+        for (String tableName : Arrays.asList(null, "", "1message", "message; DROP TABLE system_users",
+                "message -- ", "message/**/", "`message`", "a.b.c", "message(id)")) {
+            // 准备参数
+            IotDataSinkSaveReqVO reqVO = buildReqVO(IotDataSinkTypeEnum.DATABASE.getType(), buildDatabaseConfig(tableName));
+
+            // 调用
+            Set<ConstraintViolation<IotDataSinkSaveReqVO>> violations = validator.validate(reqVO);
+            // 断言：通过 @Valid 级联校验到配置的表名
+            assertTrue(!violations.isEmpty(), String.valueOf(tableName));
+            violations.forEach(violation -> assertEquals("config.tableName", violation.getPropertyPath().toString()));
+        }
+    }
+
+    @Test
+    public void testValidate_otherConfigNotAffected() {
+        // 准备参数：其它类型的配置，没有新增的校验约束
+        IotDataSinkHttpConfig config = new IotDataSinkHttpConfig();
+        config.setUrl("https://www.iocoder.cn");
+        IotDataSinkSaveReqVO reqVO = buildReqVO(IotDataSinkTypeEnum.HTTP.getType(), config);
+
+        // 调用，并断言
+        assertTrue(validator.validate(reqVO).isEmpty());
+    }
+
+    private static IotDataSinkSaveReqVO buildReqVO(Integer type, IotAbstractDataSinkConfig config) {
+        IotDataSinkSaveReqVO reqVO = new IotDataSinkSaveReqVO();
+        reqVO.setName("数据目的");
+        reqVO.setStatus(CommonStatusEnum.ENABLE.getStatus());
+        reqVO.setType(type);
+        reqVO.setConfig(config);
+        return reqVO;
+    }
+
+    private static IotDataSinkDatabaseConfig buildDatabaseConfig(String tableName) {
+        IotDataSinkDatabaseConfig config = new IotDataSinkDatabaseConfig();
+        config.setJdbcUrl("jdbc:mysql://127.0.0.1:3306/iot_data");
+        config.setTableName(tableName);
+        return config;
+    }
+
+}
 package cn.iocoder.yudao.module.iot.service.rule.data.action;
  
+import cn.iocoder.yudao.module.iot.core.mq.message.IotDeviceMessage;
+import cn.iocoder.yudao.module.iot.dal.dataobject.rule.config.IotDataSinkDatabaseConfig;
 import cn.iocoder.yudao.module.iot.enums.rule.IotDataSinkTypeEnum;
 import com.zaxxer.hikari.HikariDataSource;
 import org.junit.jupiter.api.BeforeEach;
@@ -9,9 +11,12 @@ import org.mockito.MockitoAnnotations;
 import org.springframework.jdbc.core.JdbcTemplate;
  
 import javax.sql.DataSource;
+import java.time.LocalDateTime;
+import java.util.Arrays;
  
 import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.Mockito.*;
  
 /**
@@ -39,6 +44,41 @@ class IotDatabaseDataRuleActionTest {
     }
  
     @Test
+    public void testExecute_tableNameValid() throws Exception {
+        // 准备参数：普通表名，以及 schema.table 格式
+        for (String tableName : Arrays.asList("iot_device_message", "_message01", "iot_data.device_message")) {
+            IotDataSinkDatabaseConfig config = buildConfig(tableName);
+            IotDatabaseDataRuleAction action = spy(databaseDataRuleAction);
+            doReturn(jdbcTemplate).when(action).getProducer(same(config));
+
+            // 调用
+            action.execute(buildMessage(), config);
+
+            // 断言：表名原样拼接，值使用参数绑定
+            verify(jdbcTemplate).update(eq("INSERT INTO " + tableName
+                            + " (id, device_id, tenant_id, method, report_time, data, create_time) VALUES (?, ?, ?, ?, ?, ?, NOW())"),
+                    any(), any(), any(), any(), any(), any());
+        }
+    }
+
+    @Test
+    public void testExecute_tableNameInvalid() throws Exception {
+        // 准备参数：模拟历史已保存的非法表名，执行时也必须拦截
+        for (String tableName : Arrays.asList(null, "", "1message", "device message",
+                "message; DROP TABLE system_users", "message -- ", "message/**/", "`message`", "\"message\"",
+                "message(id)", "a.b.c", ".message", "message.", "表名")) {
+            IotDataSinkDatabaseConfig config = buildConfig(tableName);
+            IotDatabaseDataRuleAction action = spy(databaseDataRuleAction);
+
+            // 调用，并断言异常
+            assertThrows(IllegalArgumentException.class, () -> action.execute(buildMessage(), config));
+            // 断言：校验在获取连接之前，不会创建连接池、执行 SQL
+            verify(action, never()).getProducer(any());
+        }
+        verifyNoInteractions(jdbcTemplate);
+    }
+
+    @Test
     public void testCloseProducer_whenHikari() throws Exception {
         // 准备：底层是 HikariDataSource
         HikariDataSource hikari = mock(HikariDataSource.class);
@@ -62,4 +102,18 @@ class IotDatabaseDataRuleActionTest {
         verifyNoInteractions(other);
     }
  
+    private static IotDataSinkDatabaseConfig buildConfig(String tableName) {
+        IotDataSinkDatabaseConfig config = new IotDataSinkDatabaseConfig();
+        config.setJdbcUrl("jdbc:mysql://127.0.0.1:3306/iot_data");
+        config.setUsername("root");
+        config.setPassword("123456");
+        config.setTableName(tableName);
+        return config;
+    }
+
+    private static IotDeviceMessage buildMessage() {
+        return new IotDeviceMessage().setId("message_01").setDeviceId(1L).setTenantId(1L)
+                .setMethod("thing.property.post").setReportTime(LocalDateTime.now());
+    }
+
 }