websocket sync handshakes

shenyu-admin/src/main/java/org/apache/shenyu/admin/config/properties/WebsocketSyncProperties.java

@@ -40,6 +40,11 @@ public class WebsocketSyncProperties {
      */
     private String allowOrigins;
 
+    /**
+     * WebSocket sync token.
+     */
+    private String token;
+
     /**
      * Gets the value of enabled.
      *
@@ -91,4 +96,22 @@ public String getAllowOrigins() {
     public void setAllowOrigins(final String allowOrigins) {
         this.allowOrigins = allowOrigins;
     }
+
+    /**
+     * get token.
+     *
+     * @return token
+     */
+    public String getToken() {
+        return token;
+    }
+
+    /**
+     * set token.
+     *
+     * @param token token
+     */
+    public void setToken(final String token) {
+        this.token = token;
+    }
 }

shenyu-admin/src/main/java/org/apache/shenyu/admin/listener/websocket/WebsocketConfigurator.java

@@ -26,13 +26,21 @@
 import org.apache.shenyu.admin.config.properties.WebsocketSyncProperties;
 import org.apache.shenyu.admin.spring.SpringBeanUtils;
 import org.apache.shenyu.common.constant.Constants;
+import org.apache.shenyu.common.exception.ShenyuException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.web.servlet.ServletContextInitializer;
 import org.springframework.context.annotation.Configuration;
 
+import java.nio.charset.StandardCharsets;
+import java.security.MessageDigest;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
 import static org.apache.tomcat.websocket.server.Constants.BINARY_BUFFER_SIZE_SERVLET_CONTEXT_INIT_PARAM;
 import static org.apache.tomcat.websocket.server.Constants.TEXT_BUFFER_SIZE_SERVLET_CONTEXT_INIT_PARAM;
 
@@ -52,6 +60,7 @@ public class WebsocketConfigurator extends ServerEndpointConfig.Configurator imp
 
     @Override
     public void modifyHandshake(final ServerEndpointConfig sec, final HandshakeRequest request, final HandshakeResponse response) {
+        checkSyncToken(request);
         HttpSession httpSession = (HttpSession) request.getHttpSession();
         sec.getUserProperties().put(WebsocketListener.CLIENT_IP_NAME, httpSession.getAttribute(WebsocketListener.CLIENT_IP_NAME));
         sec.getUserProperties().put(Constants.CLIENT_PORT_NAME, httpSession.getAttribute(Constants.CLIENT_PORT_NAME));
@@ -85,4 +94,34 @@ public void onStartup(final ServletContext servletContext) {
                     String.valueOf(messageMaxSize));
         }
     }
+
+    private void checkSyncToken(final HandshakeRequest request) {
+        String configuredToken = websocketSyncProperties.getToken();
+        if (StringUtils.isBlank(configuredToken)) {
+            throw new ShenyuException("websocket sync token is not configured");
+        }
+        String requestToken = getHeader(request.getHeaders(), Constants.SHENYU_WEBSOCKET_SYNC_TOKEN);
+        if (StringUtils.isBlank(requestToken) || !isSameToken(configuredToken, requestToken)) {
+            throw new ShenyuException("websocket sync token is invalid");
+        }
+    }
+
+    private boolean isSameToken(final String configuredToken, final String requestToken) {
+        return MessageDigest.isEqual(
+                configuredToken.getBytes(StandardCharsets.UTF_8),
+                requestToken.getBytes(StandardCharsets.UTF_8));
+    }
+
+    private String getHeader(final Map<String, List<String>> headers, final String name) {
+        return Optional.ofNullable(headers)
+                .orElse(Collections.emptyMap())
+                .entrySet()
+                .stream()
+                .filter(entry -> StringUtils.equalsIgnoreCase(entry.getKey(), name))
+                .map(Map.Entry::getValue)
+                .filter(values -> !values.isEmpty())
+                .map(values -> values.get(0))
+                .findFirst()
+                .orElse(null);
+    }
 }

