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 8e082dd..58554fa 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 @@ -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); } diff --git a/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentTerminalWebSocketHandler.java b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentTerminalWebSocketHandler.java new file mode 100644 index 0000000..d94e0cb --- /dev/null +++ b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentTerminalWebSocketHandler.java @@ -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().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); + } +} diff --git a/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketInterceptor.java b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentWebSocketInterceptor.java similarity index 77% rename from src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketInterceptor.java rename to src/main/java/net/geedge/asw/common/config/websocket/EnvironmentWebSocketInterceptor.java index 2b1316f..abcd09c 100644 --- a/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentNovncWebSocketInterceptor.java +++ b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentWebSocketInterceptor.java @@ -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 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 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); diff --git a/src/main/java/net/geedge/asw/common/config/websocket/WebSocketConfig.java b/src/main/java/net/geedge/asw/common/config/websocket/WebSocketConfig.java index 773875b..3010dea 100644 --- a/src/main/java/net/geedge/asw/common/config/websocket/WebSocketConfig.java +++ b/src/main/java/net/geedge/asw/common/config/websocket/WebSocketConfig.java @@ -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("*"); } } 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 f05c198..3632cf9 100644 --- a/src/main/java/net/geedge/asw/common/util/Constants.java +++ b/src/main/java/net/geedge/asw/common/util/Constants.java @@ -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 ENV_WEBSOCKET_SESSION = T.MapUtil.newHashMap(); + public static final Map ENV_NOVNC_WEBSOCKET_SESSION = T.MapUtil.newHashMap(); + + /** + * terminal websocket 连接信息对应的 env session id 用以进行主动断开服务器连接功能 + */ + public static final Map ENV_TERMINAL_WEBSOCKET_SESSION = T.MapUtil.newHashMap(); + 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 d2142e1..1212233 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 @@ -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(); }