fix: 修复 session已关闭的 vnc连接未断开

1.修复 session已关闭的 vnc连接未断开
2.创建 session 时，检查 env 状态

7 files changed, 156 insertions, 19 deletions

diff --git a/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketHandler.java b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketHandler.java
index 0227935..8e082dd 100644
--- a/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketHandler.java
+++ b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketHandler.java
@@ -1,6 +1,8 @@
 package net.geedge.asw.common.config.websocket;
 
 
+import cn.hutool.core.io.IoUtil;
+import cn.hutool.core.util.ObjectUtil;
 import cn.hutool.json.JSONObject;
 import cn.hutool.log.Log;
 import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
@@ -54,42 +56,52 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
         this.envId = (String) session.getAttributes().get("envId");
         this.sessionId = (String) session.getAttributes().get("sessionId");
         this.userId = (String) session.getAttributes().get("userId");
+        Constants.ENV_WEBSOCKET_SESSION.put(sessionId, session);
     }
 
     @Override
     public void afterConnectionEstablished(WebSocketSession session) throws Exception {
         super.afterConnectionEstablished(session);
-
         this.initFieldVal(session);
-        log.info("WebSocket connectioned. after connection established open env begin...  env id: {}", envId);
+        // token
+        if (T.StrUtil.isEmpty(userId)) {
+            log.warn("Websocket token authentication failed");
+            session.close(CloseStatus.NORMAL.withReason("Websocket token authentication failed"));
+            return;
+        }
 
+        // env session
         EnvironmentSessionEntity environmentSession = environmentSessionService.getOne(new LambdaQueryWrapper<EnvironmentSessionEntity>().eq(EnvironmentSessionEntity::getId, sessionId).eq(EnvironmentSessionEntity::getStatus, 1));
         if (environmentSession == null) {
             log.warn("environment session does not exist. session id: {}", sessionId);
             session.close(CloseStatus.NORMAL.withReason("Environment session does not exist"));
             return;
         }
-
-        if (T.StrUtil.isEmpty(userId)) {
-            log.warn("Websocket token authentication failed");
-            session.close(CloseStatus.NORMAL.withReason("Websocket token authentication failed"));
-            return;
-        }
-
+        log.info("WebSocket connectioned. after connection established open environment begin...  environment id: {}", envId);
         EnvironmentEntity deviceEntity = environmentService.queryInfo(envId);
         JSONObject paramJSONObject = deviceEntity.getParamJSONObject();
 
         String urlStr = String.format("%s%s", paramJSONObject.getStr("url"), Constants.ENV_API_WEBSOCKET_PATH);
         urlStr = urlStr.replace("http", "ws");
-
-        HttpClient client = HttpClient.newHttpClient();
-        WebSocket webSocket = client.newWebSocketBuilder()
-                .buildAsync(URI.create(urlStr), new WebSocketListener(session))
-                .get();
-
-        log.info("[afterConnectionEstablished] [env server: {}]", T.JSONUtil.toJsonStr(paramJSONObject));
+        WebSocket webSocket = null;
+        try {
+            HttpClient client = HttpClient.newHttpClient();
+            webSocket = client.newWebSocketBuilder()
+                    .buildAsync(URI.create(urlStr), new WebSocketListener(session))
+                    .get();
+        } catch (Exception e) {
+            log.error(e, "Environment WebSocket connectioned. after connection established open environment error. session id: {}", sessionId);
+            if (ObjectUtil.isNotNull(webSocket)) {
+                webSocket.sendClose(WebSocket.NORMAL_CLOSURE, "Normal closure");
+            }
+            if (session != null) {
+                session.close(CloseStatus.NORMAL.withReason("Environment WebSocket connectioned. after connection established open environment error!"));
+                IoUtil.close(session);
+                Constants.ENV_WEBSOCKET_SESSION.remove(session);
+            }
+        }
+        log.info("[afterConnectionEstablished] [environment server: {}]", T.JSONUtil.toJsonStr(paramJSONObject));
         session.getAttributes().put("envWebsocket", webSocket);
-
     }
 
     // WebSocket 监听器实现
@@ -115,7 +127,7 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
 
         @Override
         public CompletionStage<?> onClose(WebSocket webSocket, int statusCode, String reason) {
-            log.info("Env webSocket connection closed, Status: " + statusCode + ", Reason: " + reason);
+            log.info("Environment webSocket connection closed, Status: " + statusCode + ", Reason: " + reason);
             return WebSocket.Listener.super.onClose(webSocket, statusCode, reason);
         }
     }
