From 4dcba9f3743c5cfe3bd33e4192665043fb898fd7 Mon Sep 17 00:00:00 2001
From: gsd <mamu@ebal.net>
Date: Wed, 10 Jan 2024 18:54:27 +0300
Subject: [PATCH] websocket support v3

---
 .../app/controllers/debug/WebSocketPage.java  |   4 +-
 .../other/ExternalVIPController.java          |  11 +-
 .../server/ServerUpdaterController.java       |  12 +-
 src/main/java/app/utils/SaltedCookie.java     |  12 +-
 .../app/websocket/BaseWebsocketHandler.java   |   5 +
 .../IntegrationWebSockerHandler.java          | 113 -----------------
 .../java/app/websocket/WebSocketConfig.java   |  28 -----
 .../configuration/WebSocketConfig.java        |  30 +++++
 .../handlers/ExternalVIPHandler.java          | 115 ++++++++++++++++++
 .../ServersHandler.java}                      |  24 +++-
 10 files changed, 193 insertions(+), 161 deletions(-)
 create mode 100644 src/main/java/app/websocket/BaseWebsocketHandler.java
 delete mode 100644 src/main/java/app/websocket/IntegrationWebSockerHandler.java
 delete mode 100644 src/main/java/app/websocket/WebSocketConfig.java
 create mode 100644 src/main/java/app/websocket/configuration/WebSocketConfig.java
 create mode 100644 src/main/java/app/websocket/handlers/ExternalVIPHandler.java
 rename src/main/java/app/websocket/{WebSocketHandler.java => handlers/ServersHandler.java} (79%)
