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}