enhance: revalidation after ending cooperation

zyyzyykk · Feb 27, 2025 · afecfcb · afecfcb
1 parent 8ae2b12
commit afecfcb
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 3 deletions.
diff --git a/backend/terminal/src/main/java/com/kkbpro/terminal/consumer/WebSocketServer.java b/backend/terminal/src/main/java/com/kkbpro/terminal/consumer/WebSocketServer.java
@@ -100,9 +100,11 @@ public void onOpen(Session sessionSocket, @PathParam("env") String env) throws I
             Integer state = ResultCodeEnum.COOPERATE_KEY_INVALID.getState();
             String msg = ResultCodeEnum.COOPERATE_KEY_INVALID.getDesc();
             try {
-                String sshKey = AesUtil.aesDecrypt(StringUtil.changeStrBase64(cooperateKey), AdvanceController.COOPERATE_SECRET_KEY);
+                String[] keyInfo = AesUtil.aesDecrypt(StringUtil.changeStrBase64(cooperateKey), AdvanceController.COOPERATE_SECRET_KEY).split("\\^");
+                String cooperateId = keyInfo[0];
+                String sshKey = keyInfo[1];
                 WebSocketServer webSocketServer = WebSocketServer.webSocketServerMap.get(sshKey);
-                if(webSocketServer == null || webSocketServer.cooperateInfo == null)
+                if(webSocketServer == null || webSocketServer.cooperateInfo == null || !webSocketServer.cooperateInfo.getId().equals(cooperateId))
                     throw new RuntimeException();
 
                 List<Session> sessions = WebSocketServer.cooperateMap.computeIfAbsent(sshKey, k -> new ArrayList<>());

diff --git a/backend/terminal/src/main/java/com/kkbpro/terminal/controller/AdvanceController.java b/backend/terminal/src/main/java/com/kkbpro/terminal/controller/AdvanceController.java
@@ -17,6 +17,7 @@
 import java.io.IOException;
 import java.io.InputStreamReader;
 import java.util.List;
+import java.util.UUID;
 
 /**
  * 高级功能接口类
@@ -64,12 +65,14 @@ public Result cooperateKey(String sshKey, Boolean readOnly, Integer maxHeadCount
         if(webSocketServer.getCooperateInfo() != null)
             return Result.error(errorMsg);
 
+        String cooperateId = UUID.randomUUID().toString();
         CooperateInfo cooperateInfo = new CooperateInfo();
+        cooperateInfo.setId(cooperateId);
         cooperateInfo.setReadOnly(readOnly);
         cooperateInfo.setMaxHeadCount(maxHeadCount);
         webSocketServer.setCooperateInfo(cooperateInfo);
 
-        String key = StringUtil.changeBase64Str(AesUtil.aesEncrypt(sshKey, COOPERATE_SECRET_KEY));
+        String key = StringUtil.changeBase64Str(AesUtil.aesEncrypt(cooperateId + "^" + sshKey, COOPERATE_SECRET_KEY));
 
         return Result.success(successMsg, key);
     }

diff --git a/backend/terminal/src/main/java/com/kkbpro/terminal/pojo/dto/CooperateInfo.java b/backend/terminal/src/main/java/com/kkbpro/terminal/pojo/dto/CooperateInfo.java
@@ -12,6 +12,8 @@
 @AllArgsConstructor
 public class CooperateInfo {
 
+    private String id;
+
     private Boolean readOnly;
 
     private Integer maxHeadCount;