valite input to fix security issue

OCRTool/pom.xml

@@ -38,7 +38,7 @@
         <dependency>
             <groupId>com.google.code.gson</groupId>
             <artifactId>gson</artifactId>
-            <version>2.8.6</version>
+            <version>2.8.9</version>
         </dependency>
     </dependencies>
 

OCRTool/src/main/java/com/bonus/core/OCRRecognition.java

@@ -1,17 +1,16 @@
 package com.bonus.core;
-
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Component;
 
 import java.io.*;
 import java.nio.charset.StandardCharsets;
+import java.sql.Time;
 import java.util.ArrayList;
 import java.util.List;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-
-
+import java.util.Base64;
 /**
  * @author wangvivi
  */
@@ -22,25 +21,21 @@ public class OCRRecognition {
     private String pythonEnv;
     @Value("${recognition.scriptEnv}")
     private String scriptEnv;
+    @Value("${recognition.allowExtensions}")
+    private String allowExtensions;
+
+    private static final String BASE64_PATTERN = "^[A-Za-z0-9+/]+={0,2}$";
 
     public String extractInfo(IdentifyRecognitionParams recognitionParams){
         Process proc;
         List<String> stringList = new ArrayList<>();
         String lastElement = "";
         try {
-            logger.info("开始打印从配置里读取的值：");
             logger.info(pythonEnv);
             logger.info(scriptEnv);
-            logger.info("开始打印从请求里带过来的参数：");
-
             String typeStr = Integer.toString(recognitionParams.getType());
             logger.info(typeStr);
-
-            int type = recognitionParams.getType();
-            if (type < 0 || type >= 2) {
-                logger.info("参数检查错误");
-                return ErrorHandler.getErrorInfo(ErrorHandler.RECOGNITION_PARAM_ERROR);
-            }
+            //数据参数验证
             File file = new File(pythonEnv);
             if (!file.exists()) {
                 logger.info("pythonEnv不存在");
@@ -51,7 +46,17 @@ public class OCRRecognition {
                 logger.info("scriptEnv 不存在");
                 return ErrorHandler.getErrorInfo(ErrorHandler.PYTHON_SCRIPT_NO_EXIST);
             }
+            int type = recognitionParams.getType();
+            if (type < 0 || type >= 2) {
+                logger.info("参数检查错误");
+                return ErrorHandler.getErrorInfo(ErrorHandler.RECOGNITION_PARAM_ERROR);
+            }
+            if (!isValidParam(recognitionParams)){
+                logger.info("传入参数错误");
+                return ErrorHandler.getErrorInfo(ErrorHandler.RECOGNITION_PARAM_ERROR);
+            }
 
+            //调用Python脚本执行身份证识别任务
             String[] str = new String[]{pythonEnv,scriptEnv, String.valueOf(recognitionParams.getType())};
             Runtime runtime = Runtime.getRuntime();
             if (runtime == null) {
@@ -59,22 +64,23 @@ public class OCRRecognition {
                 return ErrorHandler.getErrorInfo(ErrorHandler.PYTHON_PROCESS_START_ERROR);
             }
             proc = runtime.exec(str);
-            logger.info("proc = runtime.exec(str)");
             OutputStream outputStream = proc.getOutputStream();
-            if (!recognitionParams.getRecognitionFrontData().isEmpty()) {
-                byte[] input = recognitionParams.getRecognitionFrontData().getBytes(StandardCharsets.UTF_8);
+            String frontData = recognitionParams.getRecognitionFrontData();
+            String backData = recognitionParams.getRecognitionBackData();
+            if (!frontData.isEmpty()) {
+                byte[] input = frontData.getBytes(StandardCharsets.UTF_8);
                 passParameter(outputStream, input);
             }
-            if (!recognitionParams.getRecognitionBackData().isEmpty()) {
+            if (!backData.isEmpty()) {
                 outputStream.write(System.lineSeparator().getBytes());
-                byte[] input = recognitionParams.getRecognitionBackData().getBytes(StandardCharsets.UTF_8);
+                byte[] input = backData.getBytes(StandardCharsets.UTF_8);
                 passParameter(outputStream, input);
             }
             outputStream.flush();
             outputStream.close();
-            logger.info("outputStream.close()");
+
+            //获取Python脚本返回值
             BufferedReader in = new BufferedReader(new InputStreamReader(proc.getInputStream(), StandardCharsets.UTF_8));
-            logger.info("BufferedReader in = new BufferedReade");
             String line;
             while ((line = in.readLine()) != null) {
                 logger.info(line);
@@ -107,6 +113,54 @@ public class OCRRecognition {
         }
     }
 
+    private boolean isValidParam(IdentifyRecognitionParams recognitionParams){
+        int type = recognitionParams.getType();
+        String frontData = recognitionParams.getRecognitionFrontData();
+        String backData = recognitionParams.getRecognitionBackData();
+        String[] allowedExtensions = allowExtensions.split(",");
+        boolean bFrontValidate;
+        boolean bBackValidate;
+        if (type == 0){
+            bFrontValidate = isAllowedFileExtension(frontData, allowedExtensions);
+            bBackValidate = isAllowedFileExtension(backData, allowedExtensions);
+            return bFrontValidate && bBackValidate;
+        }
+
+        if (type == 1){
+            bFrontValidate = isValidBase64(frontData);
+            bBackValidate = isValidBase64(backData);
+            return bFrontValidate && bBackValidate;
+        }
+
+        return false;
+    }
+
+
+    public boolean isValidBase64(String str) {
+        if (str.isEmpty()) return true;
+
+        if (!str.matches(BASE64_PATTERN)) {
+            return false;
+        }
+        try {
+            Base64.getDecoder().decode(str);
+            return true;
+        } catch (IllegalArgumentException e) {
+            return false;
+        }
+    }
+
+    public boolean isAllowedFileExtension(String filePath,  String[] allowedExtensions) {
+        if (filePath.isEmpty()) return true;
+        for (String extension : allowedExtensions) {
+            if (filePath.toLowerCase().endsWith(extension)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
+
 }
 
 

OCRTool/src/main/resources/application.yml

@@ -1,3 +1,4 @@
 recognition:
   pythonEnv: /Users/wangvivi/miniconda3/envs/ocr/bin/python
   scriptEnv: /Users/wangvivi/Desktop/Code/Component/OCRPython/maincopy.py
+  allowExtensions: .jpg,.png,.jpeg