shenyu-admin/src/main/resources/application.yml

@@ -70,6 +70,7 @@ shenyu:
     websocket:
       enabled: true
       messageMaxSize: 10240
+      token: ${SHENYU_SYNC_WEBSOCKET_TOKEN:}
       allowOrigins: ws://localhost:9095;ws://localhost:9195;
 #    apollo:
 #      meta: http://localhost:8080

shenyu-admin/src/test/java/org/apache/shenyu/admin/config/properties/WebsocketSyncPropertiesTest.java

@@ -43,9 +43,11 @@ public void testWebsocketSyncPropertiesSetValue() {
         WebsocketSyncProperties websocketSyncProperties = getContext().getBean(WebsocketSyncProperties.class);
         websocketSyncProperties.setMessageMaxSize(0);
         websocketSyncProperties.setAllowOrigins("allowOrigins");
+        websocketSyncProperties.setToken("token");
         assertThat(websocketSyncProperties.isEnabled(), comparesEqualTo(false));
         Assertions.assertEquals(websocketSyncProperties.getMessageMaxSize(), 0);
         Assertions.assertEquals(websocketSyncProperties.getAllowOrigins(), "allowOrigins");
+        Assertions.assertEquals(websocketSyncProperties.getToken(), "token");
     }
 
     @Configuration

shenyu-admin/src/test/java/org/apache/shenyu/admin/listener/websocket/WebsocketConfiguratorTest.java

@@ -25,18 +25,22 @@
 import org.apache.shenyu.admin.config.properties.WebsocketSyncProperties;
 import org.apache.shenyu.admin.spring.SpringBeanUtils;
 import org.apache.shenyu.common.constant.Constants;
+import org.apache.shenyu.common.exception.ShenyuException;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.springframework.context.ApplicationContext;
 import org.springframework.test.util.ReflectionTestUtils;
 