diff --git a/src/main/java/app/controllers/debug/WebSocketPage.java b/src/main/java/app/controllers/debug/WebSocketPage.java
index a56faed..150792b 100644
--- a/src/main/java/app/controllers/debug/WebSocketPage.java
+++ b/src/main/java/app/controllers/debug/WebSocketPage.java
@@ -10,7 +10,7 @@ import org.springframework.web.bind.annotation.RestController;
 public class WebSocketPage {
 
     @GetMapping
-    public String getPage(@RequestParam(required = false, defaultValue = "ws") String route){
+    public String getPage(@RequestParam(required = false, defaultValue = "test") String route){
         return "<!DOCTYPE html>\n" +
                 "<html lang=\"en\">\n" +
                 "<head>\n" +
@@ -57,7 +57,7 @@ public class WebSocketPage {
                 "\n" +
                 "\n" +
                 "<script>\n" +
-                "    const socketConn = new WebSocket('ws://localhost:8080/" + route + "');\n" +
+                "    const socketConn = new WebSocket('ws://localhost:8080/ws/" + route + "');\n" +
                 "\n" +
                 "    function send() {\n" +
                 "        const clientMsg = document.getElementById('myMessage');\n" +
diff --git a/src/main/java/app/controllers/other/ExternalVIPController.java b/src/main/java/app/controllers/other/ExternalVIPController.java
index ba35370..2243a4f 100644
--- a/src/main/java/app/controllers/other/ExternalVIPController.java
+++ b/src/main/java/app/controllers/other/ExternalVIPController.java
@@ -15,6 +15,7 @@ import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.*;
 
 import java.util.*;
+import java.util.stream.Collectors;
 
 @RestController
 @RequestMapping("api/external/vip")
@@ -66,10 +67,14 @@ public class ExternalVIPController {
     }
 
     @GetMapping
-    public ResponseEntity getPrices() {
-        return new ResponseEntity(prices.stream()
+    public ResponseEntity getPricesResponse() {
+        return new ResponseEntity(getPrices(), HttpStatus.OK);
+    }
+
+    public List<VipPrice> getPrices() {
+        return prices.stream()
                 .peek(price -> price.setQiwi(qiwi))
-                .peek(price -> price.setSteam(steam)), HttpStatus.OK);
+                .peek(price -> price.setSteam(steam)).collect(Collectors.toList());
     }
 
     @GetMapping("/change")
diff --git a/src/main/java/app/controllers/server/ServerUpdaterController.java b/src/main/java/app/controllers/server/ServerUpdaterController.java
index 06b0f0c..56171c9 100644
--- a/src/main/java/app/controllers/server/ServerUpdaterController.java
+++ b/src/main/java/app/controllers/server/ServerUpdaterController.java
@@ -4,7 +4,7 @@ import app.annotations.enums.AuthMethod;
 import app.annotations.interfaces.CheckWebAccess;
 import app.entities.Stats;
 import app.entities.server.request.ServerRequestBody;
-import app.websocket.WebSocketHandler;
+import app.websocket.handlers.ServersHandler;
 import jakarta.servlet.http.HttpServletRequest;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.http.HttpStatus;
@@ -15,12 +15,12 @@ import org.springframework.web.bind.annotation.*;
 @RequestMapping("api/server")
 public class ServerUpdaterController {
     private Stats stats;
-    private WebSocketHandler webSocketHandler;
+    private ServersHandler serversHandler;
 
     @Autowired
-    public ServerUpdaterController(Stats stats, WebSocketHandler webSocketHandler) {
+    public ServerUpdaterController(Stats stats, ServersHandler serversHandler) {
         this.stats = stats;
-        this.webSocketHandler = webSocketHandler;
+        this.serversHandler = serversHandler;
     }
 
     @PostMapping(value = "/{srv}")
@@ -28,7 +28,7 @@ public class ServerUpdaterController {
     public ResponseEntity updateServer(HttpServletRequest request, @PathVariable String srv, @RequestBody ServerRequestBody serverRequestBody) {
         if (!stats.getServers().containsKey(srv)) return new ResponseEntity<>(HttpStatus.NOT_FOUND);
         stats.getServers().get(srv).RefreshServerFromRequest(serverRequestBody);
-        webSocketHandler.pushServer(srv, stats.getServers().get(srv));
+        serversHandler.pushServer(srv, stats.getServers().get(srv));
         return new ResponseEntity(HttpStatus.OK);
     }
 
@@ -37,7 +37,7 @@ public class ServerUpdaterController {
     public ResponseEntity downServer(HttpServletRequest request, @PathVariable String srv) {
         if (!stats.getServers().containsKey(srv)) return new ResponseEntity<>(HttpStatus.NOT_FOUND);
         stats.getServers().get(srv).RefreshServerFromRequest(null);
-        webSocketHandler.pushServer(srv, stats.getServers().get(srv));
+        serversHandler.pushServer(srv, stats.getServers().get(srv));
         return new ResponseEntity(HttpStatus.OK);
     }
 }
diff --git a/src/main/java/app/utils/SaltedCookie.java b/src/main/java/app/utils/SaltedCookie.java
index 7277fad..3cb7dcf 100644
--- a/src/main/java/app/utils/SaltedCookie.java
+++ b/src/main/java/app/utils/SaltedCookie.java
@@ -1,16 +1,15 @@
 package app.utils;
 
-import app.updates.CountriesUpdater;
-import jakarta.annotation.PostConstruct;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
+import org.springframework.http.HttpHeaders;
 import org.springframework.stereotype.Component;
 import org.springframework.util.DigestUtils;
 
-import java.util.Arrays;
 import java.util.HashMap;
+import java.util.List;
 
 @Component
 public class SaltedCookie {
@@ -79,4 +78,11 @@ public class SaltedCookie {
         }
         return false;
     }
+
+    public boolean validateHeaders(HttpHeaders httpHeaders) {
+        List<String> cookies = httpHeaders.get("cookie");
+        if (cookies == null) return false;
+        String secret_key = cookies.stream().filter(s -> s.contains("secretkey=")).map(s -> s.split("=")[1]).findFirst().orElse(null);
+        return ValidateSecretKey(secret_key);
+    }
 }
diff --git a/src/main/java/app/websocket/BaseWebsocketHandler.java b/src/main/java/app/websocket/BaseWebsocketHandler.java
new file mode 100644
index 0000000..6fbd5e7
--- /dev/null
+++ b/src/main/java/app/websocket/BaseWebsocketHandler.java
@@ -0,0 +1,5 @@
+package app.websocket;
+
+public interface BaseWebsocketHandler {
+    String getPath();
+}
diff --git a/src/main/java/app/websocket/IntegrationWebSockerHandler.java b/src/main/java/app/websocket/IntegrationWebSockerHandler.java
deleted file mode 100644
index 9f082ce..0000000
--- a/src/main/java/app/websocket/IntegrationWebSockerHandler.java
+++ /dev/null
@@ -1,113 +0,0 @@
-package app.websocket;
-
-import app.entities.VipGiveMethod;
-import app.services.db.VIPService;
-import app.utils.SaltedCookie;
-import app.utils.SteamIDConverter;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import lombok.Data;
-import lombok.NoArgsConstructor;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-import org.springframework.beans.factory.annotation.Autowired;
-import org.springframework.context.annotation.Scope;
-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.time.Instant;
-import java.util.HashSet;
-
-@Component
-@Scope(value = "singleton")
-public class IntegrationWebSockerHandler extends TextWebSocketHandler {
-    private final Logger logger = LoggerFactory.getLogger(getClass());
-    private ObjectMapper objectMapper = new ObjectMapper();
-    private SaltedCookie saltedCookie;
-    private VIPService vipService;
-    private HashSet<String> unique_set;
-
-    @Autowired
-    public IntegrationWebSockerHandler(SaltedCookie saltedCookie, VIPService vipService) {
-        super();
-        this.saltedCookie = saltedCookie;
-        this.vipService = vipService;
-        this.unique_set = new HashSet();
-    }
-
-    @Override
-    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
-        logger.info("Session {} open", session.getId());
-        super.afterConnectionEstablished(session);
-    }
-
-    @Override
-    public void handleTextMessage(WebSocketSession session, TextMessage message) throws IOException {
-        IntegrationPayload payload = objectMapper.readValue(message.getPayload(), IntegrationPayload.class);
-        if (payload.test) {
-            session.sendMessage(new TextMessage("%d".formatted(Instant.now().getEpochSecond())));
-            return;
-        }
-        if (payload.isVip()) {
-            if (!saltedCookie.ValidateSecretKey(payload.getSecret_key())) {
-                logger.warn("Cannot add vip without secret key");
-                return;
-            }
-
-            if (!payload.vip.isFilled()) {
-                logger.warn("Can try add VIP but, payload is not full filled!");
-                return;
-            }
-
-            if (!payload.vip.unique.isEmpty()) {
-                if (unique_set.contains(payload.vip.unique)) {
-                    logger.warn("VIP uniq payload value already exist! " + payload.vip.unique);
-                    return;
-                } else unique_set.add(payload.vip.unique);
-            }
-
-            Integer result = vipService.addVIP(
-                    SteamIDConverter.getSteamID(payload.vip.getSteam()),
-                    payload.vip.getAmount(),
-                    VipGiveMethod.valueOf(payload.vip.service.toUpperCase()),
-                    payload.vip.getExtra());
-            logger.info("Success add VIP to {}, result: {}", payload.vip.getSteam(), result);
-            return;
-        }
-    }
-
-    @Override
-    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
-        logger.info("Session {} close", session.getId());
-        super.afterConnectionClosed(session, status);
-    }
-
-    @Data
-    @NoArgsConstructor
-    public static class IntegrationPayload {
-        String secret_key = null;
-        ExternalVIP vip = null;
-        Boolean test = false;
-
-        public boolean isVip() {
-            return vip != null && secret_key != null;
-        }
-    }
-
-    @Data
-    @NoArgsConstructor
-    public static class ExternalVIP {
-        String steam;
-        Integer amount;
-        String service;
-        String extra = "";
-        String unique = "";
-
-        public boolean isFilled() {
-            return steam != null && amount != null && service != null;
-        }
-    }
-}
diff --git a/src/main/java/app/websocket/WebSocketConfig.java b/src/main/java/app/websocket/WebSocketConfig.java
deleted file mode 100644
index 27bc2d7..0000000
--- a/src/main/java/app/websocket/WebSocketConfig.java
+++ /dev/null
@@ -1,28 +0,0 @@
-package app.websocket;
-
-import org.springframework.context.annotation.Configuration;
-import org.springframework.web.socket.config.annotation.EnableWebSocket;
-import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
-import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
-
-@Configuration
-@EnableWebSocket
-public class WebSocketConfig implements WebSocketConfigurer {
-
-    private WebSocketHandler webSocketHandler;
-    private IntegrationWebSockerHandler integrationWebSockerHandler;
-
-    public WebSocketConfig(
-            WebSocketHandler webSocketHandler,
-            IntegrationWebSockerHandler integrationWebSockerHandler
-    ) {
-        this.webSocketHandler = webSocketHandler;
-        this.integrationWebSockerHandler = integrationWebSockerHandler;
-    }
-
-    @Override
-    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
-        registry.addHandler(webSocketHandler, "/ws");
-        registry.addHandler(integrationWebSockerHandler, "/ws_integrations");
-    }
-}
diff --git a/src/main/java/app/websocket/configuration/WebSocketConfig.java b/src/main/java/app/websocket/configuration/WebSocketConfig.java
new file mode 100644
index 0000000..6397190
--- /dev/null
+++ b/src/main/java/app/websocket/configuration/WebSocketConfig.java
@@ -0,0 +1,30 @@
+package app.websocket.configuration;
+
+import app.websocket.BaseWebsocketHandler;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.context.annotation.Configuration;
+import org.springframework.web.socket.WebSocketHandler;
+import org.springframework.web.socket.config.annotation.EnableWebSocket;
+import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
+import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
+
+import java.util.List;
+
+@Configuration
+@EnableWebSocket
+public class WebSocketConfig implements WebSocketConfigurer {
+
+    @Autowired
+    private List<BaseWebsocketHandler> handlers;
+    private final Logger logger = LoggerFactory.getLogger(getClass());
+
+    @Override
+    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
+        handlers.forEach(hndl -> {
+            logger.info("Found websocket bean: {}", hndl.getPath());
+            registry.addHandler((WebSocketHandler) hndl, "/ws/" + hndl.getPath());
+        });
+    }
+}
diff --git a/src/main/java/app/websocket/handlers/ExternalVIPHandler.java b/src/main/java/app/websocket/handlers/ExternalVIPHandler.java
new file mode 100644
index 0000000..e5b312b
--- /dev/null
+++ b/src/main/java/app/websocket/handlers/ExternalVIPHandler.java
@@ -0,0 +1,115 @@
+package app.websocket.handlers;
+
+import app.controllers.other.ExternalVIPController;
+import app.entities.VipGiveMethod;
+import app.services.db.VIPService;
+import app.utils.SaltedCookie;
+import app.utils.SteamIDConverter;
+import app.websocket.BaseWebsocketHandler;
+import com.fasterxml.jackson.core.JsonParseException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import lombok.Data;
+import lombok.NoArgsConstructor;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.annotation.Autowired;
+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.util.HashSet;
+
+@Component
+public class ExternalVIPHandler extends TextWebSocketHandler implements BaseWebsocketHandler {
+    private final Logger logger = LoggerFactory.getLogger(getClass());
+    private ObjectMapper objectMapper = new ObjectMapper();
+    private SaltedCookie saltedCookie;
+    private VIPService vipService;
+    private HashSet<String> unique_set;
+    private ExternalVIPController externalVIPController;
+
+    @Autowired
+    public ExternalVIPHandler(SaltedCookie saltedCookie, VIPService vipService, ExternalVIPController externalVIPController) {
+        super();
+        this.saltedCookie = saltedCookie;
+        this.vipService = vipService;
+        this.externalVIPController = externalVIPController;
+        this.unique_set = new HashSet();
+    }
+
+    @Override
+    public String getPath() {
+        return "integration/vip";
+    }
+
+    @Override
+    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())));
+    }
+
+    @Override
+    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;
+        }
+
+        ExternalVIP vip;
+        try {
+            vip = objectMapper.readValue(message.getPayload(), ExternalVIP.class);
+        } catch (JsonParseException jpe) {
+            session.sendMessage(new TextMessage("Invalid payload struct"));
+            return;
+        }
+
+        if (!vip.isFilled()) {
+            logger.warn("Can try add VIP but, payload is not full filled!");
+            return;
+        }
+
+        if (!vip.unique.isEmpty()) {
+            if (unique_set.contains(vip.unique)) {
+                logger.warn("VIP uniq payload value already exist! " + vip.unique);
+                return;
+            } else unique_set.add(vip.unique);
+        }
+
+        Integer result = vipService.addVIP(
+                SteamIDConverter.getSteamID(vip.getSteam()),
+                vip.getAmount(),
+                VipGiveMethod.valueOf(vip.service.toUpperCase()),
+                vip.getExtra());
+        logger.info("Success add VIP to {}, result: {}", vip.getSteam(), result);
+    }
+
+    @Override
+    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
+        logger.info("Session {} close", session.getId());
+        super.afterConnectionClosed(session, status);
+    }
+
+    @Data
+    @NoArgsConstructor
+    public static class ExternalVIP {
+        String steam;
+        Integer amount;
+        String service;
+        String extra = "";
+        String unique = "";
+
+        public boolean isFilled() {
+            return steam != null && amount != null && service != null;
+        }
+    }
+}
diff --git a/src/main/java/app/websocket/WebSocketHandler.java b/src/main/java/app/websocket/handlers/ServersHandler.java
similarity index 79%
rename from src/main/java/app/websocket/WebSocketHandler.java
rename to src/main/java/app/websocket/handlers/ServersHandler.java
index a572cc8..b41421c 100644
--- a/src/main/java/app/websocket/WebSocketHandler.java
+++ b/src/main/java/app/websocket/handlers/ServersHandler.java
@@ -1,12 +1,12 @@
-package app.websocket;
+package app.websocket.handlers;
 