diff --git a/src/main/java/net/geedge/asw/common/util/Constants.java b/src/main/java/net/geedge/asw/common/util/Constants.java
index b5ab040..f05c198 100644
--- a/src/main/java/net/geedge/asw/common/util/Constants.java
+++ b/src/main/java/net/geedge/asw/common/util/Constants.java
@@ -1,5 +1,7 @@
 package net.geedge.asw.common.util;
 
+import org.springframework.web.socket.WebSocketSession;
+
 import java.io.File;
 import java.util.HashMap;
 import java.util.List;
@@ -74,7 +76,22 @@ public class Constants {
 	 */
 	public static final String ENV_API_WEBSOCKET_PATH = "/api/v1/env/novnc";
 
+	/**
+	 * env api stop tcpdump path
+	 */
 	public static final String ENV_API_TCPDUMP_PATH = "/api/v1/env/pcap";
 
+	/**
+	 * env api status path
+	 */
+	public static final String ENV_API_STATUS_PATH = "/api/v1/env/status";
+
+	/**
+	 * novnc websocket 连接信息对应的 env session id  用以进行主动断开服务器连接功能
+	 */
+	public static final Map<String, WebSocketSession> ENV_WEBSOCKET_SESSION = T.MapUtil.newHashMap();
+
+
+
 	public static final String EMPTY_FILE_MD5 = "d41d8cd98f00b204e9800998ecf8427e";
 }