+import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 import static org.apache.tomcat.websocket.server.Constants.BINARY_BUFFER_SIZE_SERVLET_CONTEXT_INIT_PARAM;
 import static org.apache.tomcat.websocket.server.Constants.TEXT_BUFFER_SIZE_SERVLET_CONTEXT_INIT_PARAM;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -64,11 +68,13 @@ void setUp() {
 
     @Test
     void testModifyHandshake() {
+        websocketSyncProperties.setToken("websocket-sync-token");
         ServerEndpointConfig sec = mock(ServerEndpointConfig.class);
         Map<String, Object> userProperties = new HashMap<>();
         when(sec.getUserProperties()).thenReturn(userProperties);
 
         HandshakeRequest request = mock(HandshakeRequest.class);
+        when(request.getHeaders()).thenReturn(Collections.singletonMap(Constants.SHENYU_WEBSOCKET_SYNC_TOKEN, List.of("websocket-sync-token")));
         HttpSession httpSession = mock(HttpSession.class);
         when(request.getHttpSession()).thenReturn(httpSession);
         when(httpSession.getAttribute(WebsocketListener.CLIENT_IP_NAME)).thenReturn("192.168.1.1");
@@ -85,11 +91,13 @@ void testModifyHandshake() {
 
     @Test
     void testModifyHandshakePutsAllAttributes() {
+        websocketSyncProperties.setToken("websocket-sync-token");
         ServerEndpointConfig sec = mock(ServerEndpointConfig.class);
         Map<String, Object> userProperties = new HashMap<>();
         when(sec.getUserProperties()).thenReturn(userProperties);
 
         HandshakeRequest request = mock(HandshakeRequest.class);
+        when(request.getHeaders()).thenReturn(Collections.singletonMap(Constants.SHENYU_WEBSOCKET_SYNC_TOKEN, List.of("websocket-sync-token")));
         HttpSession httpSession = mock(HttpSession.class);
         when(request.getHttpSession()).thenReturn(httpSession);
 
@@ -104,6 +112,39 @@ void testModifyHandshakePutsAllAttributes() {
         assertTrue(userProperties.containsKey(Constants.SHENYU_NAMESPACE_ID));
     }
 
+    @Test
+    void testModifyHandshakeRejectsMissingToken() {
+        websocketSyncProperties.setToken("websocket-sync-token");
+        ServerEndpointConfig sec = mock(ServerEndpointConfig.class);
+        HandshakeRequest request = mock(HandshakeRequest.class);
+        when(request.getHeaders()).thenReturn(Collections.emptyMap());
+        HandshakeResponse response = mock(HandshakeResponse.class);
+
+        assertThrows(ShenyuException.class, () -> websocketConfigurator.modifyHandshake(sec, request, response));
+    }
+
+    @Test
+    void testModifyHandshakeRejectsInvalidToken() {
+        websocketSyncProperties.setToken("websocket-sync-token");
+        ServerEndpointConfig sec = mock(ServerEndpointConfig.class);
+        HandshakeRequest request = mock(HandshakeRequest.class);
+        when(request.getHeaders()).thenReturn(Collections.singletonMap(Constants.SHENYU_WEBSOCKET_SYNC_TOKEN, List.of("invalid-token")));
+        HandshakeResponse response = mock(HandshakeResponse.class);
+
+        assertThrows(ShenyuException.class, () -> websocketConfigurator.modifyHandshake(sec, request, response));
+    }
+
+    @Test
+    void testModifyHandshakeRejectsBlankConfiguredToken() {
+        websocketSyncProperties.setToken("");
+        ServerEndpointConfig sec = mock(ServerEndpointConfig.class);
+        HandshakeRequest request = mock(HandshakeRequest.class);
+        when(request.getHeaders()).thenReturn(Collections.singletonMap(Constants.SHENYU_WEBSOCKET_SYNC_TOKEN, List.of("websocket-sync-token")));
+        HandshakeResponse response = mock(HandshakeResponse.class);
+
+        assertThrows(ShenyuException.class, () -> websocketConfigurator.modifyHandshake(sec, request, response));
+    }
+
     @Test
     void testCheckOriginAllowedWhenAllowOriginsEmpty() {
         websocketSyncProperties.setAllowOrigins("");

shenyu-bootstrap/src/main/resources/application.yml

@@ -223,6 +223,7 @@ shenyu:
   sync:
     websocket:
       urls: ws://localhost:9095/websocket
+      token: ${SHENYU_SYNC_WEBSOCKET_TOKEN:}
       allowOrigin: ws://localhost:9195
 #    apollo:
 #      appId: shenyu

shenyu-common/src/main/java/org/apache/shenyu/common/constant/Constants.java

@@ -713,6 +713,11 @@ public interface Constants {
      */
     String X_ACCESS_TOKEN = "X-Access-Token";
 
+    /**
+     * X-Shenyu-Sync-Token.
+     */
+    String SHENYU_WEBSOCKET_SYNC_TOKEN = "X-Shenyu-Sync-Token";
+
     /**
      * X-API-KEY; AI proxy key header.
      */

shenyu-sync-data-center/shenyu-sync-data-websocket/src/main/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataService.java

@@ -17,12 +17,12 @@
 
 package org.apache.shenyu.plugin.sync.data.websocket;
 
-import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Lists;
 
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.shenyu.common.config.ShenyuConfig;
+import org.apache.shenyu.common.constant.Constants;
 import org.apache.shenyu.common.enums.RunningModeEnum;
 import org.apache.shenyu.common.timer.AbstractRoundTask;
 import org.apache.shenyu.common.timer.Timer;
@@ -41,6 +41,7 @@
 import org.springframework.boot.autoconfigure.web.ServerProperties;
 
 import java.net.URI;
+import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -119,34 +120,7 @@ public WebsocketSyncDataService(
         LOG.info("start init connecting...");
         List<String> urls = websocketConfig.getUrls();
         for (String url : urls) {
-            if (StringUtils.isNotEmpty(websocketConfig.getAllowOrigin())) {
-                Map<String, String> headers =
-                        ImmutableMap.of(ORIGIN_HEADER_NAME, websocketConfig.getAllowOrigin());
-                clients.add(
-                        new ShenyuWebsocketClient(
-                                URI.create(url),
-                        headers,
-                        Objects.requireNonNull(pluginDataSubscriber),
-                        metaDataSubscribers,
-                        authDataSubscribers,
-                        proxySelectorDataSubscribers,
-                        discoveryUpstreamDataSubscribers,
-                                this.aiProxyApiKeyDataSubscribers,
-                        namespaceId,
-                        serverProperties.getPort()));
-            } else {
-                clients.add(
-                        new ShenyuWebsocketClient(
-                                URI.create(url),
-                        Objects.requireNonNull(pluginDataSubscriber),
-                        metaDataSubscribers,
-                        authDataSubscribers,
-                        proxySelectorDataSubscribers,
-                        discoveryUpstreamDataSubscribers,
-                                this.aiProxyApiKeyDataSubscribers,
-                        namespaceId,
-                        serverProperties.getPort()));
-            }
+            clients.add(createClient(url));
         }
         LOG.info("start check task...");
         this.timer.add(timerTask = new AbstractRoundTask(null, TimeUnit.SECONDS.toMillis(60)) {
@@ -164,32 +138,7 @@ private void masterCheck() {
         if (CollectionUtils.isEmpty(clients)) {
             List<String> urls = websocketConfig.getUrls();
             for (String url : urls) {
-                if (StringUtils.isNotEmpty(websocketConfig.getAllowOrigin())) {
-                    Map<String, String> headers =
-                            ImmutableMap.of(ORIGIN_HEADER_NAME, websocketConfig.getAllowOrigin());
-                    clients.add(
-                            new ShenyuWebsocketClient(
-                                    URI.create(url),
-                                    headers,
-                                    Objects.requireNonNull(pluginDataSubscriber),
-                                    metaDataSubscribers,
-                                    authDataSubscribers,
-                                    proxySelectorDataSubscribers,
-                                    discoveryUpstreamDataSubscribers,
-                                    this.aiProxyApiKeyDataSubscribers,
-                                    namespaceId, serverProperties.getPort()));
-                } else {
-                    clients.add(
-                            new ShenyuWebsocketClient(
-                                    URI.create(url),
-                                    Objects.requireNonNull(pluginDataSubscriber),
-                                    metaDataSubscribers,
-                                    authDataSubscribers,
-                                    proxySelectorDataSubscribers,
-                                    discoveryUpstreamDataSubscribers,
-                                    this.aiProxyApiKeyDataSubscribers,
-                                    namespaceId, serverProperties.getPort()));
-                }
+                clients.add(createClient(url));
             }
         }
         Iterator<ShenyuWebsocketClient> iterator = clients.iterator();
@@ -228,6 +177,27 @@ public void close() {
         }
         timer.shutdown();
     }
+
+    private ShenyuWebsocketClient createClient(final String url) {
+        Map<String, String> headers = new HashMap<>();
+        if (StringUtils.isNotEmpty(websocketConfig.getAllowOrigin())) {
+            headers.put(ORIGIN_HEADER_NAME, websocketConfig.getAllowOrigin());
+        }
+        if (StringUtils.isNotBlank(websocketConfig.getToken())) {
+            headers.put(Constants.SHENYU_WEBSOCKET_SYNC_TOKEN, websocketConfig.getToken());
+        }
+        return new ShenyuWebsocketClient(
+                URI.create(url),
+                headers,
+                Objects.requireNonNull(pluginDataSubscriber),
+                metaDataSubscribers,
+                authDataSubscribers,
+                proxySelectorDataSubscribers,
+                discoveryUpstreamDataSubscribers,
+                this.aiProxyApiKeyDataSubscribers,
+                namespaceId,
+                serverProperties.getPort());
+    }
 
     /**
      * get websocket config.