Skip to content

Commit d9c381a

Browse files
authored
feat: enhance dataset file download functionality to zip all files and improve path validation (#285)
1 parent 2383f63 commit d9c381a

1 file changed

Lines changed: 27 additions & 13 deletions

File tree

backend/services/data-management-service/src/main/java/com/datamate/datamanagement/application/DatasetFileApplicationService.java

Lines changed: 27 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -255,33 +255,47 @@ public Resource downloadFile(String datasetId, String fileId) {
255255
}
256256

257257
/**
258-
* 下载文件
258+
* 下载数据集所有文件为 ZIP
259259
*/
260260
@Transactional(readOnly = true)
261261
public void downloadDatasetFileAsZip(String datasetId, HttpServletResponse response) {
262262
Dataset dataset = datasetRepository.getById(datasetId);
263263
if (Objects.isNull(dataset)) {
264264
throw BusinessException.of(DataManagementErrorCode.DATASET_NOT_FOUND);
265265
}
266-
List<DatasetFile> allByDatasetId = datasetFileRepository.findAllByDatasetId(datasetId);
267-
Set<String> filePaths = allByDatasetId.stream().map(DatasetFile::getFilePath).collect(Collectors.toSet());
268266
String datasetPath = dataset.getPath();
269-
Path downloadPath = Path.of(datasetPath);
267+
Path downloadPath = Paths.get(datasetPath).normalize();
268+
269+
// 检查路径是否存在
270+
if (!Files.exists(downloadPath) || !Files.isDirectory(downloadPath)) {
271+
throw BusinessException.of(DataManagementErrorCode.DATASET_NOT_FOUND);
272+
}
273+
270274
response.setContentType("application/zip");
271-
String zipName = String.format("dataset_%s.zip",
275+
String zipName = String.format("dataset_%s_%s.zip",
276+
dataset.getName() != null ? dataset.getName().replaceAll("[^a-zA-Z0-9_-]", "_") : "dataset",
272277
LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyyMMddHHmmss")));
273-
response.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=" + zipName);
278+
response.setHeader(HttpHeaders.CONTENT_DISPOSITION, "attachment; filename=\"" + zipName + "\"");
279+
274280
try (ZipArchiveOutputStream zos = new ZipArchiveOutputStream(response.getOutputStream())) {
275281
try (Stream<Path> pathStream = Files.walk(downloadPath)) {
276-
List<Path> allPaths = pathStream.filter(path -> path.toString().startsWith(datasetPath))
277-
.filter(path -> filePaths.stream().anyMatch(filePath -> filePath.startsWith(path.toString())))
278-
.toList();
279-
for (Path path : allPaths) {
280-
addToZipFile(path, downloadPath, zos);
281-
}
282+
pathStream
283+
.filter(path -> {
284+
// 确保路径在数据集目录内,防止路径遍历攻击
285+
Path normalized = path.normalize();
286+
return normalized.startsWith(downloadPath);
287+
})
288+
.forEach(path -> {
289+
try {
290+
addToZipFile(path, downloadPath, zos);
291+
} catch (IOException e) {
292+
log.error("Failed to add file to zip: {}", path, e);
293+
}
294+
});
282295
}
296+
zos.finish();
283297
} catch (IOException e) {
284-
log.error("Failed to download files in batches.", e);
298+
log.error("Failed to download dataset files as zip for dataset {}", datasetId, e);
285299
throw BusinessException.of(SystemErrorCode.FILE_SYSTEM_ERROR);
286300
}
287301
}

0 commit comments

Comments
 (0)