wangpeng/yidaima
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Harden AI database generation safety

Browse Source
This commit is contained in:
王鹏
2026-05-22 14:55:17 +08:00
parent 4936d31d3b
commit 948e664f22
6 changed files with 69 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
RuoYi-Vue/ruoyi-generator/src/main/java/com/ruoyi/generator/service/front/AiDatabaseSchemaValidator.java
View File

@@ -5,6 +5,7 @@ import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.springframework.stereotype.Component;
import com.ruoyi.common.exception.ServiceException;
@@ -21,7 +22,7 @@ public class AiDatabaseSchemaValidator
private static final Pattern INT = Pattern.compile("^int(\\(11\\))?$");
private static final Pattern VARCHAR = Pattern.compile("^varchar\\(([1-9][0-9]{0,2}|1000)\\)$");
private static final Pattern CHAR = Pattern.compile("^char\\(([1-9][0-9]?)\\)$");
private static final Pattern DECIMAL = Pattern.compile("^decimal\\(([1-9][0-9]?),([1-9][0-9]?)\\)$");
private static final Pattern DECIMAL = Pattern.compile("^decimal\\(([1-9][0-9]?),([0-9][0-9]?)\\)$");
private static final Set<String> EXACT_TYPES = new HashSet<String>(Arrays.asList("datetime", "date", "text", "longtext"));
private static final List<String> DANGEROUS_SQL = Arrays.asList(
"drop", "truncate", "alter user", "grant", "revoke", "load_file", "into outfile");
@@ -39,7 +40,19 @@ public class AiDatabaseSchemaValidator
}
String type = columnType.trim().toLowerCase(Locale.ENGLISH);
return EXACT_TYPES.contains(type) || BIGINT.matcher(type).matches() || INT.matcher(type).matches()
|| VARCHAR.matcher(type).matches() || CHAR.matcher(type).matches() || DECIMAL.matcher(type).matches();
|| VARCHAR.matcher(type).matches() || CHAR.matcher(type).matches() || isAllowedDecimal(type);
}
private boolean isAllowedDecimal(String type)
{
Matcher matcher = DECIMAL.matcher(type);
if (!matcher.matches())
{
return false;
}
int precision = Integer.parseInt(matcher.group(1));
int scale = Integer.parseInt(matcher.group(2));
return scale <= precision;
}
public boolean isSafeSql(String sql)

18
RuoYi-Vue/ruoyi-generator/src/main/java/com/ruoyi/generator/service/front/AiGenerateServiceImpl.java
View File

@@ -25,6 +25,8 @@ import com.ruoyi.generator.mapper.front.FrontProjectTableMapper;
@Service
public class AiGenerateServiceImpl implements IAiGenerateService
{
private static final int MAX_ERROR_MESSAGE_LENGTH = 1000;
@Autowired
private IDeepSeekClient deepSeekClient;
@Autowired
@@ -68,7 +70,7 @@ public class AiGenerateServiceImpl implements IAiGenerateService
catch (RuntimeException e)
{
generation.setSuccess("0");
generation.setErrorMessage(e.getMessage());
generation.setErrorMessage(sanitizeErrorMessage(e.getMessage()));
generation.setElapsedMs(System.currentTimeMillis() - start);
generationRecordService.insert(generation);
throw e;
@@ -208,6 +210,20 @@ public class AiGenerateServiceImpl implements IAiGenerateService
return generation;
}
private String sanitizeErrorMessage(String message)
{
if (message == null)
{
return null;
}
String sanitized = message.replaceAll("[\\p{Cntrl}]", " ");
if (sanitized.length() <= MAX_ERROR_MESSAGE_LENGTH)
{
return sanitized;
}
return sanitized.substring(0, MAX_ERROR_MESSAGE_LENGTH);
}
private FrontProjectTable toFrontProjectTable(Long projectId, Long userId, DatabaseTableDesign source)
{
FrontProjectTable table = new FrontProjectTable();

5
RuoYi-Vue/ruoyi-generator/src/main/java/com/ruoyi/generator/service/front/MysqlDdlBuilder.java
View File

@@ -53,6 +53,9 @@ public class MysqlDdlBuilder
private String escapeComment(String comment)
{
return StringUtils.defaultString(comment).replace("'", "''");
return StringUtils.defaultString(comment)
.replaceAll("[\\p{Cntrl}]", " ")
.replace("\\", "\\\\")
.replace("'", "''");
}
}

2
RuoYi-Vue/ruoyi-generator/src/test/java/com/ruoyi/generator/service/front/AiDatabaseSchemaValidatorTest.java
View File

@@ -17,6 +17,8 @@ public class AiDatabaseSchemaValidatorTest
assertFalse(validator.isValidName("sys user"));
assertTrue(validator.isAllowedColumnType("varchar(100)"));
assertTrue(validator.isAllowedColumnType("decimal(10,2)"));
assertTrue(validator.isAllowedColumnType("decimal(10,0)"));
assertFalse(validator.isAllowedColumnType("decimal(5,9)"));
assertFalse(validator.isAllowedColumnType("json"));
assertFalse(validator.isSafeSql("drop table sys_user"));
}

