2026-08-20 11:36:27 +08:00
3 changed files with 145 additions and 55 deletions
--- a/easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/document/support/DocumentSourceLoader.java
+++ b/easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/document/support/DocumentSourceLoader.java
@@ -1,5 +1,6 @@
 package tech.easyflow.ai.document.support;

+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -7,12 +8,14 @@ import org.springframework.stereotype.Component;
 import org.springframework.util.StringUtils;
 import tech.easyflow.ai.document.exception.DocumentParseBridgeException;
 import tech.easyflow.ai.document.model.DocumentSourceRef;
+import tech.easyflow.ai.easyagentsflow.upload.WorkflowApiUploadedFileReader;
 import tech.easyflow.common.filestorage.FileStorageService;
 import tech.easyflow.common.filestorage.utils.PathGeneratorUtil;

 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URLConnection;
+import java.util.Optional;

 /**
 * 文档源加载器。
@@ -29,9 +32,29 @@ public class DocumentSourceLoader {
            LoggerFactory.getLogger(DocumentSourceLoader.class);

    private final FileStorageService fileStorageService;
+    private final WorkflowApiUploadedFileReader uploadedFileReader;

-    public DocumentSourceLoader(@Qualifier("default") FileStorageService fileStorageService) {
+    /**
+     * 创建文档源加载器。
+     *
+     * @param fileStorageService 默认文件存储服务
+     * @param uploadedFileReader 已验证的工作流 API 上传文件读取器
+     */
+    @Autowired
+    public DocumentSourceLoader(
+            @Qualifier("default") FileStorageService fileStorageService,
+            WorkflowApiUploadedFileReader uploadedFileReader) {
        this.fileStorageService = fileStorageService;
+        this.uploadedFileReader = uploadedFileReader;
+    }
+
+    /**
+     * 创建不启用工作流 API 上传识别的加载器，供同包隔离测试使用。
+     *
+     * @param fileStorageService 文件存储服务
+     */
+    DocumentSourceLoader(FileStorageService fileStorageService) {
+        this(fileStorageService, null);
    }

    /**
@@ -80,7 +103,7 @@ public class DocumentSourceLoader {

    private LoadedDocumentSource loadFromFilePath(DocumentSourceRef sourceRef, long maxBytes) {
        String fileName = resolveFileName(sourceRef);
-        try (InputStream inputStream = fileStorageService.readStream(sourceRef.getFilePath())) {
+        try (InputStream inputStream = openStoredValue(sourceRef.getFilePath(), maxBytes)) {
            byte[] contentBytes = DocumentInputStreamSupport.readBytes(inputStream, maxBytes);
            logSizeMismatch(sourceRef, contentBytes.length);
            return buildLoadedSource(
@@ -105,7 +128,7 @@ public class DocumentSourceLoader {
                                                     String remoteUrl,
                                                     long maxBytes) {
        String fileName = resolveFileName(sourceRef);
-        try (InputStream inputStream = DocumentInputStreamSupport.openRemote(remoteUrl, maxBytes)) {
+        try (InputStream inputStream = openRemoteValue(remoteUrl, maxBytes)) {
            byte[] contentBytes = DocumentInputStreamSupport.readBytes(inputStream, maxBytes);
            logSizeMismatch(sourceRef, contentBytes.length);
            return buildLoadedSource(
@@ -122,6 +145,54 @@ public class DocumentSourceLoader {
        }
    }

+    /**
+     * 优先打开经过上传记录验证的受管文件，再回退普通存储路径。
+     *
+     * @param filePath 文件路径
+     * @param maxBytes 最大允许读取字节数
+     * @return 受限输入流
+     * @throws IOException 文件无法读取时抛出
+     */
+    private InputStream openStoredValue(String filePath, long maxBytes) throws IOException {
+        Optional<InputStream> managed = openManagedValue(filePath);
+        if (managed.isPresent()) {
+            return DocumentInputStreamSupport.limit(managed.get(), maxBytes);
+        }
+        return DocumentInputStreamSupport.limit(
+                fileStorageService.readStream(filePath),
+                maxBytes);
+    }
+
+    /**
+     * 优先打开经过上传记录验证的受管 URL，再执行普通公网 URL 校验与下载。
+     *
+     * @param remoteUrl 远端 URL
+     * @param maxBytes 最大允许读取字节数
+     * @return 受限输入流
+     * @throws IOException 文件无法读取时抛出
+     */
+    private InputStream openRemoteValue(String remoteUrl, long maxBytes) throws IOException {
+        Optional<InputStream> managed = openManagedValue(remoteUrl);
+        if (managed.isPresent()) {
+            return DocumentInputStreamSupport.limit(managed.get(), maxBytes);
+        }
+        return DocumentInputStreamSupport.openRemote(remoteUrl, maxBytes);
+    }
+
+    /**
+     * 尝试打开受管工作流上传文件。
+     *
+     * @param filePath 文件路径
+     * @return 受管输入流；测试未配置读取器或普通路径时为空
+     * @throws IOException 受管文件校验或读取失败时抛出
+     */
+    private Optional<InputStream> openManagedValue(String filePath) throws IOException {
+        if (uploadedFileReader == null) {
+            return Optional.empty();
+        }
+        return uploadedFileReader.openVerified(filePath);
+    }
+
    private LoadedDocumentSource buildLoadedSource(String fileName, String contentType, Long size, byte[] contentBytes) {
        LoadedDocumentSource loadedSource = new LoadedDocumentSource();
        loadedSource.setFileName(fileName);
--- a/easyflow-modules/easyflow-module-ai/src/test/java/tech/easyflow/ai/document/service/impl/DocumentParseBridgeServiceImplTest.java
+++ b/easyflow-modules/easyflow-module-ai/src/test/java/tech/easyflow/ai/document/service/impl/DocumentParseBridgeServiceImplTest.java
@@ -221,7 +221,10 @@ public class DocumentParseBridgeServiceImplTest {
                pptxDocumentParseService,
                xlsxDocumentParseService,
                parseService,
-                new DocumentSourceLoader(new InMemoryFileStorageService()));
+                new DocumentSourceLoader(
+                        new InMemoryFileStorageService(),
+                        org.mockito.Mockito.mock(
+                                tech.easyflow.ai.easyagentsflow.upload.WorkflowApiUploadedFileReader.class)));
    }

    /**
@@ -292,7 +295,10 @@ public class DocumentParseBridgeServiceImplTest {
        private long maxBytes;

        private RecordingDocumentSourceLoader() {
-            super(new InMemoryFileStorageService());
+            super(
+                    new InMemoryFileStorageService(),
+                    org.mockito.Mockito.mock(
+                            tech.easyflow.ai.easyagentsflow.upload.WorkflowApiUploadedFileReader.class));
        }

        /**
--- a/easyflow-modules/easyflow-module-ai/src/test/java/tech/easyflow/ai/document/support/DocumentSourceLoaderTest.java
+++ b/easyflow-modules/easyflow-module-ai/src/test/java/tech/easyflow/ai/document/support/DocumentSourceLoaderTest.java
@@ -2,6 +2,8 @@ package tech.easyflow.ai.document.support;

 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.Mockito;
+import tech.easyflow.ai.easyagentsflow.upload.WorkflowApiUploadedFileReader;
 import tech.easyflow.ai.document.exception.DocumentParseBridgeException;
 import tech.easyflow.ai.document.model.DocumentSourceRef;
 import tech.easyflow.common.filestorage.FileStorageService;
@@ -9,9 +11,8 @@ import tech.easyflow.common.filestorage.FileStorageService;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.io.InputStream;
-import java.net.InetSocketAddress;
 import java.nio.charset.StandardCharsets;
-import com.sun.net.httpserver.HttpServer;
+import java.util.Optional;

 /**
 * {@link DocumentSourceLoader} 单元测试。
@@ -54,66 +55,78 @@ public class DocumentSourceLoaderTest {
    }

    /**
-     * 验证 filePath 为远端 URL 时不会误走存储读取。
+     * 验证普通 filePath 远端 URL 仍拒绝访问回环地址。
     */
    @Test
