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

feat:ASW-80 Environment terminal ws代理接口开发

Browse Source
This commit is contained in:
zhangshuai
2024-09-20 16:32:38 +08:00
parent e17feb87b2
commit d18baeab7c
6 changed files with 204 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -56,7 +56,7 @@ 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);
Constants.ENV_NOVNC_WEBSOCKET_SESSION.put(sessionId, session);
}
@Override
@@ -81,7 +81,7 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
EnvironmentEntity deviceEntity = environmentService.queryInfo(envId);
JSONObject paramJSONObject = deviceEntity.getParamJSONObject();
String urlStr = String.format("%s%s", paramJSONObject.getStr("url"), Constants.ENV_API_WEBSOCKET_PATH);
String urlStr = String.format("%s%s", paramJSONObject.getStr("url"), Constants.ENV_NOVNC_WEBSOCKET_PATH);
urlStr = urlStr.replace("http", "ws");
WebSocket webSocket = null;
try {
@@ -97,7 +97,7 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
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);
Constants.ENV_NOVNC_WEBSOCKET_SESSION.remove(sessionId);
}
}
log.info("[afterConnectionEstablished] [environment server: {}]", T.JSONUtil.toJsonStr(paramJSONObject));
@@ -152,6 +152,7 @@ public class EnvironmentNovncWebSocketHandler extends TextWebSocketHandler {
if (envWebsocket != null) {
envWebsocket.sendClose(WebSocket.NORMAL_CLOSURE, "Normal closure");
}
Constants.ENV_NOVNC_WEBSOCKET_SESSION.remove(sessionId);
super.afterConnectionClosed(session, status);
}

160
src/main/java/net/geedge/asw/common/config/websocket/EnvironmentTerminalWebSocketHandler.java Normal file
View File

@@ -0,0 +1,160 @@
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;
import net.geedge.asw.common.util.Constants;
import net.geedge.asw.common.util.T;
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 org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.WebSocket;
import java.util.concurrent.CompletionStage;
@Component
public class EnvironmentTerminalWebSocketHandler extends TextWebSocketHandler {
private static final Log log = Log.get();
/**
* env id
*/
private String envId;
/**
* session
*/
private String sessionId;
/**
* user id
*/
private String userId;
private IEnvironmentService environmentService;
private IEnvironmentSessionService environmentSessionService;
public EnvironmentTerminalWebSocketHandler(IEnvironmentService environmentService, IEnvironmentSessionService environmentSessionService) {
this.environmentService = environmentService;
this.environmentSessionService = environmentSessionService;
}
private void initFieldVal(WebSocketSession session) {
this.envId = (String) session.getAttributes().get("envId");
this.sessionId = (String) session.getAttributes().get("sessionId");
this.userId = (String) session.getAttributes().get("userId");
Constants.ENV_TERMINAL_WEBSOCKET_SESSION.put(sessionId, session);
}
@Override
public void afterConnectionEstablished(WebSocketSession session) throws Exception {
super.afterConnectionEstablished(session);
this.initFieldVal(session);
// 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 terminal begin... environment id: {}", envId);
EnvironmentEntity deviceEntity = environmentService.queryInfo(envId);
JSONObject paramJSONObject = deviceEntity.getParamJSONObject();
String urlStr = String.format("%s%s", paramJSONObject.getStr("url"), Constants.ENV_TERMINAL_WEBSOCKET_PATH);
urlStr = urlStr.replace("http", "ws");
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 terminal webSocket connectioned. after connection established open environment terminal error. session id: {}", sessionId);
if (ObjectUtil.isNotNull(webSocket)) {
webSocket.sendClose(WebSocket.NORMAL_CLOSURE, "Normal closure");
}
if (ObjectUtil.isNotNull(session)) {
session.close(CloseStatus.NORMAL.withReason("Environment terminal webSocket connectioned. after connection established open environment terminal error!"));
IoUtil.close(session);
Constants.ENV_TERMINAL_WEBSOCKET_SESSION.remove(sessionId);
}
}
log.info("[afterConnectionEstablished] [environment terminal url: {}]", urlStr);
session.getAttributes().put("terminalWebsocket", webSocket);
}
// WebSocket 监听器实现
private static class WebSocketListener implements WebSocket.Listener {
private WebSocketSession session;
public WebSocketListener(WebSocketSession session) {
this.session = session;
}
@Override
public CompletionStage<?> onText(WebSocket webSocket, CharSequence message, boolean last) {
try {
// env -> asw
session.sendMessage(new TextMessage(message));
} catch (IOException e) {
throw new RuntimeException(e);
}
return WebSocket.Listener.super.onText(webSocket, message, last);
}
@Override
public CompletionStage<?> onClose(WebSocket webSocket, int statusCode, String reason) {
log.info("Environment terminal webSocket connection closed, Status: " + statusCode + ", Reason: " + reason);
return WebSocket.Listener.super.onClose(webSocket, statusCode, reason);
}
}
@Override
protected void handleTextMessage(WebSocketSession session, TextMessage message) {
WebSocket terminalWebsocket = (WebSocket) session.getAttributes().get("terminalWebsocket");
try {
if (terminalWebsocket != null) {
terminalWebsocket.sendText(message.getPayload(), true);
}
} catch (Exception e) {
}
}
@Override
public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
log.info("[afterConnectionClosed] [Terminal webSocket connection closed] [uri: {}]", session.getUri());
WebSocket envWebsocket = (WebSocket) session.getAttributes().get("terminalWebsocket");
if (envWebsocket != null) {
envWebsocket.sendClose(WebSocket.NORMAL_CLOSURE, "Normal closure");
}
Constants.ENV_TERMINAL_WEBSOCKET_SESSION.remove(sessionId);
super.afterConnectionClosed(session, status);
}
}