+import app.entities.Stats;
 import app.entities.server.Server;
-import app.updates.BaseUpdater;
-import com.fasterxml.jackson.core.JsonProcessingException;
+import app.websocket.BaseWebsocketHandler;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.springframework.context.annotation.Scope;
+import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.stereotype.Component;
 import org.springframework.web.socket.CloseStatus;
 import org.springframework.web.socket.TextMessage;
@@ -18,17 +18,29 @@ import java.util.HashMap;
 import java.util.Map;
 
 @Component
-@Scope(value = "singleton")
-public class WebSocketHandler extends TextWebSocketHandler {
+public class ServersHandler extends TextWebSocketHandler implements BaseWebsocketHandler {
     private final Logger logger = LoggerFactory.getLogger(getClass());
     private final Map<String, WebSocketSession> activeSessions = new HashMap<>();
     private ObjectMapper objectMapper = new ObjectMapper();
+    private Stats stats;
+
+    @Autowired
+    ServersHandler(Stats stats) {
+        super();
+        this.stats = stats;
+    }
+
+    @Override
+    public String getPath() {
+        return "servers";
+    }
 
     @Override
     public void afterConnectionEstablished(WebSocketSession session) throws Exception {
         logger.info("Session {} open", session.getId());
         activeSessions.put(session.getId(), session);
         super.afterConnectionEstablished(session);
+        session.sendMessage(new TextMessage(objectMapper.writeValueAsString(stats.getServers())));
     }
 
     @Override
