WIP: [AMQ-9637] Add web socket connection limit.

activemq-amqp/src/test/java/org/apache/activemq/transport/amqp/client/AmqpClientTestSupport.java

@@ -55,7 +55,7 @@ public void tearDown() throws Exception {
     }
 
     public String getConnectorScheme() {
-        return connectorScheme;
+        return connectorScheme.contains("ws") ? connectorScheme.replace("amqp+", "") : connectorScheme;
     }
 
     public boolean isUseSSL() {

activemq-amqp/src/test/java/org/apache/activemq/transport/amqp/interop/AmqpConfiguredMaxConnectionsTest.java

@@ -45,7 +45,9 @@ public class AmqpConfiguredMaxConnectionsTest extends AmqpClientTestSupport {
     public static Collection<Object[]> data() {
         return Arrays.asList(new Object[][] {
                 {"amqp", false},
+                {"amqp+ws", false},
                 {"amqp+nio", false},
+                {"amqp+wss", true}
             });
     }
 

activemq-http/src/main/java/org/apache/activemq/transport/ws/AbstractMQTTSocket.java

@@ -35,7 +35,7 @@
 import org.apache.activemq.util.ServiceStopper;
 import org.fusesource.mqtt.codec.MQTTFrame;
 
-public abstract class AbstractMQTTSocket extends TransportSupport implements MQTTTransport, BrokerServiceAware {
+public abstract class AbstractMQTTSocket extends AbstractWsSocket implements MQTTTransport, BrokerServiceAware {
 
     protected ReentrantLock protocolLock = new ReentrantLock();
     protected volatile MQTTProtocolConverter protocolConverter = null;

activemq-http/src/main/java/org/apache/activemq/transport/ws/AbstractStompSocket.java

@@ -38,7 +38,7 @@
 /**
  * Base implementation of a STOMP based WebSocket handler.
  */
-public abstract class AbstractStompSocket extends TransportSupport implements StompTransport {
+public abstract class AbstractStompSocket extends AbstractWsSocket implements StompTransport {
 
     private static final Logger LOG = LoggerFactory.getLogger(AbstractStompSocket.class);
 

activemq-http/src/main/java/org/apache/activemq/transport/ws/AbstractWsSocket.java

@@ -0,0 +1,43 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.activemq.transport.ws;
+
+import org.apache.activemq.transport.TransportSupport;
+import org.eclipse.jetty.websocket.api.WebSocketListener;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public abstract class AbstractWsSocket extends TransportSupport implements WebSocketListener {
+
+    private static final Logger LOG = LoggerFactory.getLogger(AbstractWsSocket.class);
+
+    @Override
+    public void onWebSocketClose(int statusCode, String reason) {
+        WebSocketListener.super.onWebSocketClose(statusCode, reason);
+
+        try {
+            stop();
+            LOG.debug("Stopped socket: {}", getRemoteAddress());
+        } catch (Exception e) {
+            LOG.debug("Could not stop socket to {}. This exception is ignored.", getRemoteAddress(), e);
+        }
+
+        doWebSocketClose(statusCode, reason);
+    }
+
+    public abstract void doWebSocketClose(int statusCode, String reason);
+}

activemq-http/src/main/java/org/apache/activemq/transport/ws/StompWSConnection.java

@@ -39,7 +39,9 @@ public class StompWSConnection extends WebSocketAdapter implements WebSocketList
     private static final Logger LOG = LoggerFactory.getLogger(StompWSConnection.class);
 
     private Session connection;
+    private Throwable connectionError;
     private final CountDownLatch connectLatch = new CountDownLatch(1);
+    private final CountDownLatch errorLatch = new CountDownLatch(1);
 
     private final BlockingQueue<String> prefetch = new LinkedBlockingDeque<String>();
     private final StompWireFormat wireFormat = new StompWireFormat();
@@ -49,7 +51,7 @@ public class StompWSConnection extends WebSocketAdapter implements WebSocketList
 
     @Override
     public boolean isConnected() {
-        return connection != null ? connection.isOpen() : false;
+        return connection != null && connection.isOpen();
     }
 
     public void close() {
@@ -62,6 +64,9 @@ protected Session getConnection() {
         return connection;
     }
 
+    public Throwable getConnectionError() {
+        return connectionError;
+    }
     //---- Send methods ------------------------------------------------------//
 
     public synchronized void sendRawFrame(String rawFrame) throws Exception {
@@ -106,6 +111,10 @@ public boolean awaitConnection(long time, TimeUnit unit) throws InterruptedExcep
         return connectLatch.await(time, unit);
     }
 
+    public boolean awaitError(long time, TimeUnit unit) throws InterruptedException {
+        return errorLatch.await(time, unit);
+    }
+
     //----- Property Accessors -----------------------------------------------//
 
     public int getCloseCode() {
@@ -148,6 +157,13 @@ public void onWebSocketConnect(org.eclipse.jetty.websocket.api.Session session)
         this.connectLatch.countDown();
     }
 
+    @Override
+    public void onWebSocketError(Throwable cause) {
+        this.connection = null;
+        this.connectionError = cause;
+        this.errorLatch.countDown();
+    }
+
     //----- Internal implementation ------------------------------------------//
 
     private void checkConnected() throws IOException {

activemq-http/src/main/java/org/apache/activemq/transport/ws/WSTransportProxy.java

@@ -43,7 +43,7 @@
  * A proxy class that manages sending WebSocket events to the wrapped protocol level
  * WebSocket Transport.
  */
-public final class WSTransportProxy extends TransportSupport implements Transport, WebSocketListener, BrokerServiceAware, WSTransportSink {
+public final class WSTransportProxy extends AbstractWsSocket implements Transport, WebSocketListener, BrokerServiceAware, WSTransportSink {
 
     private static final Logger LOG = LoggerFactory.getLogger(WSTransportProxy.class);
 
@@ -201,7 +201,7 @@ public void onWebSocketText(String data) {
     }
 
     @Override
-    public void onWebSocketClose(int statusCode, String reason) {
+    public void doWebSocketClose(int statusCode, String reason) {
         try {
             if (protocolLock.tryLock() || protocolLock.tryLock(ORDERLY_CLOSE_TIMEOUT, TimeUnit.SECONDS)) {
                 LOG.debug("WebSocket closed: code[{}] message[{}]", statusCode, reason);

activemq-http/src/main/java/org/apache/activemq/transport/ws/WSTransportServer.java

@@ -20,6 +20,7 @@
 import java.net.InetSocketAddress;
 import java.net.URI;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import jakarta.servlet.Servlet;
 
@@ -53,6 +54,11 @@ public class WSTransportServer extends WebTransportServerSupport implements Brok
     private BrokerService brokerService;
     private WSServlet servlet;
 
+    /**
+     * The maximum number of sockets allowed for this server
+     */
+    protected int maximumConnections = Integer.MAX_VALUE;
+
     public WSTransportServer(URI location) {
         super(location);
         this.bindAddress = location;
@@ -122,6 +128,7 @@ private Servlet createWSServlet() throws Exception {
         servlet = new WSServlet();
         servlet.setTransportOptions(transportOptions);
         servlet.setBrokerService(brokerService);
+        servlet.setMaximumConnections(maximumConnections);
 
         return servlet;
     }
@@ -176,12 +183,35 @@ public void setBrokerService(BrokerService brokerService) {
 
     @Override
     public long getMaxConnectionExceededCount() {
-        // Max Connection Count not supported for ws
-        return -1l;
+        if (servlet != null) {
+            return servlet.getMaxConnectionExceededCount();
+        }
+        return 0;
     }
 
     @Override
     public void resetStatistics() {
-        // Statistics not implemented for ws
+        if(servlet != null) {
+            servlet.resetStatistics();
+        }
     }
+
+    public int getMaximumConnections() {
+        return maximumConnections;
+    }
+
+    public void setMaximumConnections(int maximumConnections) {
+        this.maximumConnections = maximumConnections;
+        if (servlet != null) {
+            servlet.setMaximumConnections(maximumConnections);
+        }
+    }
+
+    public AtomicInteger getCurrentTransportCount() {
+        if (servlet != null) {
+            return servlet.getCurrentTransportCount();
+        }
+        return new AtomicInteger(0);
+    }
+
 }

activemq-http/src/main/java/org/apache/activemq/transport/ws/jetty11/MQTTSocket.java

@@ -95,7 +95,7 @@ public void onWebSocketBinary(byte[] bytes, int offset, int length) {
     }
 
     @Override
-    public void onWebSocketClose(int arg0, String arg1) {
+    public void doWebSocketClose(int arg0, String arg1) {
         try {
             if (protocolLock.tryLock() || protocolLock.tryLock(ORDERLY_CLOSE_TIMEOUT, TimeUnit.SECONDS)) {
                 LOG.debug("MQTT WebSocket closed: code[{}] message[{}]", arg0, arg1);
@@ -120,14 +120,6 @@ public void onWebSocketConnect(Session session) {
         this.session.setIdleTimeout(Duration.ZERO);
     }
 
-    @Override
-    public void onWebSocketError(Throwable arg0) {
-
-    }
-
-    @Override
-    public void onWebSocketText(String arg0) {
-    }
 
     private static int getDefaultSendTimeOut() {
         return Integer.getInteger("org.apache.activemq.transport.ws.MQTTSocket.sendTimeout", 30);

activemq-http/src/main/java/org/apache/activemq/transport/ws/jetty11/StompSocket.java

@@ -69,7 +69,7 @@ public void onWebSocketBinary(byte[] arg0, int arg1, int arg2) {
     }
 
     @Override
-    public void onWebSocketClose(int arg0, String arg1) {
+    public void doWebSocketClose(int arg0, String arg1) {
         try {
             if (protocolLock.tryLock() || protocolLock.tryLock(ORDERLY_CLOSE_TIMEOUT, TimeUnit.SECONDS)) {
                 LOG.debug("Stomp WebSocket closed: code[{}] message[{}]", arg0, arg1);

activemq-http/src/main/java/org/apache/activemq/transport/ws/jetty11/WSServlet.java

@@ -26,18 +26,23 @@
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 
 import jakarta.servlet.ServletException;
 import jakarta.servlet.http.HttpServletRequest;
 import jakarta.servlet.http.HttpServletResponse;
 
+import org.apache.activemq.Service;
 import org.apache.activemq.broker.BrokerService;
 import org.apache.activemq.broker.BrokerServiceAware;
 import org.apache.activemq.transport.Transport;
 import org.apache.activemq.transport.TransportAcceptListener;
 import org.apache.activemq.transport.TransportFactory;
+import org.apache.activemq.transport.tcp.ExceededMaximumConnectionsException;
 import org.apache.activemq.transport.util.HttpTransportUtils;
 import org.apache.activemq.transport.ws.WSTransportProxy;
+import org.apache.activemq.util.ServiceListener;
 import org.eclipse.jetty.websocket.api.WebSocketListener;
 import org.eclipse.jetty.websocket.server.JettyServerUpgradeRequest;
 import org.eclipse.jetty.websocket.server.JettyServerUpgradeResponse;
@@ -48,7 +53,7 @@
 /**
  * Handle connection upgrade requests and creates web sockets
  */
-public class WSServlet extends JettyWebSocketServlet implements BrokerServiceAware {
+public class WSServlet extends JettyWebSocketServlet implements BrokerServiceAware, ServiceListener {
 
     private static final long serialVersionUID = -4716657876092884139L;
 
@@ -60,6 +65,10 @@ public class WSServlet extends JettyWebSocketServlet implements BrokerServiceAwa
     private Map<String, Object> transportOptions;
     private BrokerService brokerService;
 
+    private int maximumConnections = Integer.MAX_VALUE;
+    protected final AtomicLong maximumConnectionsExceededCount = new AtomicLong();
+    private final AtomicInteger currentTransportCount = new AtomicInteger();
+
     private enum Protocol {
         MQTT, STOMP, UNKNOWN
     }
@@ -93,6 +102,23 @@ public void configure(JettyWebSocketServletFactory factory) {
         factory.setCreator(new JettyWebSocketCreator() {
             @Override
             public Object createWebSocket(JettyServerUpgradeRequest req, JettyServerUpgradeResponse resp) {
+                int currentCount;
+                do {
+                    currentCount = currentTransportCount.get();
+                    if (currentCount >= maximumConnections) {
+                        maximumConnectionsExceededCount.incrementAndGet();
+                        listener.onAcceptError(new ExceededMaximumConnectionsException(
+                                "Exceeded the maximum number of allowed client connections. See the '" +
+                                        "maximumConnections' property on the WS transport configuration URI " +
+                                        "in the ActiveMQ configuration file (e.g., activemq.xml)"));
+                        return null;
+                    }
+
+                    //Increment this value before configuring the transport
+                    //This is necessary because some of the transport servers must read from the
+                    //socket during configureTransport() so we want to make sure this value is
+                    //accurate as the transport server could pause here waiting for data to be sent from a client
+                } while(!currentTransportCount.compareAndSet(currentCount, currentCount + 1));
                 WebSocketListener socket;
                 Protocol requestedProtocol = Protocol.UNKNOWN;
 
@@ -114,6 +140,7 @@ public Object createWebSocket(JettyServerUpgradeRequest req, JettyServerUpgradeR
                         socket = new MQTTSocket(HttpTransportUtils.generateWsRemoteAddress(req.getHttpServletRequest()));
                         ((MQTTSocket) socket).setTransportOptions(new HashMap<>(transportOptions));
                         ((MQTTSocket) socket).setPeerCertificates(req.getCertificates());
+                        ((MQTTSocket) socket).addServiceListener(WSServlet.this);
                         resp.setAcceptedSubProtocol(getAcceptedSubProtocol(mqttProtocols, req.getSubProtocols(), "mqtt"));
                         break;
                     case UNKNOWN:
@@ -124,11 +151,13 @@ public Object createWebSocket(JettyServerUpgradeRequest req, JettyServerUpgradeR
                     case STOMP:
                         socket = new StompSocket(HttpTransportUtils.generateWsRemoteAddress(req.getHttpServletRequest()));
                         ((StompSocket) socket).setPeerCertificates(req.getCertificates());
+                        ((StompSocket) socket).addServiceListener(WSServlet.this);
                         resp.setAcceptedSubProtocol(getAcceptedSubProtocol(stompProtocols, req.getSubProtocols(), "stomp"));
                         break;
                     default:
                         socket = null;
                         listener.onAcceptError(new IOException("Unknown protocol requested"));
+                        currentTransportCount.decrementAndGet();
                         break;
                 }
 
@@ -160,6 +189,7 @@ private WebSocketListener findWSTransport(JettyServerUpgradeRequest request, Jet
                 proxy = new WSTransportProxy(remoteAddress, transport);
                 proxy.setPeerCertificates(request.getCertificates());
                 proxy.setTransportOptions(new HashMap<>(transportOptions));
+                proxy.addServiceListener(this);
 
                 response.setAcceptedSubProtocol(proxy.getSubProtocol());
             } catch (Exception e) {
@@ -217,4 +247,42 @@ public void setTransportOptions(Map<String, Object> transportOptions) {
     public void setBrokerService(BrokerService brokerService) {
         this.brokerService = brokerService;
     }
+
+
+    @Override
+    public void started(Service service) {
+    }
+
+    @Override
+    public void stopped(Service service) {
+        this.currentTransportCount.decrementAndGet();
+    }
+
+    /**
+     * @return the maximumConnections
+     */
+    public int getMaximumConnections() {
+        return maximumConnections;
+    }
+
+
+    public long getMaxConnectionExceededCount() {
+        return this.maximumConnectionsExceededCount.get();
+    }
+
+    public void resetStatistics() {
+        this.maximumConnectionsExceededCount.set(0L);
+    }
+
+    /**
+     * @param maximumConnections
+     *            the maximumConnections to set
+     */
+    public void setMaximumConnections(int maximumConnections) {
+        this.maximumConnections = maximumConnections;
+    }
+
+    public AtomicInteger getCurrentTransportCount() {
+        return currentTransportCount;
+    }
 }