diff --git a/src/main/java/net/geedge/asw/common/util/RCode.java b/src/main/java/net/geedge/asw/common/util/RCode.java
index 2182b40..ead37d0 100644
--- a/src/main/java/net/geedge/asw/common/util/RCode.java
+++ b/src/main/java/net/geedge/asw/common/util/RCode.java
@@ -85,6 +85,7 @@ public enum RCode {
 	ENVIRONMENT_SESSION_NOT_EXIST(601001, "environment session does not exist"),
 	ENVIRONMENT_NOT_EXIST(601002, "environment does not exist"),
 	ENVIRONMENT_USED(601003, "The environment is already in use"),
+	ENVIRONMENT_STATUS_ERROR(601004, "The environment status is unavailable"),
 
 
 
diff --git a/src/main/java/net/geedge/asw/module/environment/controller/EnvironmentController.java b/src/main/java/net/geedge/asw/module/environment/controller/EnvironmentController.java
index 7a5ad4a..d2142e1 100644
--- a/src/main/java/net/geedge/asw/module/environment/controller/EnvironmentController.java
+++ b/src/main/java/net/geedge/asw/module/environment/controller/EnvironmentController.java
@@ -26,7 +26,10 @@ import net.geedge.asw.module.workspace.entity.WorkspaceEntity;
 import net.geedge.asw.module.workspace.service.IWorkspaceService;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.*;
+import org.springframework.transaction.annotation.Transactional;
 import org.springframework.web.bind.annotation.*;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.WebSocketSession;
 
 import java.io.File;
 import java.io.FileOutputStream;
@@ -164,11 +167,22 @@ public class EnvironmentController {
 
 
     @DeleteMapping("/{envId}/session/{sessionId}")
+    @Transactional
     public R removeSession(@PathVariable("envId") String envId, @PathVariable("sessionId") String sessionId, @RequestParam String workspaceId){
         EnvironmentSessionEntity session = environmentSessionService.getById(sessionId);
+        WebSocketSession webSocketSession = Constants.ENV_WEBSOCKET_SESSION.get(sessionId);
+        // 根据 session 找到 webSocketSession，更新状态，设置结束时间
         session.setEndTimestamp(System.currentTimeMillis());
         session.setStatus(2);
         environmentSessionService.updateById(session);
+        if (T.ObjectUtil.isNotEmpty(webSocketSession)) {
+            try {
+                Constants.ENV_WEBSOCKET_SESSION.remove(sessionId);
+                webSocketSession.close(CloseStatus.NORMAL.withReason("Administrator disconnected."));
+            } catch (IOException e) {
+                log.error(e, "RemoveSession send exit prompt error sessionId: {}", sessionId);
+            }
+        }
         return R.ok();
     }
 
diff --git a/src/main/java/net/geedge/asw/module/environment/service/impl/EnvironmentSessionServiceImpl.java b/src/main/java/net/geedge/asw/module/environment/service/impl/EnvironmentSessionServiceImpl.java
index bb5a61f..3611bb2 100644
--- a/src/main/java/net/geedge/asw/module/environment/service/impl/EnvironmentSessionServiceImpl.java
+++ b/src/main/java/net/geedge/asw/module/environment/service/impl/EnvironmentSessionServiceImpl.java
@@ -5,20 +5,29 @@ import cn.hutool.log.Log;
 import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
 import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
 import net.geedge.asw.common.util.ASWException;
+import net.geedge.asw.common.util.Constants;
 import net.geedge.asw.common.util.RCode;
 import net.geedge.asw.common.util.T;
 import net.geedge.asw.module.environment.dao.EnvironmentSessionDao;
+import net.geedge.asw.module.environment.entity.EnvironmentEntity;
 import net.geedge.asw.module.environment.entity.EnvironmentSessionEntity;
+import net.geedge.asw.module.environment.service.IEnvironmentService;
 import net.geedge.asw.module.environment.service.IEnvironmentSessionService;
+import net.geedge.asw.module.environment.util.EnvironmentUtil;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Service;
 
 import java.util.List;
+import java.util.Map;
 
 @Service
 public class EnvironmentSessionServiceImpl extends ServiceImpl<EnvironmentSessionDao, EnvironmentSessionEntity> implements IEnvironmentSessionService {
 
     private static final Log log = Log.get();
 
+    @Autowired
+    private IEnvironmentService environmentService;
+
     @Override
     public EnvironmentSessionEntity saveSession(String envId, String workspaceId) {
 
@@ -30,7 +39,10 @@ public class EnvironmentSessionServiceImpl extends ServiceImpl<EnvironmentSessio
         if (T.CollectionUtil.isNotEmpty(sessionEntityList)) {
             throw new ASWException(RCode.ENVIRONMENT_USED);
         }
-
+        boolean isFree = this.checkEnvironmentStatus(envId);
+        if (!isFree) {
+            throw new ASWException(RCode.ENVIRONMENT_STATUS_ERROR);
+        }
         EnvironmentSessionEntity session = new EnvironmentSessionEntity();
         session.setEnvId(envId);
         session.setWorkspaceId(workspaceId);
@@ -42,6 +54,39 @@ public class EnvironmentSessionServiceImpl extends ServiceImpl<EnvironmentSessio
         return session;
     }
 
+    private boolean checkEnvironmentStatus(String envId) {
+        boolean isFree = true;
+        EnvironmentEntity environment = environmentService.getById(envId);
+        if (T.ObjectUtil.isNull(environment)) {
+            throw new ASWException(RCode.ENVIRONMENT_NOT_EXIST);
+        }
+        if (environment.getStatus() != 1){
+            isFree = false;
+        }
+        String resultJsonStr = T.StrUtil.EMPTY_JSON;
+        try {
+            resultJsonStr = EnvironmentUtil.requestGet(environment, Constants.ENV_API_STATUS_PATH, null, String.class);
+        }catch (Exception e){
+            log.error(e, "CheckEnvironmentStatus. request environment status api error environment: {}]", T.JSONUtil.toJsonStr(environment));
+            isFree = false;
+        }
+        log.info("CheckEnvironmentStatus. environment status api result: {}", resultJsonStr);
+
+        Map resultObj = T.JSONUtil.toBean(resultJsonStr, Map.class);
+        if (T.BooleanUtil.or(
+                T.MapUtil.isEmpty(resultObj),
+                T.ObjectUtil.notEqual(RCode.SUCCESS.getCode(), resultObj.get("code")))) {
+            isFree = false;
+        } else {
+            Map data = T.MapUtil.get(resultObj, "data", Map.class);
+            String status = T.MapUtil.getStr(data, "status");
+            if (!T.StrUtil.equalsIgnoreCase(status, "online")){
+                isFree = false;
+            }
+        }
+        return isFree;
+    }
+
     @Override
     public List<EnvironmentSessionEntity> queryListByUsed() {
         List<EnvironmentSessionEntity> sessionEntityList = this.getBaseMapper().queryListByUsed();
diff --git a/src/main/java/net/geedge/asw/module/environment/util/EnvironmentUtil.java b/src/main/java/net/geedge/asw/module/environment/util/EnvironmentUtil.java
index a725d92..9f02460 100644
--- a/src/main/java/net/geedge/asw/module/environment/util/EnvironmentUtil.java
+++ b/src/main/java/net/geedge/asw/module/environment/util/EnvironmentUtil.java
@@ -27,6 +27,7 @@ import org.springframework.web.client.RestTemplate;
 
 import java.io.IOException;
 import java.net.URI;
+import java.net.URISyntaxException;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.*;
@@ -38,6 +39,51 @@ public class EnvironmentUtil {
 	private static Log log = Log.get();
 	private static RestTemplate restTemplate;
 
+	public static <T> T requestGet(EnvironmentEntity environment, String path, String queryString, Class<T> responseType) {
+		return request(environment, HttpMethod.GET, path, queryString, null, responseType);
+	}
+
+	public static <T> T request(EnvironmentEntity environment, HttpMethod method, String path, String queryString, Object body,
+								Class<T> responseType) {
+		JSONObject jsonObject = environment.getParamJSONObject();
+		String url = jsonObject.getStr("url");
+		String token = jsonObject.getStr("token");
+
+		String urlString = UrlBuilder.of(url)
+				.setPath(UrlPath.of(path, Charset.forName("UTF-8")))
+				.setQuery(UrlQuery.of(queryString, Charset.forName("UTF-8"), false, true))
+				.setCharset(StandardCharsets.UTF_8).toString();
+
+		HttpHeaders headers = new HttpHeaders();
+		headers.add(HttpHeaders.AUTHORIZATION,token);
+		HttpEntity httpEntity = body == null ? new HttpEntity(headers) : new HttpEntity(body, headers);
+		// 发送 请求
+		return request(urlString, method, token, body, responseType);
+	}
+
+	public static <T> T request(String url, HttpMethod method, String token, Object body, Class<T> responseType) {
+		HttpHeaders headers = new HttpHeaders();
+		headers.add(HttpHeaders.AUTHORIZATION, token);
+		HttpEntity httpEntity = body == null ? new HttpEntity(headers) : new HttpEntity(body, headers);
+		// 发送 请求
+		ResponseEntity<T> exchange = null;
+		try {
+			exchange = restTemplate.exchange(new URI(url), method, httpEntity, responseType);
+		} catch (URISyntaxException e) {
+			log.error(e);
+		}
+		return exchange.getBody();
+	}
+
+	public static <T> T requestGet(String url, String token, Class<T> responseType) {
+		HttpHeaders headers = new HttpHeaders();
+		headers.add(HttpHeaders.AUTHORIZATION, token);
+		HttpEntity httpEntity = new HttpEntity(headers);
+		// 发送 请求
+		ResponseEntity<T> exchange = restTemplate.exchange(url, HttpMethod.GET, httpEntity, responseType);
+		return exchange.getBody();
+	}
+
 	/**
 	 * agent stop tcpdump
 	 * @param environment
diff --git a/src/main/resources/db/migration/R__AZ_sys_i18n.sql b/src/main/resources/db/migration/R__AZ_sys_i18n.sql
index f034215..8c2aaf2 100644
--- a/src/main/resources/db/migration/R__AZ_sys_i18n.sql
+++ b/src/main/resources/db/migration/R__AZ_sys_i18n.sql
@@ -131,6 +131,8 @@ INSERT INTO `sys_i18n`(`id`, `name`, `code`, `value`, `lang`, `remark`, `update_
 INSERT INTO `sys_i18n`(`id`, `name`, `code`, `value`, `lang`, `remark`, `update_user_id`, `update_timestamp`) VALUES (219, '401011', 'WORKSPACE_MEMBER_USER_ID_REPEAT', '工作空间用户重复', 'zh', '', 'admin', 1724030366000);
 INSERT INTO `sys_i18n`(`id`, `name`, `code`, `value`, `lang`, `remark`, `update_user_id`, `update_timestamp`) VALUES (221, '601003', 'ENVIRONMENT_USED', 'The environment is already in use', 'en', '', 'admin', 1724030366000);
 INSERT INTO `sys_i18n`(`id`, `name`, `code`, `value`, `lang`, `remark`, `update_user_id`, `update_timestamp`) VALUES (223, '601003', 'ENVIRONMENT_USED', '环境已在使用中', 'zh', '', 'admin', 1724030366000);
+INSERT INTO `sys_i18n`(`id`, `name`, `code`, `value`, `lang`, `remark`, `update_user_id`, `update_timestamp`) VALUES (225, '601004', 'ENVIRONMENT_STATUS_ERROR', 'The environment status is unavailable', 'en', '', 'admin', 1724030366000);
+INSERT INTO `sys_i18n`(`id`, `name`, `code`, `value`, `lang`, `remark`, `update_user_id`, `update_timestamp`) VALUES (227, '601004', 'ENVIRONMENT_STATUS_ERROR', '环境状态不可用', 'zh', '', 'admin', 1724030366000);
 
 
 SET FOREIGN_KEY_CHECKS = 1;