content: /**
   * 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.jetty11;
  
  import java.io.IOException;
  import java.net.URI;
  import java.util.ArrayList;
  import java.util.Collections;
  import java.util.Comparator;
  import java.util.HashMap;
  import java.util.List;
  import java.util.Map;
  import java.util.concurrent.ConcurrentHashMap;
  
  import jakarta.servlet.ServletException;
  import jakarta.servlet.http.HttpServletRequest;
  import jakarta.servlet.http.HttpServletResponse;
  
  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.util.HttpTransportUtils;
  import org.apache.activemq.transport.ws.WSTransportProxy;
  import org.eclipse.jetty.websocket.api.WebSocketListener;
  import org.eclipse.jetty.websocket.server.JettyServerUpgradeRequest;
  import org.eclipse.jetty.websocket.server.JettyServerUpgradeResponse;
  import org.eclipse.jetty.websocket.server.JettyWebSocketCreator;
  import org.eclipse.jetty.websocket.server.JettyWebSocketServlet;
  import org.eclipse.jetty.websocket.server.JettyWebSocketServletFactory;
  
  /**
   * Handle connection upgrade requests and creates web sockets
   */
  public class WSServlet extends JettyWebSocketServlet implements BrokerServiceAware {
  
      private static final long serialVersionUID = -4716657876092884139L;
  
      private TransportAcceptListener listener;
  
      private final static Map<String, Integer> stompProtocols = new ConcurrentHashMap<>();
      private final static Map<String, Integer> mqttProtocols = new ConcurrentHashMap<>();
  
      private Map<String, Object> transportOptions;
      private BrokerService brokerService;
  
      private enum Protocol {
          MQTT, STOMP, UNKNOWN
      }
  
      static {
          stompProtocols.put("v12.stomp", 3);
          stompProtocols.put("v11.stomp", 2);
          stompProtocols.put("v10.stomp", 1);
          stompProtocols.put("stomp", 0);
  
          mqttProtocols.put("mqttv3.1", 1);
          mqttProtocols.put("mqtt", 0);
      }
  
      @Override
      public void init() throws ServletException {
          super.init();
          listener = (TransportAcceptListener) getServletContext().getAttribute("acceptListener");
          if (listener == null) {
              throw new ServletException("No such attribute 'acceptListener' available in the ServletContext");
          }
      }
  
      @Override
      protected void doGet(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
          //return empty response - AMQ-6491
      }
  
      @Override
      public void configure(JettyWebSocketServletFactory factory) {
          factory.setCreator(new JettyWebSocketCreator() {
              @Override
              public Object createWebSocket(JettyServerUpgradeRequest req, JettyServerUpgradeResponse resp) {
                  WebSocketListener socket;
                  Protocol requestedProtocol = Protocol.UNKNOWN;
  
                  // When no sub-protocol is requested we default to STOMP for legacy reasons.
                  if (!req.getSubProtocols().isEmpty()) {
                      for (String subProtocol : req.getSubProtocols()) {
                          if (subProtocol.startsWith("mqtt")) {
                              requestedProtocol = Protocol.MQTT;
                          } else if (subProtocol.contains("stomp")) {
                              requestedProtocol = Protocol.STOMP;
                          }
                      }
                  } else {
                      requestedProtocol = Protocol.STOMP;
                  }
  
                  switch (requestedProtocol) {
                      case MQTT:
                          socket = new MQTTSocket(HttpTransportUtils.generateWsRemoteAddress(req.getHttpServletRequest()));
                          ((MQTTSocket) socket).setTransportOptions(new HashMap<>(transportOptions));
                          ((MQTTSocket) socket).setPeerCertificates(req.getCertificates());
                          resp.setAcceptedSubProtocol(getAcceptedSubProtocol(mqttProtocols, req.getSubProtocols(), "mqtt"));
                          break;
                      case UNKNOWN:
                          socket = findWSTransport(req, resp);
                          if (socket != null) {
                              break;
                          }
                      case STOMP:
                          socket = new StompSocket(HttpTransportUtils.generateWsRemoteAddress(req.getHttpServletRequest()));
                          ((StompSocket) socket).setPeerCertificates(req.getCertificates());
                          resp.setAcceptedSubProtocol(getAcceptedSubProtocol(stompProtocols, req.getSubProtocols(), "stomp"));
                          break;
                      default:
                          socket = null;
                          listener.onAcceptError(new IOException("Unknown protocol requested"));
                          break;
                  }
  
                  if (socket != null) {
                      listener.onAccept((Transport) socket);
                  }
  
                  return socket;
              }
          });
      }
  
      private WebSocketListener findWSTransport(JettyServerUpgradeRequest request, JettyServerUpgradeResponse response) {
          WSTransportProxy proxy = null;
  
          for (String subProtocol : request.getSubProtocols()) {
              try {
                  String remoteAddress = HttpTransportUtils.generateWsRemoteAddress(request.getHttpServletRequest(), subProtocol);
                  URI remoteURI = new URI(remoteAddress);
  
                  TransportFactory factory = TransportFactory.findTransportFactory(remoteURI);
  
                  if (factory instanceof BrokerServiceAware) {
                      ((BrokerServiceAware) factory).setBrokerService(brokerService);
                  }
  
                  Transport transport = factory.doConnect(remoteURI);
  
                  proxy = new WSTransportProxy(remoteAddress, transport);
                  proxy.setPeerCertificates(request.getCertificates());
                  proxy.setTransportOptions(new HashMap<>(transportOptions));
  
                  response.setAcceptedSubProtocol(proxy.getSubProtocol());
              } catch (Exception e) {
                  proxy = null;
  
                  // Keep going and try any other sub-protocols present.
                  continue;
              }
          }
  
          return proxy;
      }
  
      private String getAcceptedSubProtocol(final Map<String, Integer> protocols, List<String> subProtocols, String defaultProtocol) {
          List<SubProtocol> matchedProtocols = new ArrayList<>();
          if (subProtocols != null && subProtocols.size() > 0) {
              // detect which subprotocols match accepted protocols and add to the
              // list
              for (String subProtocol : subProtocols) {
                  Integer priority = protocols.get(subProtocol);
                  if (subProtocol != null && priority != null) {
                      // only insert if both subProtocol and priority are not null
                      matchedProtocols.add(new SubProtocol(subProtocol, priority));
                  }
              }
              // sort the list by priority
              if (matchedProtocols.size() > 0) {
                  Collections.sort(matchedProtocols, new Comparator<SubProtocol>() {
                      @Override
                      public int compare(SubProtocol s1, SubProtocol s2) {
                          return s2.priority.compareTo(s1.priority);
                      }
                  });
                  return matchedProtocols.get(0).protocol;
              }
          }
          return defaultProtocol;
      }
  
      private class SubProtocol {
          private String protocol;
          private Integer priority;
  
          public SubProtocol(String protocol, Integer priority) {
              this.protocol = protocol;
              this.priority = priority;
          }
      }
  
      public void setTransportOptions(Map<String, Object> transportOptions) {
          this.transportOptions = transportOptions;
      }
  
      @Override
      public void setBrokerService(BrokerService brokerService) {
          this.brokerService = brokerService;
      }
  }
  