fix: Correctly handle ws connections for new path pattern (#1701)

balhar-jakub · web-flow · commit 2db9a7ee0588 · 2021-08-13T13:40:24.000+02:00
* Correctly handle ws connections for new path pattern .
* Handle new and old path patterns in unit tests

Signed-off-by: Jakub Balhar &lt;jakub@balhar.net&gt;
diff --git a/gateway-service/src/main/java/org/zowe/apiml/gateway/ws/GatewayWebSocketConfigurer.java b/gateway-service/src/main/java/org/zowe/apiml/gateway/ws/GatewayWebSocketConfigurer.java
@@ -27,7 +27,7 @@ public GatewayWebSocketConfigurer(WebSocketProxyServerHandler webSocketProxyServ
 
     @Override
     public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
-        String webSocketPath = "/ws/**";    // NOSONAR
+        String webSocketPath = "/**/ws/**";    // NOSONAR
         log.debug("Registering WebSocket proxy handler to " + webSocketPath);
         registry.addHandler(webSocketProxyServerHandler, webSocketPath);
     }
diff --git a/gateway-service/src/main/java/org/zowe/apiml/gateway/ws/WebSocketProxyServerHandler.java b/gateway-service/src/main/java/org/zowe/apiml/gateway/ws/WebSocketProxyServerHandler.java
@@ -89,36 +89,50 @@ public Map<String, WebSocketRoutedSession> getRoutedSessions() {
     }
 
     @Override
-    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception {
+    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws IOException {
         String[] uriParts = getUriParts(webSocketSession);
-        if (uriParts != null && uriParts.length == 5) {
-            String majorVersion = uriParts[2];
-            String serviceId = uriParts[3];
-            String path = uriParts[4];
-
-            RoutedServices routedServices = routedServicesMap.get(serviceId);
-
-            if (routedServices != null) {
-                RoutedService service = routedServices.findServiceByGatewayUrl("ws/" + majorVersion);
-                if (service == null) {
-                    closeWebSocket(webSocketSession, CloseStatus.NOT_ACCEPTABLE,
-                        String.format("Requested ws/%s url is not known by the gateway", majorVersion));
-                    return;
-                }
-
-                ServiceInstance serviceInstance = findServiceInstance(serviceId);
-                if (serviceInstance != null) {
-                    openWebSocketConnection(service, serviceInstance, serviceInstance, path, webSocketSession);
-                } else {
-                    closeWebSocket(webSocketSession, CloseStatus.SERVICE_RESTARTED,
-                        String.format("Requested service %s does not have available instance", serviceId));
-                }
-            } else {
-                closeWebSocket(webSocketSession, CloseStatus.NOT_ACCEPTABLE,
-                    String.format("Requested service %s is not known by the gateway", serviceId));
-            }
-        } else {
+        if (uriParts == null || uriParts.length != 5) {
             closeWebSocket(webSocketSession, CloseStatus.NOT_ACCEPTABLE, "Invalid URL format");
+            return;
+        }
+
+        String majorVersion;
+        String serviceId;
+        String path = uriParts[4];
+
+        if (uriParts[1].equals("ws")) {
+            majorVersion = uriParts[2];
+            serviceId = uriParts[3];
+        } else {
+            majorVersion = uriParts[3];
+            serviceId = uriParts[1];
+        }
+
+        routeToService(webSocketSession, serviceId, majorVersion, path);
+    }
+
+    private void routeToService(WebSocketSession webSocketSession, String serviceId, String majorVersion, String path) throws IOException {
+        RoutedServices routedServices = routedServicesMap.get(serviceId);
+
+        if (routedServices == null) {
+            closeWebSocket(webSocketSession, CloseStatus.NOT_ACCEPTABLE,
+                String.format("Requested service %s is not known by the gateway", serviceId));
+            return;
+        }
+
+        RoutedService service = routedServices.findServiceByGatewayUrl("ws/" + majorVersion);
+        if (service == null) {
+            closeWebSocket(webSocketSession, CloseStatus.NOT_ACCEPTABLE,
+                String.format("Requested ws/%s url is not known by the gateway", majorVersion));
+            return;
+        }
+
+        ServiceInstance serviceInstance = findServiceInstance(serviceId);
+        if (serviceInstance != null) {
+            openWebSocketConnection(service, serviceInstance, serviceInstance, path, webSocketSession);
+        } else {
+            closeWebSocket(webSocketSession, CloseStatus.SERVICE_RESTARTED,
+                String.format("Requested service %s does not have available instance", serviceId));
         }
     }
 
diff --git a/gateway-service/src/main/resources/application.yml b/gateway-service/src/main/resources/application.yml
@@ -122,7 +122,7 @@ zuul:
     includeDebugHeader: false
     sensitiveHeaders: Expires,Date
     ignoredPatterns:
-        - /ws/**
+        - /**/ws/**
     host:
         connectTimeoutMillis: ${apiml.gateway.timeoutMillis}
         socketTimeoutMillis: ${apiml.gateway.timeoutMillis}
diff --git a/gateway-service/src/test/java/org/zowe/apiml/gateway/ws/WebSocketProxyServerHandlerTest.java b/gateway-service/src/test/java/org/zowe/apiml/gateway/ws/WebSocketProxyServerHandlerTest.java
@@ -11,7 +11,10 @@
 package org.zowe.apiml.gateway.ws;
 
 import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.springframework.cloud.client.ServiceInstance;
 import org.springframework.cloud.client.discovery.DiscoveryClient;
 import org.springframework.web.socket.CloseStatus;
@@ -52,36 +55,7 @@ public void setup() {
         );
     }
 
-    /**
-     * Happy Path
-     * <p>
-     * The Handler is properly created
-     * Specified Route is added to the list
-     * The connection is established
-     * The URI contains the valid service Id
-     * The service associated with given URI is retrieved
-     * Proper WebSocketSession is stored.
-     */
-    @Test
-    void givenValidRoute_whenTheConnectionIsEstablished_thenTheValidSessionIsStoredInternally() throws Exception {
-        RoutedServices routesForSpecificValidService = mock(RoutedServices.class);
-        when(routesForSpecificValidService.findServiceByGatewayUrl("ws/1"))
-            .thenReturn(new RoutedService("api-v1", "api/v1", "/api-v1/api/v1"));
-        ServiceInstance foundService = validServiceInstance();
-        when(discoveryClient.getInstances("api-v1")).thenReturn(Collections.singletonList(foundService));
-        underTest.addRoutedServices("api-v1", routesForSpecificValidService);
-        when(webSocketRoutedSessionFactory.session(any(), any(), any())).thenReturn(mock(WebSocketRoutedSession.class));
-
-        WebSocketSession establishedSession = mock(WebSocketSession.class);
-        String establishedSessionId = "validAndUniqueId";
-        when(establishedSession.getId()).thenReturn(establishedSessionId);
-        when(establishedSession.getUri()).thenReturn(new URI("wss://gatewayHost:1443/gateway/1/api-v1/api/v1"));
-        underTest.afterConnectionEstablished(establishedSession);
-
-        verify(webSocketRoutedSessionFactory).session(any(), any(), any());
-        WebSocketRoutedSession preparedSession = routedSessions.get(establishedSessionId);
-        assertThat(preparedSession, is(notNullValue()));
-    }
+
 
     private ServiceInstance validServiceInstance() {
         ServiceInstance validService = mock(ServiceInstance.class);
@@ -92,95 +66,159 @@ private ServiceInstance validServiceInstance() {
         return validService;
     }
 
-    /**
-     * Error Path
-     * <p>
-     * The Handler is properly created
-     * The connection is established
-     * The URI doesn't contain all needed parts
-     * The WebSocketSession is closed
-     */
-    @Test
-    void givenInvalidURI_whenTheConnectionIsEstablished_thenTheSocketIsClosedAsNotAcceptable() throws Exception {
-        WebSocketSession establishedSession = mock(WebSocketSession.class);
-        when(establishedSession.isOpen()).thenReturn(true);
-        when(establishedSession.getUri()).thenReturn(new URI("wss://gatewayHost:1443/invalidUrl"));
-
-        underTest.afterConnectionEstablished(establishedSession);
-
-        verify(establishedSession).close(new CloseStatus(CloseStatus.NOT_ACCEPTABLE.getCode(), "Invalid URL format"));
+    @Nested
+    class WhenTheConnectionIsEstablished {
+        WebSocketSession establishedSession;
+
+        @BeforeEach
+        void prepareSessionMock() {
+            establishedSession = mock(WebSocketSession.class);
+        }
+
+        @Nested
+        class ThenTheValidSessionIsStoredInternally {
+            @BeforeEach
+            void prepareRoutedService() {
+                String serviceId = "valid-service";
+
+                RoutedServices routesForSpecificValidService = mock(RoutedServices.class);
+                when(routesForSpecificValidService.findServiceByGatewayUrl("ws/v1"))
+                    .thenReturn(new RoutedService("ws-v1", "ws/v1", "/valid-service/ws/v1"));
+                ServiceInstance foundService = validServiceInstance();
+                when(discoveryClient.getInstances(serviceId)).thenReturn(Collections.singletonList(foundService));
+
+                underTest.addRoutedServices(serviceId, routesForSpecificValidService);
+            }
+
+            /**
+             * Happy Path
+             * <p>
+             * The Handler is properly created
+             * Specified Route is added to the list
+             * The connection is established
+             * The URI contains the valid service Id
+             * The service associated with given URI is retrieved
+             * Proper WebSocketSession is stored.
+             */
+            @ParameterizedTest(name = "WhenTheConnectionIsEstablished.ThenTheValidSessionIsStoredInternally#givenValidRoute {0}")
+            @ValueSource(strings = {"wss://gatewayHost:1443/valid-service/ws/v1/valid-path", "wss://gatewayHost:1443/ws/v1/valid-service/valid-path"})
+            void givenValidRoute(String path) throws Exception {
+                when(webSocketRoutedSessionFactory.session(any(), any(), any())).thenReturn(mock(WebSocketRoutedSession.class));
+
+                String establishedSessionId = "validAndUniqueId";
+                when(establishedSession.getId()).thenReturn(establishedSessionId);
+                when(establishedSession.getUri()).thenReturn(new URI(path));
+
+                underTest.afterConnectionEstablished(establishedSession);
+
+                verify(webSocketRoutedSessionFactory).session(any(), any(), any());
+                WebSocketRoutedSession preparedSession = routedSessions.get(establishedSessionId);
+                assertThat(preparedSession, is(notNullValue()));
+            }
+        }
+
+
+        @Nested
+        class ThenTheSocketIsClosed {
+            @BeforeEach
+            void sessionIsOpen() {
+                when(establishedSession.isOpen()).thenReturn(true);
+            }
+
+            /**
+             * Error Path
+             * <p>
+             * The Handler is properly created
+             * The connection is established
+             * The URI doesn't contain all needed parts
+             * The WebSocketSession is closed
+             */
+            @Test
+            void givenInvalidURI() throws Exception {
+                when(establishedSession.getUri()).thenReturn(new URI("wss://gatewayHost:1443/invalidUrl"));
+
+                underTest.afterConnectionEstablished(establishedSession);
+
+                verify(establishedSession).close(new CloseStatus(CloseStatus.NOT_ACCEPTABLE.getCode(), "Invalid URL format"));
+            }
+
+            /**
+             * Error Path
+             * <p>
+             * The Handler is properly created
+             * The connection is established
+             * The URI contains the service Id for which there is no service
+             * The WebSocketSession is closed
+             */
+            @Test
+            void givenInvalidRoute() throws Exception {
+                when(establishedSession.getUri()).thenReturn(new URI("wss://gatewayHost:1443/ws/v1/non_existent_service/valid-path"));
+
+                underTest.afterConnectionEstablished(establishedSession);
+
+                verify(establishedSession).close(new CloseStatus(CloseStatus.NOT_ACCEPTABLE.getCode(), "Requested service non_existent_service is not known by the gateway"));
+            }
+
+            /**
+             * Error Path
+             * <p>
+             * The Handler is properly created
+             * Specified Route is added to the list
+             * The connection is established
+             * The URI contains the valid service Id
+             * The service associated with given URI is retrieved
+             * The service isn't available in the Discovery Service
+             * Proper WebSocketSession is stored.
+             */
+            @Test
+            void givenNoInstanceOfTheServiceIsInTheRepository() throws Exception {
+                when(establishedSession.getUri()).thenReturn(new URI("wss://gatewayHost:1443/service-without-instance/ws/v1/valid-path"));
+
+                RoutedServices routesForSpecificValidService = mock(RoutedServices.class);
+                when(routesForSpecificValidService.findServiceByGatewayUrl("ws/v1"))
+                    .thenReturn(new RoutedService("api-v1", "api/v1", "/api-v1/api/v1"));
+                underTest.addRoutedServices("service-without-instance", routesForSpecificValidService);
+
+                underTest.afterConnectionEstablished(establishedSession);
+
+                verify(establishedSession).close(new CloseStatus(CloseStatus.SERVICE_RESTARTED.getCode(), "Requested service service-without-instance does not have available instance"));
+            }
+        }
     }
 
-    /**
-     * Error Path
-     * <p>
-     * The Handler is properly created
-     * The connection is established
-     * The URI contains the service Id for which there is no service
-     * The WebSocketSession is closed
-     */
-    @Test
-    void givenInvalidRoute_whenTheConnectionIsEstablished_thenTheSocketIsClosedAsNotAcceptable() throws Exception {
-        WebSocketSession establishedSession = mock(WebSocketSession.class);
-        when(establishedSession.isOpen()).thenReturn(true);
-        when(establishedSession.getUri()).thenReturn(new URI("wss://gatewayHost:1443/api/v1/non_existent_service/api/v1"));
-
-        underTest.afterConnectionEstablished(establishedSession);
-
-        verify(establishedSession).close(new CloseStatus(CloseStatus.NOT_ACCEPTABLE.getCode(), "Requested service non_existent_service is not known by the gateway"));
-    }
+    @Nested
+    class GivenValidExistingSession {
+        WebSocketSession establishedSession;
+        WebSocketRoutedSession internallyStoredSession;
 
-    /**
-     * Error Path
-     * <p>
-     * The Handler is properly created
-     * Specified Route is added to the list
-     * The connection is established
-     * The URI contains the valid service Id
-     * The service associated with given URI is retrieved
-     * The service isn't available in the Discovery Service
-     * Proper WebSocketSession is stored.
-     */
-    @Test
-    void givenNoInstanceOfTheServiceIsInTheRepository_whenTheConnectionIsEstablished_thenTheSocketIsClosedAsServiceRestarted() throws Exception {
-        WebSocketSession establishedSession = mock(WebSocketSession.class);
-        when(establishedSession.isOpen()).thenReturn(true);
-        when(establishedSession.getUri()).thenReturn(new URI("wss://gatewayHost:1443/api/v1/api-v1/api/v1"));
-        RoutedServices routesForSpecificValidService = mock(RoutedServices.class);
-        when(routesForSpecificValidService.findServiceByGatewayUrl("ws/v1"))
-            .thenReturn(new RoutedService("api-v1", "api/v1", "/api-v1/api/v1"));
-        underTest.addRoutedServices("api-v1", routesForSpecificValidService);
-
-        underTest.afterConnectionEstablished(establishedSession);
-
-        verify(establishedSession).close(new CloseStatus(CloseStatus.SERVICE_RESTARTED.getCode(), "Requested service api-v1 does not have available instance"));
-    }
+        @BeforeEach
+        void prepareSessionMock() {
+            establishedSession = mock(WebSocketSession.class);
+            String validSessionId = "123";
+            when(establishedSession.getId()).thenReturn(validSessionId);
 
-    @Test
-    void givenValidSession_whenTheConnectionIsClosed_thenTheSessionIsClosedAndRemovedFromRepository() throws Exception {
-        CloseStatus normalClose = CloseStatus.NORMAL;
-        WebSocketSession establishedSession = mock(WebSocketSession.class);
-        String validSessionId = "123";
-        when(establishedSession.getId()).thenReturn(validSessionId);
-        routedSessions.put(validSessionId, mock(WebSocketRoutedSession.class));
+            internallyStoredSession = mock(WebSocketRoutedSession.class);
+            routedSessions.put(validSessionId, internallyStoredSession);
+        }
 
-        underTest.afterConnectionClosed(establishedSession, normalClose);
+        @Test
+        void whenTheConnectionIsClosed_thenTheSessionIsClosedAndRemovedFromRepository() throws Exception {
+            CloseStatus normalClose = CloseStatus.NORMAL;
 
-        verify(establishedSession).close(normalClose);
-        assertThat(routedSessions.entrySet(), hasSize(0));
-    }
+            underTest.afterConnectionClosed(establishedSession, normalClose);
+
+            verify(establishedSession).close(normalClose);
+            assertThat(routedSessions.entrySet(), hasSize(0));
+        }
+
+        @Test
+        void whenTheMessageIsReceived_thenTheMessageIsPassedToTheSession() throws Exception {
+            WebSocketMessage<String> passedMessage = mock(WebSocketMessage.class);
 
-    @Test
-    void givenValidSession_whenTheMessageIsReceived_thenTheMessageIsPassedToTheSession() throws Exception {
-        WebSocketSession establishedSession = mock(WebSocketSession.class);
-        String validSessionId = "123";
-        when(establishedSession.getId()).thenReturn(validSessionId);
-        WebSocketRoutedSession internallyStoredSession = mock(WebSocketRoutedSession.class);
-        routedSessions.put(validSessionId, internallyStoredSession);
-        WebSocketMessage<String> passedMessage = mock(WebSocketMessage.class);
+            underTest.handleMessage(establishedSession, passedMessage);
 
-        underTest.handleMessage(establishedSession, passedMessage);
+            verify(internallyStoredSession).sendMessageToServer(passedMessage);
+        }
 
-        verify(internallyStoredSession).sendMessageToServer(passedMessage);
     }
 }
diff --git a/integration-tests/src/test/java/org/zowe/apiml/integration/proxy/WebSocketProxyTest.java b/integration-tests/src/test/java/org/zowe/apiml/integration/proxy/WebSocketProxyTest.java

Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,7 @@ public GatewayWebSocketConfigurer(WebSocketProxyServerHandler webSocketProxyServ`
`27`	`27`
`28`	`28`	`@Override`
`29`	`29`	`public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {`
`30`		`- String webSocketPath = "/ws/**"; // NOSONAR`
	`30`	`+ String webSocketPath = "//ws/"; // NOSONAR`
`31`	`31`	`log.debug("Registering WebSocket proxy handler to " + webSocketPath);`
`32`	`32`	`registry.addHandler(webSocketProxyServerHandler, webSocketPath);`
`33`	`33`	`}`