diff --git a/src/main/java/app/annotations/impl/WebAccessAspect.java b/src/main/java/app/annotations/impl/WebAccessAspect.java index d0412d0..34baa03 100644 --- a/src/main/java/app/annotations/impl/WebAccessAspect.java +++ b/src/main/java/app/annotations/impl/WebAccessAspect.java @@ -17,6 +17,11 @@ import org.slf4j.LoggerFactory; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.context.annotation.Configuration; import org.springframework.core.env.Environment; +import org.springframework.web.socket.CloseStatus; +import org.springframework.web.socket.WebSocketSession; + +import java.io.IOException; +import java.util.List; @Aspect @Configuration @@ -36,6 +41,36 @@ public class WebAccessAspect { } } + @Before("@annotation(app.annotations.interfaces.CheckWebAccess) && args(session,..)") + public void before(JoinPoint joinPoint, WebSocketSession session) throws IOException { + AuthMethod auth_method = ((MethodSignature) joinPoint.getSignature()).getMethod().getAnnotation(CheckWebAccess.class).auth_method(); + + if(!(session instanceof WebSocketSession)) { + logger.error("Invalid websocket session"); + session.close(CloseStatus.NOT_ACCEPTABLE); + return; + } + + if(session.getHandshakeHeaders().get("cookie") == null) { + logger.warn("Session {} [{}] Request not contain cookies", session.getId(), session.getHandshakeHeaders().get("X-Forwarded-For")); + session.close(CloseStatus.NOT_ACCEPTABLE); + return; + } + + List rawCookie = session.getHandshakeHeaders().get("cookie"); + switch (auth_method) { + case SECRET_KEY -> { + String secret_key = rawCookie.stream().filter(s -> s.contains("secretkey=")).map(s -> s.split("=")[1]).findFirst().orElse(null); + if(!saltedCookie.ValidateSecretKey(secret_key)){ + logger.error("Invalid secret key on session {}", session.getId()); + session.close(CloseStatus.NOT_ACCEPTABLE); + } + return; + } + default -> {session.close(CloseStatus.NOT_ACCEPTABLE);return;} + } + } + @Before("@annotation(app.annotations.interfaces.CheckWebAccess) && args(request,..)") public void before(JoinPoint joinPoint, HttpServletRequest request){ if (!enabled) return; diff --git a/src/main/java/app/websocket/handlers/ExternalVIPHandler.java b/src/main/java/app/websocket/handlers/ExternalVIPHandler.java index e5b312b..59c960e 100644 --- a/src/main/java/app/websocket/handlers/ExternalVIPHandler.java +++ b/src/main/java/app/websocket/handlers/ExternalVIPHandler.java @@ -1,5 +1,7 @@ package app.websocket.handlers; +import app.annotations.enums.AuthMethod; +import app.annotations.interfaces.CheckWebAccess; import app.controllers.other.ExternalVIPController; import app.entities.VipGiveMethod; import app.services.db.VIPService; @@ -46,24 +48,19 @@ public class ExternalVIPHandler extends TextWebSocketHandler implements BaseWebs } @Override + @CheckWebAccess(auth_method = AuthMethod.SECRET_KEY) public void afterConnectionEstablished(WebSocketSession session) throws Exception { - logger.info("Session {} open", session.getId()); - if (!saltedCookie.validateHeaders(session.getHandshakeHeaders())) { - logger.warn("Session {} not acceptable, invalid secret key in header!", session.getId()); - session.close(CloseStatus.NOT_ACCEPTABLE); - return; - } super.afterConnectionEstablished(session); - session.sendMessage(new TextMessage(objectMapper.writeValueAsString(externalVIPController.getPrices()))); + if (session.isOpen()) { + logger.info("Session {} open", session.getId()); + session.sendMessage(new TextMessage(objectMapper.writeValueAsString(externalVIPController.getPrices()))); + } } @Override + @CheckWebAccess(auth_method = AuthMethod.SECRET_KEY) public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException { - if (!saltedCookie.validateHeaders(session.getHandshakeHeaders())) { - logger.warn("Session {} closed, invalid secret key in header!", session.getId()); - session.close(CloseStatus.NOT_ACCEPTABLE); - return; - } + if (!session.isOpen()) return; ExternalVIP vip; try { diff --git a/src/main/java/app/websocket/handlers/ServersHandler.java b/src/main/java/app/websocket/handlers/ServersHandler.java index 282b7a4..cee77f6 100644 --- a/src/main/java/app/websocket/handlers/ServersHandler.java +++ b/src/main/java/app/websocket/handlers/ServersHandler.java @@ -3,6 +3,7 @@ package app.websocket.handlers; import app.entities.Stats; import app.entities.server.Server; import app.websocket.BaseWebsocketHandler; +import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -14,6 +15,8 @@ import org.springframework.web.socket.WebSocketSession; import org.springframework.web.socket.handler.TextWebSocketHandler; import java.io.IOException; +import java.util.Collection; +import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -40,7 +43,7 @@ public class ServersHandler extends TextWebSocketHandler implements BaseWebsocke logger.info("Session {} open", session.getId()); activeSessions.put(session.getId(), session); super.afterConnectionEstablished(session); - session.sendMessage(new TextMessage(objectMapper.writeValueAsString(stats.getServers()))); + session.sendMessage(getServersPayload()); } @Override @@ -55,22 +58,45 @@ public class ServersHandler extends TextWebSocketHandler implements BaseWebsocke super.afterConnectionClosed(session, status); } - private void push2All(String msg) throws IOException { - for (Map.Entry otherSession : activeSessions.entrySet()) { - otherSession.getValue().sendMessage(new TextMessage(msg)); + private void push2All(TextMessage msg) throws IOException { + for (Map.Entry otherSession : Collections.synchronizedSet(activeSessions.entrySet())) { + try { + otherSession.getValue().sendMessage(msg); + } catch (Exception err) { + logger.error("Cannot send update servers status to session: {}", otherSession.getValue().getId()); + } } } + private TextMessage getServersPayload() throws JsonProcessingException { + return new TextMessage(objectMapper.writeValueAsString( + new HashMap() {{ + put("servers", stats.getServers()); + put("state", State.INIT.ordinal()); + }} + )); + } + + private TextMessage getServersPayload(String srv, Server server) throws JsonProcessingException { + return new TextMessage(objectMapper.writeValueAsString( + new HashMap() {{ + put("servers", new HashMap(){{ + put(srv, server); + }}); + put("state", State.AFTER.ordinal()); + }} + )); + } + public void pushServer(String srv, Server server) { try { - push2All( - objectMapper.writeValueAsString( - new HashMap() {{ - put("servers", new HashMap(){{ - put(srv, server);}}); - }})); + push2All(getServersPayload(srv, server)); } catch (IOException err) { logger.error("Cannot send message"); } } + + public enum State { + INIT, AFTER; + } }