001/*
002 * Copyright 2002-2018 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.springframework.web.socket.config.annotation;
018
019import java.util.ArrayList;
020import java.util.LinkedHashMap;
021import java.util.List;
022import java.util.Map;
023import java.util.Objects;
024
025import org.springframework.lang.Nullable;
026import org.springframework.scheduling.TaskScheduler;
027import org.springframework.util.MultiValueMap;
028import org.springframework.web.HttpRequestHandler;
029import org.springframework.web.servlet.handler.AbstractHandlerMapping;
030import org.springframework.web.servlet.handler.SimpleUrlHandlerMapping;
031import org.springframework.web.socket.WebSocketHandler;
032import org.springframework.web.socket.server.support.WebSocketHandlerMapping;
033import org.springframework.web.util.UrlPathHelper;
034
035/**
036 * {@link WebSocketHandlerRegistry} with Spring MVC handler mappings for the
037 * handshake requests.
038 *
039 * @author Rossen Stoyanchev
040 * @since 4.0
041 */
042public class ServletWebSocketHandlerRegistry implements WebSocketHandlerRegistry {
043
044        private final List<ServletWebSocketHandlerRegistration> registrations = new ArrayList<>(4);
045
046        private int order = 1;
047
048        @Nullable
049        private UrlPathHelper urlPathHelper;
050
051
052        public ServletWebSocketHandlerRegistry() {
053        }
054
055
056        @Override
057        public WebSocketHandlerRegistration addHandler(WebSocketHandler handler, String... paths) {
058                ServletWebSocketHandlerRegistration registration = new ServletWebSocketHandlerRegistration();
059                registration.addHandler(handler, paths);
060                this.registrations.add(registration);
061                return registration;
062        }
063
064        /**
065         * Set the order for the resulting {@link SimpleUrlHandlerMapping} relative to
066         * other handler mappings configured in Spring MVC.
067         * <p>The default value is 1.
068         */
069        public void setOrder(int order) {
070                this.order = order;
071        }
072
073        public int getOrder() {
074                return this.order;
075        }
076
077        /**
078         * Set the UrlPathHelper to configure on the {@code SimpleUrlHandlerMapping}
079         * used to map handshake requests.
080         */
081        public void setUrlPathHelper(@Nullable UrlPathHelper urlPathHelper) {
082                this.urlPathHelper = urlPathHelper;
083        }
084
085        @Nullable
086        public UrlPathHelper getUrlPathHelper() {
087                return this.urlPathHelper;
088        }
089
090
091        /**
092         * Whether there are any endpoint SockJS registrations without a TaskScheduler.
093         * This method should be invoked just before {@link #getHandlerMapping()} to
094         * allow for registrations to be made first.
095         */
096        protected boolean requiresTaskScheduler() {
097                return this.registrations.stream()
098                                .anyMatch(r -> r.getSockJsServiceRegistration() != null &&
099                                                r.getSockJsServiceRegistration().getTaskScheduler() == null);
100        }
101
102        /**
103         * Provide the TaskScheduler to use for SockJS endpoints for which a task
104         * scheduler has not been explicitly registered. This method must be called
105         * prior to {@link #getHandlerMapping()}.
106         */
107        protected void setTaskScheduler(TaskScheduler scheduler) {
108                this.registrations.stream()
109                                .map(ServletWebSocketHandlerRegistration::getSockJsServiceRegistration)
110                                .filter(Objects::nonNull)
111                                .filter(r -> r.getTaskScheduler() == null)
112                                .forEach(registration -> registration.setTaskScheduler(scheduler));
113        }
114
115        public AbstractHandlerMapping getHandlerMapping() {
116                Map<String, Object> urlMap = new LinkedHashMap<>();
117                for (ServletWebSocketHandlerRegistration registration : this.registrations) {
118                        MultiValueMap<HttpRequestHandler, String> mappings = registration.getMappings();
119                        mappings.forEach((httpHandler, patterns) -> {
120                                for (String pattern : patterns) {
121                                        urlMap.put(pattern, httpHandler);
122                                }
123                        });
124                }
125                WebSocketHandlerMapping hm = new WebSocketHandlerMapping();
126                hm.setUrlMap(urlMap);
127                hm.setOrder(this.order);
128                if (this.urlPathHelper != null) {
129                        hm.setUrlPathHelper(this.urlPathHelper);
130                }
131                return hm;
132        }
133
134}