From bd19cdd255679a2eb480804ca319e97a9f5f6296 Mon Sep 17 00:00:00 2001 From: let5sne Date: Wed, 15 Jul 2026 13:06:47 +0800 Subject: [PATCH] fix(aihr): harden zip upload queue --- .../aihr/service/AihrUploadQueueService.java | 131 ++++++++++++++---- .../service/AihrUploadQueueServiceTest.java | 80 +++++++++++ 2 files changed, 185 insertions(+), 26 deletions(-) diff --git a/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/service/AihrUploadQueueService.java b/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/service/AihrUploadQueueService.java index c43939a3..390f8d11 100644 --- a/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/service/AihrUploadQueueService.java +++ b/backend/ruoyi-modules/ruoyi-aihr/src/main/java/org/dromara/aihr/service/AihrUploadQueueService.java @@ -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 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 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 extractArchive(String stagingPath, Charset charset) throws IOException { List files = new ArrayList<>(); long unpackedBytes = 0; int entries = 0; - try (var input = Files.newInputStream(Path.of(archive.stagingPath())); - var zip = new ZipInputStream(input)) { + Set 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 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 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 files) { + files.forEach(file -> deleteQuietly(file.stagedPath())); + } + + private static String uniqueFileName(String fileName, Set 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 { diff --git a/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/service/AihrUploadQueueServiceTest.java b/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/service/AihrUploadQueueServiceTest.java index 9bc67760..5c863004 100644 --- a/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/service/AihrUploadQueueServiceTest.java +++ b/backend/ruoyi-modules/ruoyi-aihr/src/test/java/org/dromara/aihr/service/AihrUploadQueueServiceTest.java @@ -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 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; + } + } }