admin/prop-ai-hr
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(aihr): harden zip upload queue

Browse Source
This commit is contained in:
admin
2026-07-15 13:06:47 +08:00
parent 5e89e4930c
commit bd19cdd255
2 changed files with 185 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/service/AihrUploadQueueService.java
+105 -26
View File
@@ -14,6 +14,8 @@ import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import java.io.IOException;
import java.nio.charset.Charset;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.StandardCopyOption;
@@ -21,7 +23,10 @@ import java.sql.ResultSet;
import java.sql.SQLException;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Locale;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ScheduledExecutorService;
@@ -47,6 +52,8 @@ public class AihrUploadQueueService {
private static final long MAX_ARCHIVE_UNPACKED_BYTES = 2L * 1024 * 1024 * 1024;
private static final int MAX_ARCHIVE_ENTRIES = 1000;
private static final int MAX_ERROR_CHARS = 480;
// ponytail: UTF-8 + GBK covers the current Chinese sources; add an explicit charset option only if another legacy encoding appears.
private static final Charset ZIP_FALLBACK_CHARSET = Charset.forName("GBK");
/** 加工并发上限：链路含 LLM 归类与 embedding 外呼，并发过高会互相争抢配额并拖慢单文件耗时。 */
private static final int WORKER_PERMITS = 2;
/** 处理中超过该时长视为孤儿（服务重启/线程异常丢失），由清扫任务重置回待处理。需大于最慢单文件加工时长。 */
@@ -284,22 +291,53 @@ public class AihrUploadQueueService {
} else {
log.info("upload item {} finished but row was re-claimed, skip status write", id);
}
} catch (ArchiveException e) {
log.warn("upload item {} archive rejected: {}", id, e.getMessage(), e);
markFailed(tenantId, id, e.getMessage());
} catch (Exception e) {
log.warn("upload item {} process failed（处理错误已隐藏）", id);
log.warn("upload item {} process failed", id, e);
markFailed(tenantId, id, "资料加工失败，请重试或联系管理员");
}
}
private int expandArchive(ItemRow archive) throws IOException {
List<ExtractedFile> files;
try {
files = extractArchiveWithEncodingFallback(archive.stagingPath());
} catch (ArchiveException e) {
throw e;
} catch (IOException | IllegalArgumentException e) {
throw new ArchiveException("ZIP解压失败，请使用未加密的 UTF-8 或 GBK ZIP 文件后重试", e);
}
try {
enqueueExtractedFiles(archive, files);
} catch (RuntimeException e) {
deleteExtractedFiles(files);
throw new ArchiveException("ZIP文件入队失败，请重试或联系管理员", e);
}
return files.size();
}
private List<ExtractedFile> extractArchiveWithEncodingFallback(String stagingPath) throws IOException {
try {
return extractArchive(stagingPath, StandardCharsets.UTF_8);
} catch (IllegalArgumentException e) {
log.info("ZIP filenames are not UTF-8, retrying with GBK");
return extractArchive(stagingPath, ZIP_FALLBACK_CHARSET);
}
}
private List<ExtractedFile> extractArchive(String stagingPath, Charset charset) throws IOException {
List<ExtractedFile> files = new ArrayList<>();
long unpackedBytes = 0;
int entries = 0;
try (var input = Files.newInputStream(Path.of(archive.stagingPath()));
var zip = new ZipInputStream(input)) {
Set<String> usedFileNames = new HashSet<>();
try (var input = Files.newInputStream(Path.of(stagingPath));
var zip = new ZipInputStream(input, charset)) {
ZipEntry entry;
while ((entry = zip.getNextEntry()) != null) {
if (++entries > MAX_ARCHIVE_ENTRIES) {
throw new IOException("ZIP文件数量不能超过 1000 个");
throw new ArchiveException("ZIP文件数量不能超过 1000 个");
}
if (entry.isDirectory()) {
continue;
@@ -307,11 +345,11 @@ public class AihrUploadQueueService {
String entryName = entry.getName() == null ? "" : entry.getName().replace('\\', '/');
Path entryPath = Path.of(entryName).normalize();
if (entryName.isBlank() || entryPath.isAbsolute() || entryPath.startsWith("..")) {
throw new IOException("ZIP包含不安全的文件路径");
throw new ArchiveException("ZIP包含不安全的文件路径");
}
String fileName = sanitizeFileName(entryPath.getFileName().toString());
String fileName = uniqueFileName(sanitizeFileName(entryPath.getFileName().toString()), usedFileNames);
if (isZip(fileName)) {
throw new IOException("不支持嵌套 ZIP 文件");
throw new ArchiveException("不支持嵌套 ZIP 文件");
}
if (!AihrSopSeedService.supportedFile(fileName)) {
continue;
@@ -326,37 +364,68 @@ public class AihrUploadQueueService {
entryBytes += read;
unpackedBytes += read;
if (entryBytes > maxBytes || unpackedBytes > MAX_ARCHIVE_UNPACKED_BYTES) {
throw new IOException("ZIP解压后文件超过大小限制");
throw new ArchiveException("ZIP解压后文件超过大小限制");
}
output.write(buffer, 0, read);
}
} catch (IOException | RuntimeException e) {
deleteQuietly(extracted);
throw e;
}
files.add(new ExtractedFile(fileName, extracted));
}
} catch (IOException e) {
files.forEach(file -> deleteQuietly(file.stagedPath()));
} catch (IOException | RuntimeException e) {
deleteExtractedFiles(files);
throw e;
}
if (files.isEmpty()) {
throw new IOException("ZIP内没有支持的资料文件");
throw new ArchiveException("ZIP内没有支持的资料文件");
}
try {
for (ExtractedFile file : files) {
jdbcTemplate.update("""
insert into aihr_knowledge_upload_item
(tenant_id, batch_id, file_name, category, staging_path, status, create_time, update_time)
values (?, ?, ?, ?, ?, 0, now(), now())
""", archive.tenantId(), archive.batchId(), file.fileName(), archive.category(), file.stagedPath().toString());
return files;
}
private void enqueueExtractedFiles(ItemRow archive, List<ExtractedFile> files) {
StringBuilder sql = new StringBuilder("""
insert into aihr_knowledge_upload_item
(tenant_id, batch_id, file_name, category, staging_path, status, create_time, update_time)
values
""");
List<Object> parameters = new ArrayList<>(files.size() * 5);
for (int index = 0; index < files.size(); index++) {
if (index > 0) {
sql.append(',');
}
} catch (RuntimeException e) {
for (ExtractedFile file : files) {
jdbcTemplate.update("delete from aihr_knowledge_upload_item where tenant_id = ? and staging_path = ?",
archive.tenantId(), file.stagedPath().toString());
deleteQuietly(file.stagedPath());
}
throw new IOException("ZIP文件入队失败", e);
sql.append(" (?, ?, ?, ?, ?, 0, now(), now())");
ExtractedFile file = files.get(index);
parameters.add(archive.tenantId());
parameters.add(archive.batchId());
parameters.add(file.fileName());
parameters.add(archive.category());
parameters.add(file.stagedPath().toString());
}
return files.size();
jdbcTemplate.update(sql.toString(), parameters.toArray());
}
private static void deleteExtractedFiles(List<ExtractedFile> files) {
files.forEach(file -> deleteQuietly(file.stagedPath()));
}
private static String uniqueFileName(String fileName, Set<String> usedFileNames) {
String candidate = fileName;
int suffix = 2;
while (!usedFileNames.add(candidate.toLowerCase(Locale.ROOT))) {
candidate = fileNameWithSuffix(fileName, suffix++);
}
return candidate;
}
private static String fileNameWithSuffix(String fileName, int suffix) {
int extensionStart = fileName.lastIndexOf('.');
String extension = extensionStart > 0 ? fileName.substring(extensionStart) : "";
String stem = extensionStart > 0 ? fileName.substring(0, extensionStart) : fileName;
String marker = "(" + suffix + ")";
int maxStemLength = Math.max(1, 200 - extension.length() - marker.length());
return stem.substring(0, Math.min(stem.length(), maxStemLength)) + marker + extension;
}
private static boolean isZip(String fileName) {
@@ -366,6 +435,16 @@ public class AihrUploadQueueService {
private record ExtractedFile(String fileName, Path stagedPath) {
}
private static final class ArchiveException extends IOException {
private ArchiveException(String message) {
super(message);
}
private ArchiveException(String message, Throwable cause) {
super(message, cause);
}
}
/** 清扫卡在「处理中」的孤儿条目（重启/异常导致），重置回待处理后再触发一轮加工；顺带回收过期暂存文件。 */
private void sweepAndTrigger() {
try {
backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/service/AihrUploadQueueServiceTest.java
+80
View File
@@ -2,16 +2,66 @@ package org.dromara.aihr.service;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import org.springframework.jdbc.core.JdbcTemplate;
import org.mockito.ArgumentCaptor;
import java.io.IOException;
import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.nio.charset.Charset;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.zip.ZipEntry;
import java.util.zip.ZipOutputStream;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@Tag("dev")
class AihrUploadQueueServiceTest {
@TempDir
Path tempDir;
@Test
void expandsLegacyGbkZipAtomicallyAndKeepsDuplicateEntries() throws Exception {
Path staging = Files.createDirectories(tempDir.resolve("staging"));
Path archive = tempDir.resolve("legacy.zip");
writeZip(archive, Charset.forName("GBK"), "目录一/附件.docx", "目录二/附件.docx");
JdbcTemplate jdbcTemplate = mock(JdbcTemplate.class);
when(jdbcTemplate.update(anyString(), any(Object[].class))).thenReturn(1);
ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
try {
AihrUploadQueueService service = new AihrUploadQueueService(null, jdbcTemplate, executor, staging.toString());
int extracted = assertDoesNotThrow(() -> expandArchive(service, itemRow(archive)));
assertEquals(2, extracted);
ArgumentCaptor<Object[]> parameters = ArgumentCaptor.forClass(Object[].class);
verify(jdbcTemplate, times(1)).update(anyString(), parameters.capture());
Object[] values = parameters.getValue();
assertEquals(10, values.length);
assertEquals("附件.docx", values[2]);
assertEquals("附件(2).docx", values[7]);
assertTrue(Files.isRegularFile(Path.of(values[4].toString())));
assertTrue(Files.isRegularFile(Path.of(values[9].toString())));
} finally {
executor.shutdownNow();
}
}
@Test
void queueKeepsTenantWhenWorkerLeavesLoginThread() throws Exception {
Path source = Path.of("src/main/java/org/dromara/aihr/service/AihrUploadQueueService.java");
@@ -31,4 +81,34 @@ class AihrUploadQueueServiceTest {
assertFalse(code.contains("\"暂存文件写入失败：\" + e.getMessage()"));
assertFalse(code.contains("where tenant_id = ? and status = 1\n and ("));
}
private static void writeZip(Path archive, Charset charset, String... entries) throws IOException {
try (ZipOutputStream output = new ZipOutputStream(Files.newOutputStream(archive), charset)) {
for (String entry : entries) {
output.putNextEntry(new ZipEntry(entry));
output.write("fixture".getBytes());
output.closeEntry();
}
}
}
private static Object itemRow(Path archive) throws ReflectiveOperationException {
Class<?> rowType = Class.forName("org.dromara.aihr.service.AihrUploadQueueService$ItemRow");
Constructor<?> constructor = rowType.getDeclaredConstructors()[0];
constructor.setAccessible(true);
return constructor.newInstance("000000", 1L, "batch", "legacy.zip", "__auto__", archive.toString(), 1);
}
private static int expandArchive(AihrUploadQueueService service, Object itemRow) throws Exception {
Method method = service.getClass().getDeclaredMethod("expandArchive", itemRow.getClass());
method.setAccessible(true);
try {
return (int) method.invoke(service, itemRow);
} catch (InvocationTargetException e) {
if (e.getCause() instanceof Exception cause) {
throw cause;
}
throw e;
}
}
}