fix(aihr): harden zip upload queue

2026-07-15 13:06:47 +08:00
parent 5e89e4930c
commit bd19cdd255
2 changed files with 185 additions and 26 deletions
@@ -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 {
@@ -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;
+        }
+    }
 }