16
RuoYi-Vue/ruoyi-generator/src/test/java/com/ruoyi/generator/service/front/AiGenerateServiceImplTest.java
View File

@@ -1,6 +1,7 @@
package com.ruoyi.generator.service.front;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;
@@ -90,7 +91,7 @@ public class AiGenerateServiceImplTest
@Test
public void generateDatabasePersistsFailureGenerationRecordBeforePropagatingClientFailure()
{
RuntimeException failure = new RuntimeException("DeepSeek unavailable");
RuntimeException failure = new RuntimeException("DeepSeek unavailable\n" + repeated("x", 1200));
when(frontProjectMapper.selectFrontProjectByUserAndId(7L, 10L)).thenReturn(project());
when(deepSeekClient.chat(anyString())).thenThrow(failure);
@@ -107,7 +108,8 @@ public class AiGenerateServiceImplTest
ArgumentCaptor<FrontProjectGeneration> generationCaptor = ArgumentCaptor.forClass(FrontProjectGeneration.class);
verify(frontProjectGenerationMapper).insertFrontProjectGeneration(generationCaptor.capture());
assertEquals("0", generationCaptor.getValue().getSuccess());
assertEquals("DeepSeek unavailable", generationCaptor.getValue().getErrorMessage());
assertTrue(generationCaptor.getValue().getErrorMessage().length() <= 1000);
assertFalse(generationCaptor.getValue().getErrorMessage().contains("\n"));
verify(frontProjectColumnMapper, never()).deleteColumnsByProjectId(10L);
verify(frontProjectTableMapper, never()).deleteTablesByProjectId(10L);
}
@@ -155,6 +157,16 @@ public class AiGenerateServiceImplTest
return request;
}
private String repeated(String value, int count)
{
StringBuilder builder = new StringBuilder();
for (int i = 0; i < count; i++)
{
builder.append(value);
}
return builder.toString();
}
private String mockResponse()
{
return "{\n"

17
RuoYi-Vue/ruoyi-generator/src/test/java/com/ruoyi/generator/service/front/MysqlDdlBuilderTest.java
View File

@@ -1,5 +1,6 @@
package com.ruoyi.generator.service.front;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import java.util.Arrays;
@@ -32,6 +33,22 @@ public class MysqlDdlBuilderTest
assertTrue(sql.endsWith("COMMENT='车辆''信息表';"));
}
@Test
public void escapesBackslashQuoteSequencesInsideComments()
{
DatabaseTableDesign table = new DatabaseTableDesign();
table.setTableName("car_info");
table.setTableComment("safe \\'; update front_project set project_name='owned'");
table.setColumns(Arrays.asList(
column("car_id", "bigint(20)", "1", "1", "1", "safe \\'; update front_project set project_name='owned'")));
String sql = builder.build(table);
assertFalse(sql.contains("\\'; update"));
assertTrue(sql.contains("COMMENT 'safe \\\\''; update front_project set project_name=''owned'''"));
assertTrue(sql.endsWith("COMMENT='safe \\\\''; update front_project set project_name=''owned''';"));
}
private DatabaseColumnDesign column(String name, String type, String isPk, String isIncrement,
String isRequired, String comment)
{