-    public void shouldPreferRemoteDownloadWhenFilePathIsRemoteUrl() throws IOException {
+    public void shouldRejectLoopbackWhenRemoteUrlIsNotManagedUpload() {
        DocumentSourceLoader loader = new DocumentSourceLoader(new FailingFileStorageService());
-        HttpServer server = HttpServer.create(new InetSocketAddress(0), 0);
-        byte[] body = "demo-pdf".getBytes(StandardCharsets.UTF_8);
-        server.createContext("/demo.pdf", exchange -> {
-            exchange.sendResponseHeaders(200, body.length);
-            exchange.getResponseBody().write(body);
-            exchange.close();
-        });
-        server.start();
+        DocumentSourceRef sourceRef = new DocumentSourceRef();
+        sourceRef.setFileName("demo.pdf");
+        sourceRef.setFilePath("http://127.0.0.1:39000/demo.pdf");
+
        try {
-            DocumentSourceRef sourceRef = new DocumentSourceRef();
-            sourceRef.setFileName("demo.pdf");
-            sourceRef.setFilePath("http://127.0.0.1:" + server.getAddress().getPort() + "/demo.pdf");
-
-            LoadedDocumentSource loadedSource = loader.load(sourceRef);
-
-            Assert.assertEquals("demo.pdf", loadedSource.getFileName());
-            Assert.assertArrayEquals(body, loadedSource.getContentBytes());
-        } finally {
-            server.stop(0);
+            loader.load(sourceRef);
+            Assert.fail("expected DocumentParseBridgeException");
+        } catch (DocumentParseBridgeException exception) {
+            Assert.assertEquals("source_load_failed", exception.getCode());
+            Assert.assertTrue(exception.getCause() instanceof java.net.UnknownHostException);
        }
    }

    /**
-     * 验证远端响应声明值不可信时仍按实际读取字节数拒绝超限内容。
+     * 验证已通过上传记录校验的内网存储 URL 会走恢复句柄读取。
     *
-     * @throws IOException 测试服务启动失败时抛出
+     * @throws IOException 测试流创建失败时抛出
     */
    @Test
-    public void shouldRejectRemoteContentThatExceedsActualByteLimit() throws IOException {
-        DocumentSourceLoader loader = new DocumentSourceLoader(new FailingFileStorageService());
-        HttpServer server = HttpServer.create(new InetSocketAddress(0), 0);
-        byte[] body = "123456789".getBytes(StandardCharsets.UTF_8);
-        server.createContext("/oversized.pdf", exchange -> {
-            exchange.sendResponseHeaders(200, 0);
-            exchange.getResponseBody().write(body);
-            exchange.close();
-        });
-        server.start();
-        try {
-            DocumentSourceRef sourceRef = new DocumentSourceRef();
-            sourceRef.setFileName("oversized.pdf");
-            sourceRef.setFilePath(
-                    "http://127.0.0.1:" + server.getAddress().getPort() + "/oversized.pdf");
-            sourceRef.setSize(1L);
+    public void shouldLoadVerifiedManagedUploadBeforeRemoteAddressGuard() throws IOException {
+        WorkflowApiUploadedFileReader uploadedFileReader =
+                Mockito.mock(WorkflowApiUploadedFileReader.class);
+        DocumentSourceLoader loader = new DocumentSourceLoader(
+                new FailingFileStorageService(),
+                uploadedFileReader);
+        String fileUrl = "http://127.0.0.1:39000/easyflow/attachment/"
+                + "workflow-api-upload/0123456789abcdef0123456789abcdef/file.pdf";
+        byte[] body = "demo-pdf".getBytes(StandardCharsets.UTF_8);
+        Mockito.when(uploadedFileReader.openVerified(fileUrl))
+                .thenReturn(Optional.of(new ByteArrayInputStream(body)));
+        DocumentSourceRef sourceRef = new DocumentSourceRef();
+        sourceRef.setFileName("demo.pdf");
+        sourceRef.setFilePath(fileUrl);

-            try {
-                loader.load(sourceRef, 8L);
-                Assert.fail("expected DocumentParseBridgeException");
-            } catch (DocumentParseBridgeException exception) {
-                Assert.assertEquals("source_load_failed", exception.getCode());
-                Assert.assertTrue(exception.getCause()
-                        instanceof DocumentInputStreamSupport.SizeLimitExceededException);
-            }
-        } finally {
-            server.stop(0);
+        LoadedDocumentSource loadedSource = loader.load(sourceRef);
+
+        Assert.assertEquals("demo.pdf", loadedSource.getFileName());
+        Assert.assertArrayEquals(body, loadedSource.getContentBytes());
+    }
+
+    /**
+     * 验证受管上传文件仍按实际读取字节数拒绝超限内容。
+     */
+    @Test
+    public void shouldRejectManagedContentThatExceedsActualByteLimit() throws IOException {
+        WorkflowApiUploadedFileReader uploadedFileReader =
+                Mockito.mock(WorkflowApiUploadedFileReader.class);
+        DocumentSourceLoader loader = new DocumentSourceLoader(
+                new FailingFileStorageService(),
+                uploadedFileReader);
+        byte[] body = "123456789".getBytes(StandardCharsets.UTF_8);
+        String fileUrl = "http://127.0.0.1:39000/easyflow/attachment/"
+                + "workflow-api-upload/0123456789abcdef0123456789abcdef/oversized.pdf";
+        Mockito.when(uploadedFileReader.openVerified(fileUrl))
+                .thenReturn(Optional.of(new ByteArrayInputStream(body)));
+        DocumentSourceRef sourceRef = new DocumentSourceRef();
+        sourceRef.setFileName("oversized.pdf");
+        sourceRef.setFilePath(fileUrl);
+        sourceRef.setSize(1L);
+
+        try {
+            loader.load(sourceRef, 8L);
+            Assert.fail("expected DocumentParseBridgeException");
+        } catch (DocumentParseBridgeException exception) {
+            Assert.assertEquals("source_load_failed", exception.getCode());
+            Assert.assertTrue(exception.getCause()
+                    instanceof DocumentInputStreamSupport.SizeLimitExceededException);
        }
    }