From 75048fa4a6b86567088955c1ca9792f848f48c0f Mon Sep 17 00:00:00 2001
From: Dallas98 <990259227@qq.com>
Date: Mon, 26 Jan 2026 17:27:06 +0800
Subject: [PATCH] feat: enhance dataset file download functionality to zip all
 files and improve path validation

---
 .../DatasetFileApplicationService.java        | 40 +++++++++++++------
 1 file changed, 27 insertions(+), 13 deletions(-)

diff --git a/backend/services/data-management-service/src/main/java/com/datamate/datamanagement/application/DatasetFileApplicationService.java b/backend/services/data-management-service/src/main/java/com/datamate/datamanagement/application/DatasetFileApplicationService.java
index c3a74070..6a51e431 100644
--- a/backend/services/data-management-service/src/main/java/com/datamate/datamanagement/application/DatasetFileApplicationService.java
+++ b/backend/services/data-management-service/src/main/java/com/datamate/datamanagement/application/DatasetFileApplicationService.java
@@ -255,7 +255,7 @@ public Resource downloadFile(String datasetId, String fileId) {
     }
 
     /**
-     * 下载文件
+     * 下载数据集所有文件为 ZIP
      */
     @Transactional(readOnly = true)
     public void downloadDatasetFileAsZip(String datasetId, HttpServletResponse response) {
@@ -263,25 +263,39 @@ public void downloadDatasetFileAsZip(String datasetId, HttpServletResponse respo
         if (Objects.isNull(dataset)) {
             throw BusinessException.of(DataManagementErrorCode.DATASET_NOT_FOUND);
         }
-        List<DatasetFile> allByDatasetId = datasetFileRepository.findAllByDatasetId(datasetId);
-        Set<String> filePaths = allByDatasetId.stream().map(DatasetFile::getFilePath).collect(Collectors.toSet());
         String datasetPath = dataset.getPath();
-        Path downloadPath = Path.of(datasetPath);
+        Path downloadPath = Paths.get(datasetPath).normalize();
+        
+        // 检查路径是否存在
+        if (!Files.exists(downloadPath) || !Files.isDirectory(downloadPath)) {
+            throw BusinessException.of(DataManagementErrorCode.DATASET_NOT_FOUND);
+        }
+        
         response.setContentType("application/zip");
-        String zipName = String.format("dataset_%s.zip",
+        String zipName = String.format("dataset_%s_%s.zip",
+                dataset.getName() != null ? dataset.getName().replaceAll("[^a-zA-Z0-9_-]", "_") : "dataset",
                 LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyyMMddHHmmss")));
-        response.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=" + zipName);
+        response.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + zipName + "\"");
+        
         try (ZipArchiveOutputStream zos = new ZipArchiveOutputStream(response.getOutputStream())) {
             try (Stream<Path> pathStream = Files.walk(downloadPath)) {
-                List<Path> allPaths = pathStream.filter(path -> path.toString().startsWith(datasetPath))
-                    .filter(path -> filePaths.stream().anyMatch(filePath -> filePath.startsWith(path.toString())))
-                    .toList();
-                for (Path path : allPaths) {
-                    addToZipFile(path, downloadPath, zos);
-                }
+                pathStream
+                    .filter(path -> {
+                        // 确保路径在数据集目录内，防止路径遍历攻击
+                        Path normalized = path.normalize();
+                        return normalized.startsWith(downloadPath);
+                    })
+                    .forEach(path -> {
+                        try {
+                            addToZipFile(path, downloadPath, zos);
+                        } catch (IOException e) {
+                            log.error("Failed to add file to zip: {}", path, e);
+                        }
+                    });
             }
+            zos.finish();
         } catch (IOException e) {
-            log.error("Failed to download files in batches.", e);
+            log.error("Failed to download dataset files as zip for dataset {}", datasetId, e);
             throw BusinessException.of(SystemErrorCode.FILE_SYSTEM_ERROR);
         }
     }