@@ -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