18
src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketInterceptor.java → src/main/java/net/geedge/asw/common/config/websocket/EnvironmentWebSocketInterceptor.java
View File

@@ -9,25 +9,29 @@ import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;
import org.springframework.web.util.UriTemplate;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
@Component
public class EnvironmentNovncWebSocketInterceptor extends HttpSessionHandshakeInterceptor {
public class EnvironmentWebSocketInterceptor extends HttpSessionHandshakeInterceptor {
private static final Log log = Log.get();
private String regex = "^/api/v1/env/([^/]+)/session/([^/]+)/(novnc|terminal)$";
@Override
public synchronized boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
if (request instanceof ServletServerHttpRequest) {
ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) request;
String servletPath = servletRequest.getServletRequest().getServletPath();
UriTemplate template = new UriTemplate("/api/v1/env/{envId}/session/{sessionId}/novnc");
Map<String, String> variables = template.match(servletPath);
attributes.put("envId", variables.get("envId"));
attributes.put("sessionId", variables.get("sessionId"));
Pattern pattern = Pattern.compile(regex);
Matcher matcher = pattern.matcher(servletPath);
if (matcher.find()) {
attributes.put("envId", matcher.group(1));
attributes.put("sessionId", matcher.group(2));
}
try {
String token = servletRequest.getServletRequest().getParameter("token");
StpUtil.setTokenValue(token);

6
src/main/java/net/geedge/asw/common/config/websocket/WebSocketConfig.java
View File

@@ -21,7 +21,11 @@ public class WebSocketConfig implements WebSocketConfigurer {
@Override
public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
registry.addHandler(new EnvironmentNovncWebSocketHandler(deviceService, environmentSessionService), "/api/v1/env/{envId}/session/{sessionId}/novnc")
.addInterceptors(new EnvironmentNovncWebSocketInterceptor())
.addInterceptors(new EnvironmentWebSocketInterceptor())
.setAllowedOrigins("*");
registry.addHandler(new EnvironmentTerminalWebSocketHandler(deviceService, environmentSessionService), "/api/v1/env/{envId}/session/{sessionId}/terminal")
.addInterceptors(new EnvironmentWebSocketInterceptor())
.setAllowedOrigins("*");
}
}

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

@@ -74,7 +74,9 @@ public class Constants {
/**
* env api novnc websocket path
*/
public static final String ENV_API_WEBSOCKET_PATH = "/api/v1/env/novnc";
public static final String ENV_NOVNC_WEBSOCKET_PATH = "/api/v1/env/novnc";
public static final String ENV_TERMINAL_WEBSOCKET_PATH = "/api/v1/env/terminal";
/**
* env api stop tcpdump path
@@ -89,7 +91,13 @@ public class Constants {
/**
* novnc websocket 连接信息对应的 env session id 用以进行主动断开服务器连接功能
*/
public static final Map<String, WebSocketSession> ENV_WEBSOCKET_SESSION = T.MapUtil.newHashMap();
public static final Map<String, WebSocketSession> ENV_NOVNC_WEBSOCKET_SESSION = T.MapUtil.newHashMap();
/**
* terminal websocket 连接信息对应的 env session id 用以进行主动断开服务器连接功能
*/
public static final Map<String, WebSocketSession> ENV_TERMINAL_WEBSOCKET_SESSION = T.MapUtil.newHashMap();

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

@@ -168,20 +168,25 @@ public class EnvironmentController {
@DeleteMapping("/{envId}/session/{sessionId}")
@Transactional
public R removeSession(@PathVariable("envId") String envId, @PathVariable("sessionId") String sessionId, @RequestParam String workspaceId){
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，更新状态，设置结束时间
WebSocketSession novncSession = Constants.ENV_NOVNC_WEBSOCKET_SESSION.get(sessionId);
WebSocketSession terminalSession = Constants.ENV_TERMINAL_WEBSOCKET_SESSION.get(sessionId);
// 根据 session 找到 novncSession&terminalSession ，更新状态，设置结束时间
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);
try {
if (T.ObjectUtil.isNotEmpty(novncSession)) {
Constants.ENV_NOVNC_WEBSOCKET_SESSION.remove(sessionId);
novncSession.close(CloseStatus.NORMAL.withReason("Administrator disconnected."));
}
if (T.ObjectUtil.isNotEmpty(terminalSession)) {
Constants.ENV_TERMINAL_WEBSOCKET_SESSION.remove(sessionId);
terminalSession.close(CloseStatus.NORMAL.withReason("Administrator disconnected."));
}
} catch (IOException e) {
log.error(e, "RemoveSession send exit prompt error sessionId: {}", sessionId);
}
return R.ok();
}