feat: websocket timeout and close server on error (#2914)

Signed-off-by: Pablo Hernán Carle <[email protected]> Co-authored-by: Pavel Jareš <[email protected]> Signed-off-by: Pablo Carle <[email protected]> --------- Signed-off-by: Pablo Hernán Carle <[email protected]> Signed-off-by: Pablo Carle <[email protected]> Co-authored-by: Pablo Hernán Carle <[email protected]> Co-authored-by: Pavel Jareš <[email protected]> Co-authored-by: Petr Weinfurt <[email protected]>
zowe · May 17, 2023 · 020da87 · 020da87
1 parent 3809622
commit 020da87
Show file tree

Hide file tree

Showing 9 changed files with 179 additions and 38 deletions.
diff --git a/gateway-package/src/main/resources/bin/start.sh b/gateway-package/src/main/resources/bin/start.sh
@@ -218,6 +218,7 @@ _BPX_JOBNAME=${ZWE_zowe_job_prefix}${GATEWAY_CODE} java \
     -Dserver.address=0.0.0.0 \
     -Dserver.maxConnectionsPerRoute=${ZWE_configs_server_maxConnectionsPerRoute:-100} \
     -Dserver.maxTotalConnections=${ZWE_configs_server_maxTotalConnections:-1000} \
+    -Dserver.webSocket.maxIdleTimeout=${ZWE_configs_server_webSocket_maxIdleTimeout:-3600000} \
     -Dserver.ssl.enabled=${ZWE_configs_server_ssl_enabled:-true} \
     -Dserver.ssl.protocol=${ZWE_configs_server_ssl_protocol:-"TLSv1.2"}  \
     -Dserver.ssl.keyStore="${keystore_location}" \

diff --git a/gateway-service/src/main/java/org/zowe/apiml/gateway/ws/WebSocketClientFactory.java b/gateway-service/src/main/java/org/zowe/apiml/gateway/ws/WebSocketClientFactory.java
@@ -10,17 +10,19 @@
 
 package org.zowe.apiml.gateway.ws;
 
-import lombok.AccessLevel;
-import lombok.RequiredArgsConstructor;
-import lombok.extern.slf4j.Slf4j;
+import javax.annotation.PreDestroy;
+
 import org.eclipse.jetty.client.HttpClient;
 import org.eclipse.jetty.util.ssl.SslContextFactory;
 import org.eclipse.jetty.websocket.client.WebSocketClient;
 import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Component;
 import org.springframework.web.socket.client.jetty.JettyWebSocketClient;
 
-import javax.annotation.PreDestroy;
+import lombok.AccessLevel;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
 
 /**
  * Factory for provisioning web socket client
@@ -35,10 +37,15 @@ public class WebSocketClientFactory {
     private final JettyWebSocketClient client;
 
     @Autowired
-    public WebSocketClientFactory(SslContextFactory.Client jettyClientSslContextFactory) {
+    public WebSocketClientFactory(
+        SslContextFactory.Client jettyClientSslContextFactory,
+        @Value("${server.webSocket.maxIdleTimeout:3600000}") int maxIdleWebSocketTimeout
+        ) {
         log.debug("Creating Jetty WebSocket client, with SslFactory: {}",
             jettyClientSslContextFactory);
-        client = new JettyWebSocketClient(new WebSocketClient(new HttpClient(jettyClientSslContextFactory)));
+        WebSocketClient wsClient = new WebSocketClient(new HttpClient(jettyClientSslContextFactory));
+        wsClient.setMaxIdleTimeout(maxIdleWebSocketTimeout);
+        client = new JettyWebSocketClient(wsClient);
         client.start();
     }
 
@@ -52,6 +59,7 @@ void closeClient() {
             log.debug("Closing Jetty WebSocket client");
             client.stop();
         }
+
     }
 
 }
diff --git a/gateway-service/src/main/java/org/zowe/apiml/gateway/ws/WebSocketProxyClientHandler.java b/gateway-service/src/main/java/org/zowe/apiml/gateway/ws/WebSocketProxyClientHandler.java
@@ -11,6 +11,10 @@
 package org.zowe.apiml.gateway.ws;
 
 import lombok.extern.slf4j.Slf4j;
+
+import java.util.concurrent.TimeoutException;
+
+import org.eclipse.jetty.websocket.api.CloseException;
 import org.springframework.web.socket.CloseStatus;
 import org.springframework.web.socket.WebSocketMessage;
 import org.springframework.web.socket.WebSocketSession;
@@ -42,5 +46,12 @@ public void afterConnectionClosed(WebSocketSession session, CloseStatus status)
     @Override
     public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
         log.warn("WebSocket transport error in session {}: {}", session.getId(), exception.getMessage());
+        if (exception instanceof CloseException && exception.getCause() instanceof TimeoutException) {
+            // Idle timeout
+            webSocketServerSession.close(CloseStatus.NORMAL);
+        } else if (exception instanceof CloseException) {
+            webSocketServerSession.close(new CloseStatus(((CloseException) exception).getStatusCode(), exception.getMessage()));
+        }
+        super.handleTransportError(session, exception);
     }
 }
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
@@ -34,6 +34,7 @@
 import java.net.URI;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
@@ -185,6 +186,16 @@ private void openWebSocketConnection(RoutedService service, ServiceInstance serv
 
     @Override
     public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
+        // if the browser closes the session, close the GWs client one as well.
+        Optional.ofNullable(routedSessions.get(session.getId()))
+            .map(WebSocketRoutedSession::getWebSocketClientSession)
+            .ifPresent(clientSession -> {
+                try {
+                    clientSession.close(status);
+                } catch (IOException e) {
+                    log.debug("Error closing WebSocket client connection {}: {}", clientSession.getId(), e.getMessage());
+                }
+            });
         routedSessions.remove(session.getId());
     }
 

diff --git a/gateway-service/src/test/java/org/zowe/apiml/gateway/ws/WebSocketClientFactoryTest.java b/gateway-service/src/test/java/org/zowe/apiml/gateway/ws/WebSocketClientFactoryTest.java
@@ -13,11 +13,16 @@
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
+import org.springframework.test.util.ReflectionTestUtils;
 import org.springframework.web.socket.client.jetty.JettyWebSocketClient;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.mockito.Mockito.*;
 
+import org.eclipse.jetty.util.ssl.SslContextFactory;
+import org.eclipse.jetty.websocket.client.WebSocketClient;
+
 class WebSocketClientFactoryTest {
 
     @Nested
@@ -52,4 +57,23 @@ void whenGetClient_thenReturnInstance() {
 
     }
 
-}
+    @Nested
+    class CreatedInstanceWithConfig {
+
+        private WebSocketClientFactory webSocketClientFactory;
+
+        @BeforeEach
+        void setUp() {
+            SslContextFactory.Client sslClient = mock(SslContextFactory.Client.class);
+            this.webSocketClientFactory = new WebSocketClientFactory(sslClient, 1234);
+        }
+
+        @Test
+        void givenInitilizedClient_thenHasNonDefaultIdleConfig() {
+            WebSocketClient wsClient = (WebSocketClient) ReflectionTestUtils.getField(webSocketClientFactory.getClientInstance(), "client");
+            assertEquals(1234, wsClient.getMaxIdleTimeout());
+        }
+
+    }
+
+}
diff --git a/gateway-service/src/test/java/org/zowe/apiml/gateway/ws/WebSocketProxyClientHandlerTest.java b/gateway-service/src/test/java/org/zowe/apiml/gateway/ws/WebSocketProxyClientHandlerTest.java
@@ -0,0 +1,75 @@
+/*
+ * This program and the accompanying materials are made available under the terms of the
+ * Eclipse Public License v2.0 which accompanies this distribution, and is available at
+ * https://www.eclipse.org/legal/epl-v20.html
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *
+ * Copyright Contributors to the Zowe Project.
+ */
+
+package org.zowe.apiml.gateway.ws;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import java.util.concurrent.TimeoutException;
+
+import org.eclipse.jetty.websocket.api.CloseException;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.web.socket.CloseStatus;
+import org.springframework.web.socket.WebSocketSession;
+
+@ExtendWith(MockitoExtension.class)
+public class WebSocketProxyClientHandlerTest {
+
+    @Mock
+    private WebSocketSession serverSession;
+
+    private WebSocketProxyClientHandler webSocketProxyClientHandler;
+
+    @BeforeEach
+    void setUp() {
+        webSocketProxyClientHandler = new WebSocketProxyClientHandler(serverSession);
+    }
+
+    @Nested
+    class GivenHandler {
+
+        @Nested
+        class AndConnectionIsClosed {
+
+            @Test
+            void thenCloseServer() throws Exception {
+                webSocketProxyClientHandler.afterConnectionClosed(mock(WebSocketSession.class), CloseStatus.NORMAL);
+                verify(serverSession, times(1)).close(CloseStatus.NORMAL);
+            }
+
+        }
+
+        @Nested
+        class AndConnectionTransportError {
+
+            @Test
+            void andTimeout_thenCloseNormal() throws Exception {
+                webSocketProxyClientHandler.handleTransportError(mock(WebSocketSession.class), new CloseException(0, new TimeoutException("null")));
+                verify(serverSession, times(1)).close(CloseStatus.NORMAL);
+            }
+
+            @Test
+            void andCloseException_thenForwardError() throws Exception {
+                webSocketProxyClientHandler.handleTransportError(mock(WebSocketSession.class), new CloseException(CloseStatus.PROTOCOL_ERROR.getCode(), new Exception("message")));
+                verify(serverSession, times(1)).close(new CloseStatus(1002, "java.lang.Exception: message"));
+            }
+
+        }
+
+    }
+
+}
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
@@ -33,7 +33,7 @@
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.notNullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.*;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.*;
@@ -59,16 +59,6 @@ public void setup() {
         ReflectionTestUtils.setField(underTest, "meAsProxy", underTest);
     }
 
-
-    private ServiceInstance validServiceInstance() {
-        ServiceInstance validService = mock(ServiceInstance.class);
-        when(validService.getHost()).thenReturn("gatewayHost");
-        when(validService.isSecure()).thenReturn(true);
-        when(validService.getPort()).thenReturn(1443);
-
-        return validService;
-    }
-
     @Nested
     class WhenTheConnectionIsEstablished {
         WebSocketSession establishedSession;
@@ -87,7 +77,6 @@ void prepareRoutedService() {
                 RoutedServices routesForSpecificValidService = mock(RoutedServices.class);
                 when(routesForSpecificValidService.findServiceByGatewayUrl("ws/v1"))
                     .thenReturn(new RoutedService("ws-v1", "ws/v1", "/valid-service/ws/v1"));
-                ServiceInstance foundService = validServiceInstance();
 
                 underTest.addRoutedServices(serviceId, routesForSpecificValidService);
             }
@@ -223,11 +212,24 @@ void prepareSessionMock() {
         void whenTheConnectionIsClosed_thenTheSessionIsClosedAndRemovedFromRepository() {
             CloseStatus normalClose = CloseStatus.NORMAL;
 
+            assertThat(routedSessions.entrySet(), not(empty()));
+
             underTest.afterConnectionClosed(establishedSession, normalClose);
 
             assertThat(routedSessions.entrySet(), hasSize(0));
         }
 
+        @Test
+        void whenTheConnectionIsClosed_thenClientSessionIsAlsoClosed() throws IOException {
+            CloseStatus normalClose = CloseStatus.NORMAL;
+            WebSocketSession clientSession = mock(WebSocketSession.class);
+            when(internallyStoredSession.getWebSocketClientSession()).thenReturn(clientSession);
+
+            underTest.afterConnectionClosed(establishedSession, normalClose);
+            verify(clientSession, times(1)).close(normalClose);
+            assertThat(routedSessions.entrySet(), hasSize(0));
+        }
+
         @Test
         void whenTheMessageIsReceived_thenTheMessageIsPassedToTheSession() throws Exception {
             underTest.handleMessage(establishedSession, passedMessage);
@@ -282,7 +284,7 @@ void thenReturnThem() {
     class WhenGettingSubProtocols {
         @Test
         void thenReturnThem() {
-            ArrayList protocol = new ArrayList();
+            List<String> protocol = new ArrayList<>();
             protocol.add("protocol");
             ReflectionTestUtils.setField(underTest, "subProtocols", protocol);
             List<String> subProtocols = underTest.getSubProtocols();

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
@@ -10,16 +10,26 @@
 
 package org.zowe.apiml.integration.proxy;
 
-import io.restassured.RestAssured;
+import static io.restassured.RestAssured.given;
+import static org.apache.http.HttpStatus.SC_OK;
+import static org.apache.tomcat.websocket.Constants.SSL_CONTEXT_PROPERTY;
+import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.zowe.apiml.util.requests.Endpoints.DISCOVERABLE_WS_HEADER;
+import static org.zowe.apiml.util.requests.Endpoints.DISCOVERABLE_WS_UPPERCASE;
+
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.util.Base64;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicInteger;
+
 import org.apache.http.client.utils.URIBuilder;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
-
-
-
-
 import org.springframework.web.socket.CloseStatus;
 import org.springframework.web.socket.TextMessage;
 import org.springframework.web.socket.WebSocketHttpHeaders;
@@ -35,19 +45,7 @@
 import org.zowe.apiml.util.http.HttpClientUtils;
 import org.zowe.apiml.util.http.HttpRequestUtils;
 
-import java.net.URI;
-import java.net.URISyntaxException;
-import java.util.Base64;
-import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicInteger;
-
-import static io.restassured.RestAssured.given;
-import static org.apache.http.HttpStatus.SC_OK;
-import static org.apache.tomcat.websocket.Constants.SSL_CONTEXT_PROPERTY;
-import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.zowe.apiml.util.requests.Endpoints.*;
-import static org.hamcrest.Matchers.is;
+import io.restassured.RestAssured;
 
 @TestsNotMeantForZowe
 @WebsocketTest

diff --git a/schemas/gateway-schema.json b/schemas/gateway-schema.json
@@ -197,6 +197,17 @@
                                                     "description": "How many connection should exists in total?",
                                                     "default": 1000
                                                 },
+                                                "webSocket": {
+                                                    "type": "object",
+                                                    "description": "Customize websocket server parameters",
+                                                    "properties": {
+                                                        "maxIdleTimeout": {
+                                                            "type": "integer",
+                                                            "description": "The gateway acts as a server and client. This parameters customizes the default idle timeout for its client role.",
+                                                            "default": 3600000
+                                                        }
+                                                    }
+                                                },
                                                 "ssl": {
                                                     "type": "object",
                                                     "description": "Network encryption for gateway service connections.",