gfwleak/appsketch-works-asw-controller
Archived
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
1.修复 session已关闭的 vnc连接未断开
2.创建 session 时，检查 env 状态
This commit is contained in:
zhangshuai
2024-09-12 14:21:42 +08:00
parent ed5dd781b4
commit 9706eee814
7 changed files with 157 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketHandler.java
View File

@@ -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);
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;
}
// 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;
}
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);
}
}

17
src/main/java/net/geedge/asw/common/util/Constants.java
View File

@@ -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";
}

1
src/main/java/net/geedge/asw/common/util/RCode.java
View File

@@ -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"),

14
src/main/java/net/geedge/asw/module/environment/controller/EnvironmentController.java
View File

@@ -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();
}

47
src/main/java/net/geedge/asw/module/environment/service/impl/EnvironmentSessionServiceImpl.java
View File

@@ -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();

46
src/main/java/net/geedge/asw/module/environment/util/EnvironmentUtil.java
View File

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

2
src/main/resources/db/migration/R__AZ_sys_i18n.sql
View File

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