Upload, file address whitelist restriction

Signed-off-by: fanyinbo <[email protected]>
xtreme1-io · Feb 14, 2025 · a146e3f · a146e3f
1 parent d49e75a
commit a146e3f
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 0 deletions.
diff --git a/backend/src/main/java/ai/basic/x1/usecase/UploadDataUseCase.java b/backend/src/main/java/ai/basic/x1/usecase/UploadDataUseCase.java
@@ -134,6 +134,9 @@ public class UploadDataUseCase {
     @Value("${file.prefix.small:small}")
     private String small;
 
+    @Value("${upload.url.whitelist}")
+    private String whitelist;
+
     private static final ExecutorService executorService = ThreadUtil.newExecutor(2);
     private static final ExecutorService parseExecutorService = ThreadUtil.newExecutor(5);
 
@@ -156,6 +159,11 @@ public class UploadDataUseCase {
     @Transactional(rollbackFor = RuntimeException.class)
     public Long upload(DataInfoUploadBO dataInfoUploadBO) {
         var uploadRecordBO = uploadUseCase.createUploadRecord(dataInfoUploadBO.getFileUrl());
+        if(!checkUrlIsValid(whitelist,dataInfoUploadBO.getFileUrl())){
+            uploadUseCase.updateUploadRecordStatus(uploadRecordBO.getId(), FAILED, DATASET_DATA_FILE_URL_ILLEGAL.getMessage());
+            log.error("File url illegal,datasetId:{},userId:{},fileUrl:{}", dataInfoUploadBO.getDatasetId(), dataInfoUploadBO.getUserId(), dataInfoUploadBO.getFileUrl());
+            return uploadRecordBO.getSerialNumber();
+        }
         var boo = DecompressionFileUtils.validateUrl(dataInfoUploadBO.getFileUrl());
         if (!boo) {
             uploadUseCase.updateUploadRecordStatus(uploadRecordBO.getId(), FAILED, DATASET_DATA_FILE_URL_ERROR.getMessage());
@@ -195,6 +203,20 @@ public Long upload(DataInfoUploadBO dataInfoUploadBO) {
         return uploadRecordBO.getSerialNumber();
     }
 
+
+    public static boolean checkUrlIsValid(String whitelist, String url) {
+        if(StrUtil.isEmpty(whitelist)){
+            return true;
+        }
+        String[] substrings = whitelist.split(",");
+        for (String substring : substrings) {
+            if (url.contains(substring.trim())) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     /**
      * Download the file and unzip the file
      *

diff --git a/backend/src/main/java/ai/basic/x1/usecase/exception/UsecaseCode.java b/backend/src/main/java/ai/basic/x1/usecase/exception/UsecaseCode.java
@@ -31,6 +31,8 @@ public enum UsecaseCode {
 
     DATASET_DATA_FILE_URL_ERROR("DATASET_DATA_FILE_URL_ERROR", "File url error"),
 
+    DATASET_DATA_FILE_URL_ILLEGAL("DATASET_DATA_FILE_URL_ILLEGAL", "File url illegal"),
+
     DATASET_DATA_FILE_FORMAT_ERROR("DATASET_DATA_FILE_FORMAT_ERROR", "Incorrect file format"),
 
     DATASET_NOT_FOUND("DATASET_NOT_FOUND", "Dataset not found"),

diff --git a/backend/src/main/resources/application.yml b/backend/src/main/resources/application.yml
@@ -146,3 +146,8 @@ pointCloud:
 dataset:
   similarity:
     url: http://image-vect-visualization:5000/api/v1/calcSimilarity
+
+upload:
+  url:
+    # If not set, no check is performed.  Multiple commas separated
+    whitelist: