diff --git a/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentWebSocketHandler.java b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentWebSocketHandler.java index 5a3ea5a..13a0c6f 100644 --- a/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentWebSocketHandler.java +++ b/src/main/java/net/geedge/asw/common/config/websocket/EnvironmentWebSocketHandler.java @@ -36,9 +36,9 @@ public class EnvironmentWebSocketHandler extends TextWebSocketHandler { private String sessionId; /** - * env token + * user id */ - private String token; + private String userId; private IEnvironmentService environmentService; @@ -53,7 +53,7 @@ public class EnvironmentWebSocketHandler extends TextWebSocketHandler { private void initFieldVal(WebSocketSession session) { this.envId = (String) session.getAttributes().get("envId"); this.sessionId = (String) session.getAttributes().get("sessionId"); - this.token = (String) session.getAttributes().get("token"); + this.userId = (String) session.getAttributes().get("userId"); } @Override @@ -66,21 +66,19 @@ public class EnvironmentWebSocketHandler extends TextWebSocketHandler { 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.sendMessage(new TextMessage("environment session does not exist")); - session.close(); + 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; } EnvironmentEntity deviceEntity = environmentService.queryInfo(envId); JSONObject paramJSONObject = deviceEntity.getParamJSONObject(); - if (!T.StrUtil.equals(token, paramJSONObject.getStr("token"))) { - log.warn("WebSocket connectioned error. env token exception. env id: {}, token: {}", envId, token); - session.sendMessage(new TextMessage("Token error, Please config env token")); - session.close(); - return; - } - String urlStr = String.format("%s%s", paramJSONObject.getStr("url"), Constants.ENV_API_WEBSOCKET_PATH); urlStr = urlStr.replace("http", "ws"); @@ -107,7 +105,7 @@ public class EnvironmentWebSocketHandler extends TextWebSocketHandler { public CompletionStage onBinary(WebSocket webSocket, ByteBuffer data, boolean last) { try { - // device -> asw + // env -> asw session.sendMessage(new BinaryMessage(data, true)); } catch (IOException e) { throw new RuntimeException(e); @@ -125,7 +123,7 @@ public class EnvironmentWebSocketHandler extends TextWebSocketHandler { @Override protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) { try { - // asw -> env api + // asw -> env WebSocket envSocket = (WebSocket) session.getAttributes().get("envWebsocket"); if (envSocket != null) { envSocket.sendBinary(message.getPayload(), true); diff --git a/src/main/java/net/geedge/asw/common/config/websocket/WebSocketInterceptor.java b/src/main/java/net/geedge/asw/common/config/websocket/WebSocketInterceptor.java index 3b09538..ef7b1e9 100644 --- a/src/main/java/net/geedge/asw/common/config/websocket/WebSocketInterceptor.java +++ b/src/main/java/net/geedge/asw/common/config/websocket/WebSocketInterceptor.java @@ -1,5 +1,8 @@ package net.geedge.asw.common.config.websocket; +import cn.dev33.satoken.stp.StpUtil; +import cn.hutool.log.Log; +import net.geedge.asw.common.util.T; import org.springframework.http.server.ServerHttpRequest; import org.springframework.http.server.ServerHttpResponse; import org.springframework.http.server.ServletServerHttpRequest; @@ -13,6 +16,7 @@ import java.util.Map; @Component public class WebSocketInterceptor extends HttpSessionHandshakeInterceptor { + private static final Log log = Log.get(); @Override public synchronized boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map attributes) throws Exception { @@ -23,7 +27,16 @@ public class WebSocketInterceptor extends HttpSessionHandshakeInterceptor { Map variables = template.match(servletPath); attributes.put("envId", variables.get("envId")); attributes.put("sessionId", variables.get("sessionId")); - attributes.put("token", servletRequest.getServletRequest().getParameter("token")); + + try { + String token = servletRequest.getServletRequest().getParameter("token"); + StpUtil.setTokenValue(token); + String userId = StpUtil.getLoginIdAsString(); + attributes.put("userId", userId); + }catch (Exception e){ + log.error("Websocket token authentication failed"); + attributes.put("userId", T.StrUtil.EMPTY); + } } return super.beforeHandshake(request, response, wsHandler, attributes); }