websocket support v3 fix 3

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

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);
+        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())) {
+        if (!session.isOpen()) return;
-            logger.warn("Session {} closed, invalid secret key in header!", session.getId());
-            session.close(CloseStatus.NOT_ACCEPTABLE);
-            return;
-        }
 
         ExternalVIP vip;
         try {

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 {
+    private void push2All(TextMessage msg) throws IOException {
-        for (Map.Entry<String, WebSocketSession> otherSession : activeSessions.entrySet()) {
+        for (Map.Entry<String, WebSocketSession> otherSession : Collections.synchronizedSet(activeSessions.entrySet())) {
-            otherSession.getValue().sendMessage(new TextMessage(msg));
+            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<String, Server>(){{
+                        put(srv, server);
+                    }});
+                    put("state", State.AFTER.ordinal());
+                }}
+        ));
     }
 
     public void pushServer(String srv, Server server) {
         try {
-            push2All(
+            push2All(getServersPayload(srv, server));
-                    objectMapper.writeValueAsString(
-                            new HashMap<String, HashMap>() {{
-                                put("servers", new HashMap<String, Server>(){{
-                                    put(srv, server);}});
-            }}));
         } catch (IOException err) {
             logger.error("Cannot send message");
         }
     }
+
+    public enum State {
+        INIT, AFTER;
+    }
 }